ext/polars/src/lazy/dsl.rs in polars-df-0.1.2 vs ext/polars/src/lazy/dsl.rs in polars-df-0.1.3

@@ -1,16 +1,17 @@
-use magnus::{RArray, RString, Value};
+use magnus::{class, RArray, RString, Value};
 use polars::chunked_array::ops::SortOptions;
 use polars::lazy::dsl;
 use polars::lazy::dsl::Operator;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 
 use crate::conversion::*;
 use crate::lazy::apply::*;
 use crate::lazy::utils::rb_exprs_to_exprs;
-use crate::RbResult;
+use crate::utils::reinterpret;
+use crate::{RbResult, RbSeries};
 
 #[magnus::wrap(class = "Polars::RbExpr")]
 #[derive(Clone)]
 pub struct RbExpr {
     pub inner: dsl::Expr,
@@ -884,10 +885,23 @@
 
     pub fn dt_round(&self, every: String, offset: String) -> Self {
         self.inner.clone().dt().round(&every, &offset).into()
     }
 
+    pub fn reinterpret(&self, signed: bool) -> Self {
+        let function = move |s: Series| reinterpret(&s, signed);
+        let dt = if signed {
+            DataType::Int64
+        } else {
+            DataType::UInt64
+        };
+        self.clone()
+            .inner
+            .map(function, GetOutput::from_type(dt))
+            .into()
+    }
+
     pub fn mode(&self) -> Self {
         self.inner.clone().mode().into()
     }
 
     pub fn keep_name(&self) -> Self {
@@ -1303,10 +1317,25 @@
             min_periods,
         };
         self.inner.clone().ewm_var(options).into()
     }
 
+    pub fn extend_constant(&self, value: Wrap<AnyValue>, n: usize) -> Self {
+        let value = Value::from(value);
+        self.inner
+            .clone()
+            .apply(
+                move |s| {
+                    let value = value.try_convert::<Wrap<AnyValue>>().unwrap().0;
+                    s.extend_constant(value, n)
+                },
+                GetOutput::same_type(),
+            )
+            .with_fmt("extend")
+            .into()
+    }
+
     pub fn any(&self) -> Self {
         self.inner.clone().any().into()
     }
 
     pub fn all(&self) -> Self {
@@ -1330,10 +1359,14 @@
     }
 
     pub fn exp(&self) -> Self {
         self.inner.clone().exp().into()
     }
+
+    pub fn entropy(&self, base: f64, normalize: bool) -> Self {
+        self.inner.clone().entropy(base, normalize).into()
+    }
 }
 
 pub fn col(name: String) -> RbExpr {
     dsl::col(&name).into()
 }
@@ -1363,18 +1396,42 @@
 
 // TODO improve
 pub fn lit(value: Value) -> RbResult<RbExpr> {
     if value.is_nil() {
         Ok(dsl::lit(Null {}).into())
+    } else if let Ok(series) = value.try_convert::<&RbSeries>() {
+        Ok(dsl::lit(series.series.borrow().clone()).into())
     } else if let Some(v) = RString::from_value(value) {
         Ok(dsl::lit(v.try_convert::<String>()?).into())
+    } else if value.is_kind_of(class::integer()) {
+        match value.try_convert::<i64>() {
+            Ok(val) => {
+                if val > 0 && val < i32::MAX as i64 || val < 0 && val > i32::MIN as i64 {
+                    Ok(dsl::lit(val as i32).into())
+                } else {
+                    Ok(dsl::lit(val).into())
+                }
+            }
+            _ => {
+                let val = value.try_convert::<u64>()?;
+                Ok(dsl::lit(val).into())
+            }
+        }
     } else {
         Ok(dsl::lit(value.try_convert::<f64>()?).into())
     }
 }
 
 pub fn arange(low: &RbExpr, high: &RbExpr, step: usize) -> RbExpr {
     polars::lazy::dsl::arange(low.inner.clone(), high.inner.clone(), step).into()
+}
+
+pub fn repeat(value: Value, n_times: &RbExpr) -> RbResult<RbExpr> {
+    if value.is_nil() {
+        Ok(polars::lazy::dsl::repeat(Null {}, n_times.inner.clone()).into())
+    } else {
+        todo!();
+    }
 }
 
 #[magnus::wrap(class = "Polars::RbWhen")]
 #[derive(Clone)]
 pub struct RbWhen {