feat: enhance WindowFrame to accept both scalar and array inputs in bounds

Author: kosiew

python/datafusion/expr.py

@@ -24,8 +24,6 @@
 
 from typing import TYPE_CHECKING, Any, Optional
 
-import pyarrow as pa
-
 try:
     from warnings import deprecated  # Python 3.13+
 except ImportError:
@@ -1206,12 +1204,20 @@ def __init__(
                 will be set to unbounded. If unit type is ``groups``, this
                 parameter must be set.
         """
-        if not isinstance(start_bound, pa.Scalar) and start_bound is not None:
-            start_bound = pa.scalar(start_bound)
+        if start_bound is not None:
+            if not isinstance(start_bound, (pa.Array, pa.ChunkedArray)):
+                if isinstance(start_bound, pa.Scalar):
+                    start_bound = pa.array([start_bound.as_py()])
+                else:
+                    start_bound = pa.array([start_bound])
             if units in ("rows", "groups"):
                 start_bound = start_bound.cast(pa.uint64())
-        if not isinstance(end_bound, pa.Scalar) and end_bound is not None:
-            end_bound = pa.scalar(end_bound)
+        if end_bound is not None:
+            if not isinstance(end_bound, (pa.Array, pa.ChunkedArray)):
+                if isinstance(end_bound, pa.Scalar):
+                    end_bound = pa.array([end_bound.as_py()])
+                else:
+                    end_bound = pa.array([end_bound])
             if units in ("rows", "groups"):
                 end_bound = end_bound.cast(pa.uint64())
         self.window_frame = expr_internal.WindowFrame(units, start_bound, end_bound)

python/tests/test_dataframe.py

@@ -767,6 +767,11 @@ def test_window_functions(partitioned_df, name, expr, result):
     ]
     + [
         ("groups", 0, 0),
+        (
+            "rows",
+            pa.array([0], type=pa.uint64()),
+            pa.array([1], type=pa.uint64()),
+        ),
     ],
 )
 def test_valid_window_frame(units, start_bound, end_bound):

python/tests/test_udwf.py

@@ -312,6 +312,18 @@ def test_register_udwf(ctx, count_window_df):
     assert result.column(0) == pa.array([0, 1, 2])
 
 
+def test_window_frame_accepts_scalar_and_none():
+    wf = WindowFrame("rows", pa.scalar(1), None)
+    assert wf.get_lower_bound().get_offset() == 1
+    assert wf.get_upper_bound().is_unbounded()
+
+
+def test_window_frame_accepts_arrays():
+    wf = WindowFrame("rows", pa.array([1]), pa.array([2]))
+    assert wf.get_lower_bound().get_offset() == 1
+    assert wf.get_upper_bound().get_offset() == 2
+
+
 smooth_default = udwf(
     ExponentialSmoothDefault,
     pa.float64(),