extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-08-13 15:27:30 +02:00
Code Issues Packages Projects Releases Wiki Activity

Creates col and nth expressions when using paths on lazy frames. (#15891)

Browse Source 
# Description
Instead of collecting the frame and returning the columns of the
collected frame when using paths $df.column_name or $df.0 this creates
expressions:

```nu
> ❯ : let df = polars open /tmp/foo.parquet
> ❯ : $df | polars select $df.pid | polars first 5 | polars collect
╭───┬───────╮
│ # │  pid  │
├───┼───────┤
│ 0 │ 45280 │
│ 1 │ 45252 │
│ 2 │ 45242 │
│ 3 │ 45241 │
│ 4 │ 45207 │
╰───┴───────╯

> ❯ : $df | polars select $df.0 | polars first 5 | polars collect
╭───┬───────╮
│ # │  pid  │
├───┼───────┤
│ 0 │ 45280 │
│ 1 │ 45252 │
│ 2 │ 45242 │
│ 3 │ 45241 │
│ 4 │ 45207 │
╰───┴───────╯
```

---------

Co-authored-by: Jack Wright <jack.wright@nike.com>
This commit is contained in:
Jack Wright
2025-06-05 12:42:27 -07:00
committed by GitHub
parent 9cc74e7a9f
commit 4a9e2ac37b
1 changed files with 8 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/custom_value.rs
View File

@ -2,12 +2,13 @@ use std::cmp::Ordering;
use nu_plugin::EngineInterface; use nu_plugin::EngineInterface;
use nu_protocol::{CustomValue, ShellError, Span, Value}; use nu_protocol::{CustomValue, ShellError, Span, Value};
use polars::prelude::{col, nth};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use uuid::Uuid; use uuid::Uuid;
use crate::{ use crate::{
Cacheable, PolarsPlugin, Cacheable, PolarsPlugin,
values::{CustomValueSupport, NuDataFrame, PolarsPluginCustomValue}, values::{CustomValueSupport, NuDataFrame, NuExpression, PolarsPluginCustomValue},
}; };
use super::NuLazyFrame; use super::NuLazyFrame;
@ -100,25 +101,22 @@ impl PolarsPluginCustomValue for NuLazyFrameCustomValue {
fn custom_value_follow_path_int( fn custom_value_follow_path_int(
&self, &self,
plugin: &PolarsPlugin, plugin: &PolarsPlugin,
_engine: &EngineInterface, engine: &EngineInterface,
_self_span: Span, _self_span: Span,
index: nu_protocol::Spanned<usize>, index: nu_protocol::Spanned<usize>,
) -> Result<Value, ShellError> { ) -> Result<Value, ShellError> {
let eager = NuLazyFrame::try_from_custom_value(plugin, self)?.collect(Span::unknown())?; let expr = NuExpression::from(nth(index.item as i64));
eager.get_value(index.item, index.span) expr.cache_and_to_value(plugin, engine, index.span)
} }
fn custom_value_follow_path_string( fn custom_value_follow_path_string(
&self, &self,
plugin: &PolarsPlugin, plugin: &PolarsPlugin,
engine: &EngineInterface, engine: &EngineInterface,
self_span: Span, _self_span: Span,
column_name: nu_protocol::Spanned<String>, column_name: nu_protocol::Spanned<String>,
) -> Result<Value, ShellError> { ) -> Result<Value, ShellError> {
let eager = NuLazyFrame::try_from_custom_value(plugin, self)?.collect(Span::unknown())?; let expr = NuExpression::from(col(column_name.item));
let column = eager.column(&column_name.item, self_span)?; expr.cache_and_to_value(plugin, engine, column_name.span)
Ok(column
.cache(plugin, engine, self_span)?
.into_value(self_span))
} }
} }