holm/nushell
1
1
Fork 0
forked from extern/nushell
Code Issues Pull Requests Packages Projects Releases Wiki Activity

displaying span information, creation time, and size with polars ls (#12472)

Browse Source 
# Description
`polars ls` is already different that `dfr ls`. Currently it just shows
the cache key, columns, rows, and type. I have added:
- creation time
- size
- span contents
-  span start and end

<img width="1471" alt="Screenshot 2024-04-10 at 17 27 06"
src="https://github.com/nushell/nushell/assets/56345/545918b7-7c96-4c25-bc01-b9e2b659a408">

# Tests + Formatting
Done

Co-authored-by: Jack Wright <jack.wright@disqo.com>
This commit is contained in:
Jack Wright
2024-04-12 07:23:46 -07:00
committed by GitHub
parent 872945ae8e
commit b9c2f9ee56
8 changed files with 128 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
crates/nu_plugin_polars/src/dataframe/eager/list.rs
View File

@@ -35,59 +35,86 @@ impl PluginCommand for ListDF {
fn run(
&self,
plugin: &Self::Plugin,
_engine: &EngineInterface,
engine: &EngineInterface,
call: &EvaluatedCall,
_input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let vals = plugin.cache.process_entries(|(key, value)| match value {
PolarsPluginObject::NuDataFrame(df) => Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::int(df.as_ref().width() as i64, call.head),
"rows" => Value::int(df.as_ref().height() as i64, call.head),
"type" => Value::string("NuDataFrame", call.head),
},
call.head,
))),
PolarsPluginObject::NuLazyFrame(lf) => {
let lf = lf.clone().collect(call.head)?;
Ok(Some(Value::record(
let vals = plugin.cache.process_entries(|(key, value)| {
let span_contents = engine.get_span_contents(value.span)?;
let span_contents = String::from_utf8_lossy(&span_contents);
match &value.value {
PolarsPluginObject::NuDataFrame(df) => Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::int(lf.as_ref().width() as i64, call.head),
"rows" => Value::int(lf.as_ref().height() as i64, call.head),
"type" => Value::string("NuLazyFrame", call.head),
"created" => Value::date(value.created, call.head),
"columns" => Value::int(df.as_ref().width() as i64, call.head),
"rows" => Value::int(df.as_ref().height() as i64, call.head),
"type" => Value::string("NuDataFrame", call.head),
"estimated_size" => Value::filesize(df.to_polars().estimated_size() as i64, call.head),
"span_contents" => Value::string(span_contents, value.span),
"span_start" => Value::int(value.span.start as i64, call.head),
"span_end" => Value::int(value.span.end as i64, call.head),
},
call.head,
)))
))),
PolarsPluginObject::NuLazyFrame(lf) => {
let lf = lf.clone().collect(call.head)?;
Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"created" => Value::date(value.created, call.head),
"columns" => Value::int(lf.as_ref().width() as i64, call.head),
"rows" => Value::int(lf.as_ref().height() as i64, call.head),
"type" => Value::string("NuLazyFrame", call.head),
"estimated_size" => Value::filesize(lf.to_polars().estimated_size() as i64, call.head),
"span_contents" => Value::string(span_contents, value.span),
"span_start" => Value::int(value.span.start as i64, call.head),
"span_end" => Value::int(value.span.end as i64, call.head),
},
call.head,
)))
}
PolarsPluginObject::NuExpression(_) => Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"created" => Value::date(value.created, call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuExpression", call.head),
"estimated_size" => Value::nothing(call.head),
"span_contents" => Value::string(span_contents, value.span),
"span_start" => Value::int(value.span.start as i64, call.head),
"span_end" => Value::int(value.span.end as i64, call.head),
},
call.head,
))),
PolarsPluginObject::NuLazyGroupBy(_) => Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuLazyGroupBy", call.head),
"estimated_size" => Value::nothing(call.head),
"span_contents" => Value::string(span_contents, call.head),
"span_start" => Value::int(call.head.start as i64, call.head),
"span_end" => Value::int(call.head.end as i64, call.head),
},
call.head,
))),
PolarsPluginObject::NuWhen(_) => Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuWhen", call.head),
"estimated_size" => Value::nothing(call.head),
"span_contents" => Value::string(span_contents.to_string(), call.head),
"span_start" => Value::int(call.head.start as i64, call.head),
"span_end" => Value::int(call.head.end as i64, call.head),
},
call.head,
))),
}
PolarsPluginObject::NuExpression(_) => Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuExpression", call.head),
},
call.head,
))),
PolarsPluginObject::NuLazyGroupBy(_) => Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuLazyGroupBy", call.head),
},
call.head,
))),
PolarsPluginObject::NuWhen(_) => Ok(Some(Value::record(
record! {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuWhen", call.head),
},
call.head,
))),
})?;
let vals = vals.into_iter().flatten().collect();
let list = Value::list(vals, call.head);

6
crates/nu_plugin_polars/src/dataframe/lazy/collect.rs
View File

@@ -66,14 +66,16 @@ impl PluginCommand for LazyCollect {
PolarsPluginObject::NuLazyFrame(lazy) => {
let eager = lazy.collect(call.head)?;
Ok(PipelineData::Value(
eager.cache(plugin, engine)?.into_value(call.head),
eager
.cache(plugin, engine, call.head)?
.into_value(call.head),
None,
))
}
PolarsPluginObject::NuDataFrame(df) => {
// just return the dataframe, add to cache again to be safe
Ok(PipelineData::Value(
df.cache(plugin, engine)?.into_value(call.head),
df.cache(plugin, engine, call.head)?.into_value(call.head),
None,
))
}

2
crates/nu_plugin_polars/src/dataframe/lazy/to_lazy.rs
View File

@@ -54,7 +54,7 @@ impl PluginCommand for ToLazyFrame {
let df = NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema)?;
let lazy = NuLazyFrame::from_dataframe(df);
Ok(PipelineData::Value(
lazy.cache(plugin, engine)?.into_value(call.head),
lazy.cache(plugin, engine, call.head)?.into_value(call.head),
None,
))
}

6
crates/nu_plugin_polars/src/dataframe/values/mod.rs
View File

@@ -318,14 +318,14 @@ pub fn cache_and_to_value(
// if it was from a lazy value, make it lazy again
PolarsPluginObject::NuDataFrame(df) if df.from_lazy => {
let df = df.lazy();
Ok(df.cache(plugin, engine)?.into_value(span))
Ok(df.cache(plugin, engine, span)?.into_value(span))
}
// if it was from an eager value, make it eager again
PolarsPluginObject::NuLazyFrame(lf) if lf.from_eager => {
let lf = lf.collect(span)?;
Ok(lf.cache(plugin, engine)?.into_value(span))
Ok(lf.cache(plugin, engine, span)?.into_value(span))
}
_ => Ok(cv.cache(plugin, engine)?.into_value(span)),
_ => Ok(cv.cache(plugin, engine, span)?.into_value(span)),
}
}

6
crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/custom_value.rs
View File

@@ -81,7 +81,7 @@ impl PolarsPluginCustomValue for NuDataFrameCustomValue {
let df = NuDataFrame::try_from_custom_value(plugin, self)?;
Ok(df
.compute_with_value(plugin, lhs_span, operator.item, operator.span, &right)?
.cache(plugin, engine)?
.cache(plugin, engine, lhs_span)?
.into_value(lhs_span))
}
@@ -105,7 +105,9 @@ impl PolarsPluginCustomValue for NuDataFrameCustomValue {
) -> Result<Value, ShellError> {
let df = NuDataFrame::try_from_custom_value(plugin, self)?;
let column = df.column(&column_name.item, self_span)?;
Ok(column.cache(plugin, engine)?.into_value(self_span))
Ok(column
.cache(plugin, engine, self_span)?
.into_value(self_span))
}
fn custom_value_partial_cmp(

14
crates/nu_plugin_polars/src/dataframe/values/nu_expression/custom_value.rs
View File

@@ -135,32 +135,32 @@ fn with_operator(
Operator::Comparison(Comparison::Equal) => Ok(left
.clone()
.apply_with_expr(right.clone(), Expr::eq)
.cache(plugin, engine)?
.cache(plugin, engine, lhs_span)?
.into_value(lhs_span)),
Operator::Comparison(Comparison::NotEqual) => Ok(left
.clone()
.apply_with_expr(right.clone(), Expr::neq)
.cache(plugin, engine)?
.cache(plugin, engine, lhs_span)?
.into_value(lhs_span)),
Operator::Comparison(Comparison::GreaterThan) => Ok(left
.clone()
.apply_with_expr(right.clone(), Expr::gt)
.cache(plugin, engine)?
.cache(plugin, engine, lhs_span)?
.into_value(lhs_span)),
Operator::Comparison(Comparison::GreaterThanOrEqual) => Ok(left
.clone()
.apply_with_expr(right.clone(), Expr::gt_eq)
.cache(plugin, engine)?
.cache(plugin, engine, lhs_span)?
.into_value(lhs_span)),
Operator::Comparison(Comparison::LessThan) => Ok(left
.clone()
.apply_with_expr(right.clone(), Expr::lt)
.cache(plugin, engine)?
.cache(plugin, engine, lhs_span)?
.into_value(lhs_span)),
Operator::Comparison(Comparison::LessThanOrEqual) => Ok(left
.clone()
.apply_with_expr(right.clone(), Expr::lt_eq)
.cache(plugin, engine)?
.cache(plugin, engine, lhs_span)?
.into_value(lhs_span)),
_ => Err(ShellError::OperatorMismatch {
op_span,
@@ -185,7 +185,7 @@ where
{
let expr: NuExpression = f(left.as_ref().clone(), right.as_ref().clone()).into();
Ok(expr.cache(plugin, engine)?.into_value(span))
Ok(expr.cache(plugin, engine, span)?.into_value(span))
}
impl PolarsPluginCustomValue for NuExpressionCustomValue {