Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions polars_vec_ops/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,17 +270,31 @@ 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)
clipped_series = df_work.select(_TEMP_VAL).select(
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, ...
Expand Down
6 changes: 3 additions & 3 deletions tests/test_join_between.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
Loading