fix lazy dtype in arg_previous_greater

pola-rs · Sep 30, 2024 · 3ea7312 · 3ea7312
1 parent df2d46e
commit 3ea7312
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/arg_previous_greater.rs b/src/arg_previous_greater.rs
@@ -1,5 +1,6 @@
 use polars::prelude::*;
 
+
 pub(crate) fn impl_arg_previous_greater<T>(ca: &ChunkedArray<T>) -> IdxCa
 where
     T: PolarsNumericType,

diff --git a/src/expressions.rs b/src/expressions.rs
@@ -83,8 +83,16 @@ fn to_julian_date(inputs: &[Series]) -> PolarsResult<Series> {
     impl_to_julian_date(s)
 }
 
-// todo: can we make this bigidx-dependent?
-#[polars_expr(output_type=UInt32)]
+
+fn list_idx_dtype(input_fields: &[Field]) -> PolarsResult<Field> {
+    let field = Field::new(
+        input_fields[0].name.clone(),
+        DataType::List(Box::new(IDX_DTYPE))
+    );
+    Ok(field.clone())
+}
+
+#[polars_expr(output_type_func=list_idx_dtype)]
 fn arg_previous_greater(inputs: &[Series]) -> PolarsResult<Series> {
     let ser = &inputs[0];
     match ser.dtype() {