From 3ea73121c923b2c02406b802115663ca47738082 Mon Sep 17 00:00:00 2001
From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com>
Date: Mon, 30 Sep 2024 20:43:42 +0100
Subject: [PATCH] fix lazy dtype in arg_previous_greater

---
 Cargo.lock                  |  2 +-
 src/arg_previous_greater.rs |  1 +
 src/expressions.rs          | 12 ++++++++++--
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ac1838f..8939a43 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1268,7 +1268,7 @@ dependencies = [
 
 [[package]]
 name = "polars_xdt"
-version = "0.15.3"
+version = "0.16.0"
 dependencies = [
  "chrono",
  "chrono-tz 0.9.0",
diff --git a/src/arg_previous_greater.rs b/src/arg_previous_greater.rs
index 2110a0d..ff8667c 100644
--- 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
index 5228e94..d84213c 100644
--- 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() {