mirror of
https://github.com/nushell/nushell.git
synced 2024-12-23 15:39:06 +01:00
displaying span information, creation time, and size with polars ls (#12472)
# 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:
parent
872945ae8e
commit
b9c2f9ee56
@ -3,19 +3,28 @@ use std::{
|
||||
sync::{Mutex, MutexGuard},
|
||||
};
|
||||
|
||||
use chrono::{DateTime, FixedOffset, Local};
|
||||
use nu_plugin::EngineInterface;
|
||||
use nu_protocol::{LabeledError, ShellError};
|
||||
use nu_protocol::{LabeledError, ShellError, Span};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::{plugin_debug, values::PolarsPluginObject, PolarsPlugin};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CacheValue {
|
||||
pub uuid: Uuid,
|
||||
pub value: PolarsPluginObject,
|
||||
pub created: DateTime<FixedOffset>,
|
||||
pub span: Span,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct Cache {
|
||||
cache: Mutex<HashMap<Uuid, PolarsPluginObject>>,
|
||||
cache: Mutex<HashMap<Uuid, CacheValue>>,
|
||||
}
|
||||
|
||||
impl Cache {
|
||||
fn lock(&self) -> Result<MutexGuard<HashMap<Uuid, PolarsPluginObject>>, ShellError> {
|
||||
fn lock(&self) -> Result<MutexGuard<HashMap<Uuid, CacheValue>>, ShellError> {
|
||||
self.cache.lock().map_err(|e| ShellError::GenericError {
|
||||
error: format!("error acquiring cache lock: {e}"),
|
||||
msg: "".into(),
|
||||
@ -31,7 +40,7 @@ impl Cache {
|
||||
&self,
|
||||
maybe_engine: Option<&EngineInterface>,
|
||||
uuid: &Uuid,
|
||||
) -> Result<Option<PolarsPluginObject>, ShellError> {
|
||||
) -> Result<Option<CacheValue>, ShellError> {
|
||||
let mut lock = self.lock()?;
|
||||
let removed = lock.remove(uuid);
|
||||
plugin_debug!("PolarsPlugin: removing {uuid} from cache: {removed:?}");
|
||||
@ -55,7 +64,8 @@ impl Cache {
|
||||
maybe_engine: Option<&EngineInterface>,
|
||||
uuid: Uuid,
|
||||
value: PolarsPluginObject,
|
||||
) -> Result<Option<PolarsPluginObject>, ShellError> {
|
||||
span: Span,
|
||||
) -> Result<Option<CacheValue>, ShellError> {
|
||||
let mut lock = self.lock()?;
|
||||
plugin_debug!("PolarsPlugin: Inserting {uuid} into cache: {value:?}");
|
||||
// turn off plugin gc the first time an entry is added to the cache
|
||||
@ -68,12 +78,18 @@ impl Cache {
|
||||
}
|
||||
_ => (),
|
||||
};
|
||||
let result = lock.insert(uuid, value);
|
||||
let cache_value = CacheValue {
|
||||
uuid,
|
||||
value,
|
||||
created: Local::now().into(),
|
||||
span,
|
||||
};
|
||||
let result = lock.insert(uuid, cache_value);
|
||||
drop(lock);
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub fn get(&self, uuid: &Uuid) -> Result<Option<PolarsPluginObject>, ShellError> {
|
||||
pub fn get(&self, uuid: &Uuid) -> Result<Option<CacheValue>, ShellError> {
|
||||
let lock = self.lock()?;
|
||||
let result = lock.get(uuid).cloned();
|
||||
drop(lock);
|
||||
@ -82,12 +98,11 @@ impl Cache {
|
||||
|
||||
pub fn process_entries<F, T>(&self, mut func: F) -> Result<Vec<T>, ShellError>
|
||||
where
|
||||
F: FnMut((&Uuid, &PolarsPluginObject)) -> Result<T, ShellError>,
|
||||
F: FnMut((&Uuid, &CacheValue)) -> Result<T, ShellError>,
|
||||
{
|
||||
let lock = self.lock()?;
|
||||
let mut vals: Vec<T> = Vec::new();
|
||||
for entry in lock.iter() {
|
||||
eprintln!("entry: {:?}", entry);
|
||||
let val = func(entry)?;
|
||||
vals.push(val);
|
||||
}
|
||||
@ -103,18 +118,24 @@ pub trait Cacheable: Sized + Clone {
|
||||
|
||||
fn from_cache_value(cv: PolarsPluginObject) -> Result<Self, ShellError>;
|
||||
|
||||
fn cache(self, plugin: &PolarsPlugin, engine: &EngineInterface) -> Result<Self, ShellError> {
|
||||
fn cache(
|
||||
self,
|
||||
plugin: &PolarsPlugin,
|
||||
engine: &EngineInterface,
|
||||
span: Span,
|
||||
) -> Result<Self, ShellError> {
|
||||
plugin.cache.insert(
|
||||
Some(engine),
|
||||
self.cache_id().to_owned(),
|
||||
self.to_cache_value()?,
|
||||
span,
|
||||
)?;
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
fn get_cached(plugin: &PolarsPlugin, id: &Uuid) -> Result<Option<Self>, ShellError> {
|
||||
if let Some(cache_value) = plugin.cache.get(id)? {
|
||||
Ok(Some(Self::from_cache_value(cache_value)?))
|
||||
Ok(Some(Self::from_cache_value(cache_value.value)?))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
|
@ -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);
|
||||
|
@ -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,
|
||||
))
|
||||
}
|
||||
|
@ -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,
|
||||
))
|
||||
}
|
||||
|
@ -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)),
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -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(
|
||||
|
@ -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 {
|
||||
|
@ -180,7 +180,7 @@ pub mod test {
|
||||
use crate::values::PolarsPluginObject;
|
||||
use nu_command::IntoDatetime;
|
||||
use nu_plugin_test_support::PluginTest;
|
||||
use nu_protocol::ShellError;
|
||||
use nu_protocol::{ShellError, Span};
|
||||
|
||||
pub fn test_polars_plugin_command(command: &impl PluginCommand) -> Result<(), ShellError> {
|
||||
let mut plugin = PolarsPlugin::default();
|
||||
@ -193,7 +193,10 @@ pub mod test {
|
||||
// if it's a polars plugin object, try to cache it
|
||||
if let Ok(obj) = PolarsPluginObject::try_from_value(&plugin, result) {
|
||||
let id = obj.id();
|
||||
plugin.cache.insert(None, id, obj).unwrap();
|
||||
plugin
|
||||
.cache
|
||||
.insert(None, id, obj, Span::test_data())
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user