NVIDIA / NVIDIA/cudf

[BUG] `IR`s that `broadcast` with `LiterColumn` expressions can fail when `max_rows_per_partition` is set with rapidsmpf

Open
#22,072 2 comments 0 reactions 0 assignees View on GitHub
bug cudf-polars
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

```python
from functools import partialmethod
import polars
from cudf_polars.utils.config import Runtime, StreamingFallbackMode
executor = "streaming"
executor_options: dict[str, Any] = {}
executor_options["max_rows_per_partition"] = 4
executor_options["target_partition_size"] = 10
# We expect many tests to fall back, so silence the warnings
executor_options["fallback_mode"] = StreamingFallbackMode.SILENT
executor_options["runtime"] = Runtime.RAPIDSMPF
collect = polars.LazyFrame.collect
engine = polars.GPUEngine(executor=executor, executor_options=executor_options)
polars.LazyFrame.collect = partialmethod(collect, engine=engine)

import polars as pl
partition_limit = 15
df = pl.DataFrame([0, 0] * partition_limit)
groups = pl.Series([0, 1] * partition_limit)
df.group_by(groups).agg(pl.all().is_not_null().sum())
```

results in an exception in `broadcast` in the IR where the `max_rows_per_partition` (4) is not broadcast-able to the total length (30)

```python

File /cudf/python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py:702, in groupby_actor(context, comm, ir, ir_context, ch_out, ch_in, target_partition_size, collective_ids)
699 decomposed = DecomposedGroupBy.from_ir(ir)
700 assert not decomposed.need_preshuffle, "Should already be shuffled."
--> 702 aggregated, input_drained, chunks_received = await _local_aggregation(
703 context,
704 decomposed,
705 ir_context,
706 ch_in,
707 target_partition_size,
708 allow_early_exit=not require_tree,
709 )
711 skip_global_comm = metadata_in.duplicated or partitioned_inter_rank
712 output_count = await _choose_strategy(
713 context,
714 comm,
(...) 722 tracer,
723 )

File /cudf/python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py:217, in _local_aggregation(context, decomposed, ir_context, ch_in, target_partition_size, allow_early_exit)
214 break
216 chunks_received += 1
--> 217 chunk = await evaluate_chunk(
218 context,
219 TableChunk.from_message(msg),
220 decomposed.piecewise_ir,
221 ir_context=ir_context,
222 )
223 chunk = _enforce_schema(chunk, decomposed.piecewise_ir.schema)
224 total_size += chunk.data_alloc_size()

File /cudf/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py:382, in evaluate_chunk(context, chunk, ir_context, *irs)
380 with opaque_memory_usage(extra):
381 for single_ir in irs:
--> 382 chunk = await asyncio.to_thread(
383 _evaluate_chunk_sync, chunk, single_ir, ir_context
384 )
385 return chunk

File /conda/envs/cudf-dev-polars-tests/lib/python3.14/asyncio/threads.py:25, in to_thread(func, *args, **kwargs)
23 ctx = contextvars.copy_context()
24 func_call = functools.partial(ctx.run, func, *args, **kwargs)
---> 25 return await loop.run_in_executor(None, func_call)

File /conda/envs/cudf-dev-polars-tests/lib/python3.14/concurrent/futures/thread.py:86, in _WorkItem.run(self, ctx)
83 return
85 try:
---> 86 result = ctx.run(self.task)
87 except BaseException as exc:
88 self.future.set_exception(exc)

File /conda/envs/cudf-dev-polars-tests/lib/python3.14/concurrent/futures/thread.py:73, in WorkerContext.run(self, task)
71 def run(self, task):
72 fn, args, kwargs = task
---> 73 return fn(*args, **kwargs)

File /cudf/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py:340, in _evaluate_chunk_sync(chunk, ir, ir_context)
338 names = list(input_schema.keys())
339 dtypes = list(input_schema.values())
--> 340 df = ir.do_evaluate(
341 *ir._non_child_args,
342 DataFrame.from_table(chunk.table_view(), names, dtypes, chunk.stream),
343 context=ir_context,
344 )
345 return TableChunk.from_pylibcudf_table(df.table, df.stream, exclusive_view=True)

File /cudf/python/cudf_polars/cudf_polars/dsl/ir.py:1867, in GroupBy.do_evaluate(cls, schema, keys_in, agg_requests, maintain_order, zlice, df, context)
1852 @classmethod
1853 @log_do_evaluate
1854 @nvtx_annotate_cudf_polars(message="GroupBy")
(...) 1864 context: IRExecutionContext,
1865 ) -> DataFrame:
1866 """Evaluate and return a dataframe."""
-> 1867 keys = broadcast(
1868 *(k.evaluate(df) for k in keys_in),
1869 target_length=df.num_rows,
1870 stream=df.stream,
1871 )
1872 keys_are_sorted = (
1873 plc.types.Sorted.YES
1874 if all(k.is_sorted for k in keys)
1875 else plc.types.Sorted.NO
1876 )
1877 grouper = plc.groupby.GroupBy(
1878 plc.Table([k.obj for k in keys]),
1879 null_handling=plc.types.NullPolicy.INCLUDE,
(...) 1882 null_precedence=[k.null_order for k in keys],
1883 )

File /cudf/python/cudf_polars/cudf_polars/dsl/utils/reshape.py:69, in broadcast(target_length, stream, *columns)
67 raise RuntimeError("Mismatching column lengths") from e
68 if target_length is not None and nrows != target_length:
---> 69 raise RuntimeError(
70 f"Cannot broadcast columns of length {nrows=} to {target_length=}"
71 )
72 return [
73 column
74 if column.size != 1
(...) 85 for column in columns
86 ]

RuntimeError: Cannot broadcast columns of length nrows=30 to target_length=4
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.