WIP This PR covers migration crates/nu-cmd-dataframes to a new plugin ./crates/nu_plugin_polars ## TODO List Other: - [X] Fix examples - [x] Fix Plugin Test Harness - [X] Move Cache to Mutex<BTreeMap> - [X] Logic for disabling/enabling plugin GC based off whether items are cached. - [x] NuExpression custom values - [X] Optimize caching (don't cache every object creation). - [x] Fix dataframe operations (in NuDataFrameCustomValue::operations) - [x] Added plugin_debug! macro that for checking an env variable POLARS_PLUGIN_DEBUG Fix duplicated commands: - [x] There are two polars median commands, one for lazy and one for expr.. there should only be one that works for both. I temporarily called on polars expr-median (inside expressions_macros.rs) - [x] polars quantile (lazy, and expr). the expr one is temporarily expr-median - [x] polars is-in (renamed one series-is-in) Commands: - [x] AppendDF - [x] CastDF - [X] ColumnsDF - [x] DataTypes - [x] Summary - [x] DropDF - [x] DropDuplicates - [x] DropNulls - [x] Dummies - [x] FilterWith - [X] FirstDF - [x] GetDF - [x] LastDF - [X] ListDF - [x] MeltDF - [X] OpenDataFrame - [x] QueryDf - [x] RenameDF - [x] SampleDF - [x] SchemaDF - [x] ShapeDF - [x] SliceDF - [x] TakeDF - [X] ToArrow - [x] ToAvro - [X] ToCSV - [X] ToDataFrame - [X] ToNu - [x] ToParquet - [x] ToJsonLines - [x] WithColumn - [x] ExprAlias - [x] ExprArgWhere - [x] ExprCol - [x] ExprConcatStr - [x] ExprCount - [x] ExprLit - [x] ExprWhen - [x] ExprOtherwise - [x] ExprQuantile - [x] ExprList - [x] ExprAggGroups - [x] ExprCount - [x] ExprIsIn - [x] ExprNot - [x] ExprMax - [x] ExprMin - [x] ExprSum - [x] ExprMean - [x] ExprMedian - [x] ExprStd - [x] ExprVar - [x] ExprDatePart - [X] LazyAggregate - [x] LazyCache - [X] LazyCollect - [x] LazyFetch - [x] LazyFillNA - [x] LazyFillNull - [x] LazyFilter - [x] LazyJoin - [x] LazyQuantile - [x] LazyMedian - [x] LazyReverse - [x] LazySelect - [x] LazySortBy - [x] ToLazyFrame - [x] ToLazyGroupBy - [x] LazyExplode - [x] LazyFlatten - [x] AllFalse - [x] AllTrue - [x] ArgMax - [x] ArgMin - [x] ArgSort - [x] ArgTrue - [x] ArgUnique - [x] AsDate - [x] AsDateTime - [x] Concatenate - [x] Contains - [x] Cumulative - [x] GetDay - [x] GetHour - [x] GetMinute - [x] GetMonth - [x] GetNanosecond - [x] GetOrdinal - [x] GetSecond - [x] GetWeek - [x] GetWeekDay - [x] GetYear - [x] IsDuplicated - [x] IsIn - [x] IsNotNull - [x] IsNull - [x] IsUnique - [x] NNull - [x] NUnique - [x] NotSeries - [x] Replace - [x] ReplaceAll - [x] Rolling - [x] SetSeries - [x] SetWithIndex - [x] Shift - [x] StrLengths - [x] StrSlice - [x] StrFTime - [x] ToLowerCase - [x] ToUpperCase - [x] Unique - [x] ValueCount --------- Co-authored-by: Jack Wright <jack.wright@disqo.com>
161 lines
5.2 KiB
Rust
161 lines
5.2 KiB
Rust
use crate::{
|
|
dataframe::values::{Column, NuDataFrame, NuLazyFrame},
|
|
values::{
|
|
cant_convert_err, to_pipeline_data, CustomValueSupport, NuExpression, PolarsPluginObject,
|
|
PolarsPluginType,
|
|
},
|
|
PolarsPlugin,
|
|
};
|
|
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
|
|
use nu_protocol::{
|
|
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
|
|
Value,
|
|
};
|
|
use polars::prelude::{lit, QuantileInterpolOptions};
|
|
|
|
#[derive(Clone)]
|
|
pub struct LazyQuantile;
|
|
|
|
impl PluginCommand for LazyQuantile {
|
|
type Plugin = PolarsPlugin;
|
|
|
|
fn name(&self) -> &str {
|
|
"polars quantile"
|
|
}
|
|
|
|
fn usage(&self) -> &str {
|
|
"Aggregates the columns to the selected quantile."
|
|
}
|
|
|
|
fn signature(&self) -> Signature {
|
|
Signature::build(self.name())
|
|
.required(
|
|
"quantile",
|
|
SyntaxShape::Number,
|
|
"quantile value for quantile operation",
|
|
)
|
|
.input_output_types(vec![
|
|
(
|
|
Type::Custom("expression".into()),
|
|
Type::Custom("expression".into()),
|
|
),
|
|
(
|
|
Type::Custom("dataframe".into()),
|
|
Type::Custom("dataframe".into()),
|
|
),
|
|
])
|
|
.category(Category::Custom("lazyframe".into()))
|
|
}
|
|
|
|
fn examples(&self) -> Vec<Example> {
|
|
vec![
|
|
Example {
|
|
description: "quantile value from columns in a dataframe",
|
|
example: "[[a b]; [6 2] [1 4] [4 1]] | polars into-df | polars quantile 0.5",
|
|
result: Some(
|
|
NuDataFrame::try_from_columns(
|
|
vec![
|
|
Column::new("a".to_string(), vec![Value::test_float(4.0)]),
|
|
Column::new("b".to_string(), vec![Value::test_float(2.0)]),
|
|
],
|
|
None,
|
|
)
|
|
.expect("simple df for test should not fail")
|
|
.into_value(Span::test_data()),
|
|
),
|
|
},
|
|
Example {
|
|
description: "Quantile aggregation for a group-by",
|
|
example: r#"[[a b]; [one 2] [one 4] [two 1]]
|
|
| polars into-df
|
|
| polars group-by a
|
|
| polars agg (polars col b | polars quantile 0.5)"#,
|
|
result: Some(
|
|
NuDataFrame::try_from_columns(
|
|
vec![
|
|
Column::new(
|
|
"a".to_string(),
|
|
vec![Value::test_string("one"), Value::test_string("two")],
|
|
),
|
|
Column::new(
|
|
"b".to_string(),
|
|
vec![Value::test_float(4.0), Value::test_float(1.0)],
|
|
),
|
|
],
|
|
None,
|
|
)
|
|
.expect("simple df for test should not fail")
|
|
.into_value(Span::test_data()),
|
|
),
|
|
},
|
|
]
|
|
}
|
|
|
|
fn run(
|
|
&self,
|
|
plugin: &Self::Plugin,
|
|
engine: &EngineInterface,
|
|
call: &EvaluatedCall,
|
|
input: PipelineData,
|
|
) -> Result<PipelineData, LabeledError> {
|
|
let value = input.into_value(call.head);
|
|
let quantile: f64 = call.req(0)?;
|
|
match PolarsPluginObject::try_from_value(plugin, &value)? {
|
|
PolarsPluginObject::NuDataFrame(df) => {
|
|
command(plugin, engine, call, df.lazy(), quantile)
|
|
}
|
|
PolarsPluginObject::NuLazyFrame(lazy) => command(plugin, engine, call, lazy, quantile),
|
|
PolarsPluginObject::NuExpression(expr) => {
|
|
let expr: NuExpression = expr
|
|
.to_polars()
|
|
.quantile(lit(quantile), QuantileInterpolOptions::default())
|
|
.into();
|
|
to_pipeline_data(plugin, engine, call.head, expr)
|
|
}
|
|
_ => Err(cant_convert_err(
|
|
&value,
|
|
&[
|
|
PolarsPluginType::NuDataFrame,
|
|
PolarsPluginType::NuLazyFrame,
|
|
PolarsPluginType::NuExpression,
|
|
],
|
|
)),
|
|
}
|
|
.map_err(LabeledError::from)
|
|
}
|
|
}
|
|
|
|
fn command(
|
|
plugin: &PolarsPlugin,
|
|
engine: &EngineInterface,
|
|
call: &EvaluatedCall,
|
|
lazy: NuLazyFrame,
|
|
quantile: f64,
|
|
) -> Result<PipelineData, ShellError> {
|
|
let lazy = NuLazyFrame::new(
|
|
lazy.from_eager,
|
|
lazy.to_polars()
|
|
.quantile(lit(quantile), QuantileInterpolOptions::default())
|
|
.map_err(|e| ShellError::GenericError {
|
|
error: "Dataframe Error".into(),
|
|
msg: e.to_string(),
|
|
help: None,
|
|
span: None,
|
|
inner: vec![],
|
|
})?,
|
|
);
|
|
|
|
to_pipeline_data(plugin, engine, call.head, lazy)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use super::*;
|
|
use crate::test::test_polars_plugin_command;
|
|
|
|
#[test]
|
|
fn test_examples() -> Result<(), ShellError> {
|
|
test_polars_plugin_command(&LazyQuantile)
|
|
}
|
|
}
|