From 98369985b1252fe18cd50c0f87d38c09d7fb474a Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Mon, 13 May 2024 16:17:31 -0700
Subject: [PATCH] Allow custom value operations to work on eager and lazy
 dataframes interchangeably.  (#12819)

Fixes Bug #12809

The example that @maxim-uvarov posted now works as expected:

<img width="1223" alt="Screenshot 2024-05-09 at 16 21 01"
src="https://github.com/nushell/nushell/assets/56345/a4df62e3-e432-4c09-8e25-9a6c198741a3">
---
 .../values/nu_dataframe/operations.rs         |  2 +-
 .../values/nu_lazyframe/custom_value.rs       | 48 +++++++++++++++++--
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/operations.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/operations.rs
index 42d803b0e4..db5409aff2 100644
--- a/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/operations.rs
+++ b/crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/operations.rs
@@ -27,7 +27,7 @@ impl NuDataFrame {
         let rhs_span = right.span();
         match right {
             Value::Custom { .. } => {
-                let rhs = NuDataFrame::try_from_value(plugin, right)?;
+                let rhs = NuDataFrame::try_from_value_coerce(plugin, right, rhs_span)?;
 
                 match (self.is_series(), rhs.is_series()) {
                     (true, true) => {
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 731d210dd8..b3cd9500dd 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
@@ -7,7 +7,7 @@ use uuid::Uuid;
 
 use crate::{
     values::{CustomValueSupport, NuDataFrame, PolarsPluginCustomValue},
-    PolarsPlugin,
+    Cacheable, PolarsPlugin,
 };
 
 use super::NuLazyFrame;
@@ -76,11 +76,49 @@ impl PolarsPluginCustomValue for NuLazyFrameCustomValue {
         _engine: &EngineInterface,
         other_value: Value,
     ) -> Result<Option<Ordering>, ShellError> {
-        // to compare, we need to convert to NuDataframe
-        let df = NuLazyFrame::try_from_custom_value(plugin, self)?;
-        let df = df.collect(other_value.span())?;
+        let eager = NuLazyFrame::try_from_custom_value(plugin, self)?.collect(Span::unknown())?;
         let other = NuDataFrame::try_from_value_coerce(plugin, &other_value, other_value.span())?;
-        let res = df.is_equal(&other);
+        let res = eager.is_equal(&other);
         Ok(res)
     }
+
+    fn custom_value_operation(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        lhs_span: Span,
+        operator: nu_protocol::Spanned<nu_protocol::ast::Operator>,
+        right: Value,
+    ) -> Result<Value, ShellError> {
+        let eager = NuLazyFrame::try_from_custom_value(plugin, self)?.collect(Span::unknown())?;
+        Ok(eager
+            .compute_with_value(plugin, lhs_span, operator.item, operator.span, &right)?
+            .cache(plugin, engine, lhs_span)?
+            .into_value(lhs_span))
+    }
+
+    fn custom_value_follow_path_int(
+        &self,
+        plugin: &PolarsPlugin,
+        _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)
+    }
+
+    fn custom_value_follow_path_string(
+        &self,
+        plugin: &PolarsPlugin,
+        engine: &EngineInterface,
+        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))
+    }
 }