From 4a9e2ac37b2c2a36d839fd49dbec125781eac5da Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Thu, 5 Jun 2025 12:42:27 -0700
Subject: [PATCH] Creates col and nth expressions when using paths on lazy
 frames. (#15891)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# 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>
---
 .../values/nu_lazyframe/custom_value.rs        | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

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