diff --git a/polars_vec_ops/frame.py b/polars_vec_ops/frame.py index ac4b59a..6be2862 100644 --- a/polars_vec_ops/frame.py +++ b/polars_vec_ops/frame.py @@ -270,10 +270,8 @@ def _is_sorted(val: Any) -> bool: kwargs={"relative": relative}, ) - # Post-process: convert to counts or cast to correct output dtype - if as_counts: - cross_clip_expr = cross_clip_expr.list.len().cast(pl.UInt32) - elif return_dtype != pl.List(pl.Float64): + # Post-process: cast to correct output dtype + if not as_counts and return_dtype != pl.List(pl.Float64): cross_clip_expr = cross_clip_expr.cast(return_dtype) # Apply cross_clip to extract only _TEMP_VAL column (avoid gathering large lists) @@ -281,6 +279,22 @@ def _is_sorted(val: Any) -> bool: cross_clip_expr.alias(val_col_name) )[val_col_name] + # Convert list results to counts after clipping, preserving nulls. + # This is done after the clip rather than inside the expression because + # list.len() returns 0 for null lists in older polars versions. + if as_counts: + clipped_series = ( + clipped_series + .to_frame() + .select( + pl.when(pl.col(val_col_name).is_null()) + .then(pl.lit(None).cast(pl.UInt32)) + .otherwise(pl.col(val_col_name).list.len().cast(pl.UInt32)) + .alias(val_col_name) + ) + [val_col_name] + ) + # ── Build the output DataFrame efficiently ──────────────────────── # We need n_units × n_intervals rows in order: # unit0×int0, unit0×int1, ..., unit0×intN, unit1×int0, ... diff --git a/tests/test_join_between.py b/tests/test_join_between.py index f896edb..d466f46 100644 --- a/tests/test_join_between.py +++ b/tests/test_join_between.py @@ -307,12 +307,12 @@ def test_join_between_as_counts_retains_columns(units, intervals): def test_join_between_vec_as_list_expr(units, intervals): """vec accepts arbitrary list expressions; clipping uses the transformed list. - Adding 1 to the list shifts all values up by 1, producing verifiably - different clipping results than the original column. + Adding 1 to each element of the list shifts all values up by 1, producing + verifiably different clipping results than the original column. """ result = units.vec.join_between( other=intervals, - values=pl.col("event_times") + 1, + values=pl.col("event_times").list.eval(pl.element() + 1), bounds=("start_time", "stop_time"), ).sort("event_id", "interval_id")