Added commands for working with the plugin cache. (#12576)

# Description
This pull request provides three new commands:
`polars store-ls` - moved from `polars ls`. It provides the list of all
object stored in the plugin cache
`polars store-rm` - deletes a cached object
`polars store-get` - gets an object from the cache. 

The addition of `polars store-get` required adding a reference_count to
cached entries. `polars get` is the only command that will increment
this value. `polars rm` will remove the value despite it's count. Calls
to PolarsPlugin::custom_value_dropped will decrement the value.

The prefix store- was chosen due to there already being a `polars cache`
command. These commands were not made sub-commands as there isn't a way
to display help for sub commands in plugins (e.g. `polars store`
displaying help) and I felt the store- seemed fine anyways.

The output of `polars store-ls` now shows the reference count for each
object.

# User-Facing Changes
polars ls has now moved to polars store-ls

---------

Co-authored-by: Jack Wright <jack.wright@disqo.com>
This commit is contained in:
Jack Wright 2024-04-21 17:43:43 -07:00 committed by GitHub
parent aad3ac11da
commit a60381a932
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
19 changed files with 324 additions and 50 deletions

View File

@ -0,0 +1,96 @@
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value,
};
use polars::{prelude::NamedFrom, series::Series};
use uuid::Uuid;
use crate::{
values::{CustomValueSupport, NuDataFrame},
PolarsPlugin,
};
#[derive(Clone)]
pub struct CacheGet;
impl PluginCommand for CacheGet {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars store-get"
}
fn usage(&self) -> &str {
"Gets a Dataframe or other object from the plugin cache."
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.required("key", SyntaxShape::String, "Key of objects to get")
.input_output_types(vec![
(Type::Any, Type::Custom("dataframe".into())),
(Type::Any, Type::Custom("expression".into())),
])
.category(Category::Custom("dataframe".into()))
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Get a stored object",
example: r#"let df = ([[a b];[1 2] [3 4]] | polars into-df);
polars store-ls | get key | first | polars store-get $in"#,
result: Some(
NuDataFrame::try_from_series_vec(
vec![Series::new("a", &[1_i64, 3]), Series::new("b", &[2_i64, 4])],
Span::test_data(),
)
.expect("could not create dataframe")
.into_value(Span::test_data()),
),
}]
}
fn run(
&self,
plugin: &Self::Plugin,
_engine: &EngineInterface,
call: &EvaluatedCall,
_input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let key = call
.req::<String>(0)
.and_then(|ref k| as_uuid(k, call.head))?;
let value = if let Some(cache_value) = plugin.cache.get(&key, true)? {
let polars_object = cache_value.value;
polars_object.into_value(call.head)
} else {
Value::nothing(call.head)
};
Ok(PipelineData::Value(value, None))
}
}
fn as_uuid(s: &str, span: Span) -> Result<Uuid, ShellError> {
Uuid::parse_str(s).map_err(|e| ShellError::GenericError {
error: format!("Failed to convert key string to UUID: {e}"),
msg: "".into(),
span: Some(span),
help: None,
inner: vec![],
})
}
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::{First, Get};
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command_with_decls(&CacheGet, vec![Box::new(Get), Box::new(First)])
}
}

View File

@ -12,7 +12,7 @@ impl PluginCommand for ListDF {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars ls"
"polars store-ls"
}
fn usage(&self) -> &str {
@ -26,8 +26,8 @@ impl PluginCommand for ListDF {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Creates a new dataframe and shows it in the dataframe list",
example: r#"let test = ([[a b];[1 2] [3 4]] | dfr into-df);
polars ls"#,
example: r#"let test = ([[a b];[1 2] [3 4]] | polars into-df);
polars store-ls"#,
result: None,
}]
}
@ -49,11 +49,12 @@ impl PluginCommand for ListDF {
"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),
"type" => Value::string("DataFrame", 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),
"reference_count" => Value::int(value.reference_count as i64, call.head),
},
call.head,
))),
@ -65,11 +66,12 @@ impl PluginCommand for ListDF {
"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),
"type" => Value::string("LazyFrame", 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),
"reference_count" => Value::int(value.reference_count as i64, call.head),
},
call.head,
)))
@ -80,11 +82,12 @@ impl PluginCommand for ListDF {
"created" => Value::date(value.created, call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuExpression", call.head),
"type" => Value::string("Expression", 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),
"reference_count" => Value::int(value.reference_count as i64, call.head),
},
call.head,
))),
@ -93,11 +96,12 @@ impl PluginCommand for ListDF {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuLazyGroupBy", call.head),
"type" => Value::string("LazyGroupBy", 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),
"reference_count" => Value::int(value.reference_count as i64, call.head),
},
call.head,
))),
@ -106,11 +110,12 @@ impl PluginCommand for ListDF {
"key" => Value::string(key.to_string(), call.head),
"columns" => Value::nothing(call.head),
"rows" => Value::nothing(call.head),
"type" => Value::string("NuWhen", call.head),
"type" => Value::string("When", 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),
"reference_count" => Value::int(value.reference_count as i64, call.head),
},
call.head,
))),

View File

@ -1,10 +1,15 @@
mod get;
mod list;
mod rm;
use std::{
collections::HashMap,
sync::{Mutex, MutexGuard},
};
use chrono::{DateTime, FixedOffset, Local};
use nu_plugin::EngineInterface;
pub use list::ListDF;
use nu_plugin::{EngineInterface, PluginCommand};
use nu_protocol::{LabeledError, ShellError, Span};
use uuid::Uuid;
@ -16,6 +21,7 @@ pub struct CacheValue {
pub value: PolarsPluginObject,
pub created: DateTime<FixedOffset>,
pub span: Span,
pub reference_count: i16,
}
#[derive(Default)]
@ -35,15 +41,32 @@ impl Cache {
}
/// Removes an item from the plugin cache.
/// The maybe_engine parameter is required outside of testing
///
/// * `maybe_engine` - Current EngineInterface reference. Required outside of testing
/// * `key` - The key of the cache entry to remove.
/// * `force` - Delete even if there are multiple references
pub fn remove(
&self,
maybe_engine: Option<&EngineInterface>,
uuid: &Uuid,
key: &Uuid,
force: bool,
) -> Result<Option<CacheValue>, ShellError> {
let mut lock = self.lock()?;
let removed = lock.remove(uuid);
plugin_debug!("PolarsPlugin: removing {uuid} from cache: {removed:?}");
let reference_count = lock.get_mut(key).map(|cache_value| {
cache_value.reference_count -= 1;
cache_value.reference_count
});
let removed = if force || reference_count.unwrap_or_default() < 1 {
let removed = lock.remove(key);
plugin_debug!("PolarsPlugin: removing {key} from cache: {removed:?}");
removed
} else {
plugin_debug!("PolarsPlugin: decrementing reference count for {key}");
None
};
// Once there are no more entries in the cache
// we can turn plugin gc back on
match maybe_engine {
@ -83,15 +106,21 @@ impl Cache {
value,
created: Local::now().into(),
span,
reference_count: 1,
};
let result = lock.insert(uuid, cache_value);
drop(lock);
Ok(result)
}
pub fn get(&self, uuid: &Uuid) -> Result<Option<CacheValue>, ShellError> {
let lock = self.lock()?;
let result = lock.get(uuid).cloned();
pub fn get(&self, uuid: &Uuid, increment: bool) -> Result<Option<CacheValue>, ShellError> {
let mut lock = self.lock()?;
let result = lock.get_mut(uuid).map(|cv| {
if increment {
cv.reference_count += 1;
}
cv.clone()
});
drop(lock);
Ok(result)
}
@ -134,10 +163,18 @@ pub trait Cacheable: Sized + Clone {
}
fn get_cached(plugin: &PolarsPlugin, id: &Uuid) -> Result<Option<Self>, ShellError> {
if let Some(cache_value) = plugin.cache.get(id)? {
if let Some(cache_value) = plugin.cache.get(id, false)? {
Ok(Some(Self::from_cache_value(cache_value.value)?))
} else {
Ok(None)
}
}
}
pub(crate) fn cache_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> {
vec![
Box::new(ListDF),
Box::new(rm::CacheRemove),
Box::new(get::CacheGet),
]
}

106
crates/nu_plugin_polars/src/cache/rm.rs vendored Normal file
View File

@ -0,0 +1,106 @@
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value,
};
use uuid::Uuid;
use crate::PolarsPlugin;
#[derive(Clone)]
pub struct CacheRemove;
impl PluginCommand for CacheRemove {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars store-rm"
}
fn usage(&self) -> &str {
"Removes a stored Dataframe or other object from the plugin cache."
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.rest("keys", SyntaxShape::String, "Keys of objects to remove")
.input_output_type(Type::Any, Type::List(Box::new(Type::String)))
.category(Category::Custom("dataframe".into()))
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Removes a stored ",
example: r#"let df = ([[a b];[1 2] [3 4]] | polars into-df);
polars store-ls | get key | first | polars store-rm $in"#,
result: None,
}]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
_input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let msgs: Vec<Value> = call
.rest::<String>(0)?
.into_iter()
.map(|ref key| remove_cache_entry(plugin, engine, key, call.head))
.collect::<Result<Vec<Value>, ShellError>>()?;
Ok(PipelineData::Value(Value::list(msgs, call.head), None))
}
}
fn remove_cache_entry(
plugin: &PolarsPlugin,
engine: &EngineInterface,
key: &str,
span: Span,
) -> Result<Value, ShellError> {
let key = as_uuid(key, span)?;
let msg = plugin
.cache
.remove(Some(engine), &key, true)?
.map(|_| format!("Removed: {key}"))
.unwrap_or_else(|| format!("No value found for key: {key}"));
Ok(Value::string(msg, span))
}
fn as_uuid(s: &str, span: Span) -> Result<Uuid, ShellError> {
Uuid::parse_str(s).map_err(|e| ShellError::GenericError {
error: format!("Failed to convert key string to UUID: {e}"),
msg: "".into(),
span: Some(span),
help: None,
inner: vec![],
})
}
#[cfg(test)]
mod test {
use nu_command::{First, Get};
use nu_plugin_test_support::PluginTest;
use nu_protocol::Span;
use super::*;
#[test]
fn test_remove() -> Result<(), ShellError> {
let plugin = PolarsPlugin::new_test_mode().into();
let pipeline_data = PluginTest::new("polars", plugin)?
.add_decl(Box::new(First))?
.add_decl(Box::new(Get))?
.eval("let df = ([[a b];[1 2] [3 4]] | polars into-df); polars store-ls | get key | first | polars store-rm $in")?;
let value = pipeline_data.into_value(Span::test_data());
let msg = value
.as_list()?
.first()
.expect("there should be a first entry")
.as_str()?;
assert!(msg.contains("Removed"));
Ok(())
}
}

View File

@ -9,7 +9,6 @@ mod filter_with;
mod first;
mod get;
mod last;
mod list;
mod melt;
mod open;
mod query_df;
@ -45,7 +44,6 @@ pub use filter_with::FilterWith;
pub use first::FirstDF;
pub use get::GetDF;
pub use last::LastDF;
pub use list::ListDF;
pub use melt::MeltDF;
use nu_plugin::PluginCommand;
pub use query_df::QueryDf;
@ -82,7 +80,6 @@ pub(crate) fn eager_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugi
Box::new(Summary),
Box::new(FirstDF),
Box::new(LastDF),
Box::new(ListDF),
Box::new(RenameDF),
Box::new(SampleDF),
Box::new(ShapeDF),

View File

@ -186,10 +186,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&AsDateTime)
test_polars_plugin_command_with_decls(&AsDateTime, vec![Box::new(IntoDatetime)])
}
}

View File

@ -93,10 +93,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetDay)
test_polars_plugin_command_with_decls(&GetDay, vec![Box::new(IntoDatetime)])
}
}

View File

@ -85,10 +85,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetHour)
test_polars_plugin_command_with_decls(&GetHour, vec![Box::new(IntoDatetime)])
}
}

View File

@ -83,10 +83,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetMinute)
test_polars_plugin_command_with_decls(&GetMinute, vec![Box::new(IntoDatetime)])
}
}

View File

@ -85,10 +85,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetMonth)
test_polars_plugin_command_with_decls(&GetMonth, vec![Box::new(IntoDatetime)])
}
}

View File

@ -85,10 +85,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetNanosecond)
test_polars_plugin_command_with_decls(&GetNanosecond, vec![Box::new(IntoDatetime)])
}
}

View File

@ -85,10 +85,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetOrdinal)
test_polars_plugin_command_with_decls(&GetOrdinal, vec![Box::new(IntoDatetime)])
}
}

View File

@ -85,10 +85,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetSecond)
test_polars_plugin_command_with_decls(&GetSecond, vec![Box::new(IntoDatetime)])
}
}

View File

@ -85,10 +85,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetWeek)
test_polars_plugin_command_with_decls(&GetWeek, vec![Box::new(IntoDatetime)])
}
}

View File

@ -85,10 +85,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetWeekDay)
test_polars_plugin_command_with_decls(&GetWeekDay, vec![Box::new(IntoDatetime)])
}
}

View File

@ -85,10 +85,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&GetYear)
test_polars_plugin_command_with_decls(&GetYear, vec![Box::new(IntoDatetime)])
}
}

View File

@ -104,10 +104,11 @@ fn command(
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use crate::test::test_polars_plugin_command_with_decls;
use nu_command::IntoDatetime;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&StrFTime)
test_polars_plugin_command_with_decls(&StrFTime, vec![Box::new(IntoDatetime)])
}
}

View File

@ -107,6 +107,16 @@ impl PolarsPluginObject {
PolarsPluginObject::NuWhen(w) => w.id,
}
}
pub fn into_value(self, span: Span) -> Value {
match self {
PolarsPluginObject::NuDataFrame(df) => df.into_value(span),
PolarsPluginObject::NuLazyFrame(lf) => lf.into_value(span),
PolarsPluginObject::NuExpression(e) => e.into_value(span),
PolarsPluginObject::NuLazyGroupBy(lg) => lg.into_value(span),
PolarsPluginObject::NuWhen(w) => w.into_value(span),
}
}
}
#[derive(Debug, Clone)]

View File

@ -1,5 +1,6 @@
use std::cmp::Ordering;
use cache::cache_commands;
pub use cache::{Cache, Cacheable};
use dataframe::{stub::PolarsCmd, values::CustomValueType};
use nu_plugin::{EngineInterface, Plugin, PluginCommand};
@ -40,6 +41,7 @@ impl Plugin for PolarsPlugin {
commands.append(&mut lazy_commands());
commands.append(&mut expr_commands());
commands.append(&mut series_commands());
commands.append(&mut cache_commands());
commands
}
@ -50,7 +52,7 @@ impl Plugin for PolarsPlugin {
) -> Result<(), LabeledError> {
if !self.disable_cache_drop {
let id = CustomValueType::try_from_custom_value(custom_value)?.id();
let _ = self.cache.remove(Some(engine), &id);
let _ = self.cache.remove(Some(engine), &id, false);
}
Ok(())
}
@ -178,9 +180,8 @@ impl Plugin for PolarsPlugin {
pub mod test {
use super::*;
use crate::values::PolarsPluginObject;
use nu_command::IntoDatetime;
use nu_plugin_test_support::PluginTest;
use nu_protocol::{ShellError, Span};
use nu_protocol::{engine::Command, ShellError, Span};
impl PolarsPlugin {
/// Creates a new polars plugin in test mode
@ -193,6 +194,13 @@ pub mod test {
}
pub fn test_polars_plugin_command(command: &impl PluginCommand) -> Result<(), ShellError> {
test_polars_plugin_command_with_decls(command, vec![])
}
pub fn test_polars_plugin_command_with_decls(
command: &impl PluginCommand,
decls: Vec<Box<dyn Command>>,
) -> Result<(), ShellError> {
let plugin = PolarsPlugin::new_test_mode();
let examples = command.examples();
@ -210,9 +218,12 @@ pub mod test {
}
}
PluginTest::new("polars", plugin.into())?
.add_decl(Box::new(IntoDatetime))?
.test_examples(&examples)?;
let mut plugin_test = PluginTest::new("polars", plugin.into())?;
for decl in decls {
let _ = plugin_test.add_decl(decl)?;
}
plugin_test.test_examples(&examples)?;
Ok(())
}