NVIDIA / NVIDIA/cudf

[BUG] MergeSorted IR assumes left/right children are Sorts in cudf_polars

Open
#23,055 0 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
import polars as pl
lf = pl.LazyFrame({"foo": [1, 3, 5], "bar": [1, 2, 3]}).sort("foo")
lf2 = pl.LazyFrame({"foo": [2, 4, 6], "bar": [4, 5, 6]}).sort("foo")
engine = pl.GPUEngine(
executor="streaming",
executor_options={"max_rows_per_partition": 2},
)
q = lf.merge_sorted(lf2, key="foo").select(pl.col("bar") + 1)
q.collect(engine=engine)

/cudf/python/cudf_polars/cudf_polars/streaming/utils.py:92: UserWarning: Class does not support multiple partitions.
_fallback_inform(msg, rec.state["config_options"])
---------------------------------------------------------------------------
KeyError Traceback (most recent call last)
File /cudf/python/cudf_polars/cudf_polars/dsl/traversal.py:222, in CachingVisitor.__call__(self, value)
221 try:
--> 222 return self.cache[value]
223 except KeyError:

KeyError: Select({'bar': )>}, (NamedExpr(bar, BinOp()>, , Col()>, 'bar'), Literal()>, 1))),), True, Projection({'bar': )>}, MergeSorted({'bar': )>, 'foo': )>}, 'foo', Projection({'bar': )>, 'foo': )>}, Sort({'foo': )>, 'bar': )>}, (NamedExpr(foo, Col()>, 'foo')),), (,), (,), False, None, DataFrameScan({'foo': )>, 'bar': )>}, , None))), Projection({'bar': )>, 'foo': )>}, Sort({'foo': )>, 'bar': )>}, (NamedExpr(foo, Col()>, 'foo')),), (,), (,), False, None, DataFrameScan({'foo': )>, 'bar': )>}, , None))))))

During handling of the above exception, another exception occurred:

KeyError Traceback (most recent call last)
File /cudf/python/cudf_polars/cudf_polars/dsl/traversal.py:222, in CachingVisitor.__call__(self, value)
221 try:
--> 222 return self.cache[value]
223 except KeyError:

KeyError: Projection({'bar': )>}, MergeSorted({'bar': )>, 'foo': )>}, 'foo', Projection({'bar': )>, 'foo': )>}, Sort({'foo': )>, 'bar': )>}, (NamedExpr(foo, Col()>, 'foo')),), (,), (,), False, None, DataFrameScan({'foo': )>, 'bar': )>}, , None))), Projection({'bar': )>, 'foo': )>}, Sort({'foo': )>, 'bar': )>}, (NamedExpr(foo, Col()>, 'foo')),), (,), (,), False, None, DataFrameScan({'foo': )>, 'bar': )>}, , None)))))

During handling of the above exception, another exception occurred:

KeyError Traceback (most recent call last)
File /cudf/python/cudf_polars/cudf_polars/dsl/traversal.py:222, in CachingVisitor.__call__(self, value)
221 try:
--> 222 return self.cache[value]
223 except KeyError:

KeyError: MergeSorted({'bar': )>, 'foo': )>}, 'foo', Projection({'bar': )>, 'foo': )>}, Sort({'foo': )>, 'bar': )>}, (NamedExpr(foo, Col()>, 'foo')),), (,), (,), False, None, DataFrameScan({'foo': )>, 'bar': )>}, , None))), Projection({'bar': )>, 'foo': )>}, Sort({'foo': )>, 'bar': )>}, (NamedExpr(foo, Col()>, 'foo')),), (,), (,), False, None, DataFrameScan({'foo': )>, 'bar': )>}, , None))))

During handling of the above exception, another exception occurred:

AssertionError Traceback (most recent call last)
Cell In[5], line 2
1 q = lf.merge_sorted(lf2, key="foo").select(pl.col("bar") + 1)
----> 2 q.collect(engine=engine)

File /miniforge3/envs/cudf-dev/lib/python3.14/site-packages/polars/_utils/deprecation.py:97, in deprecate_streaming_parameter..decorate..wrapper(*args, **kwargs)
93 kwargs["engine"] = "in-memory"
95 del kwargs["streaming"]
---> 97 return function(*args, **kwargs)

File /miniforge3/envs/cudf-dev/lib/python3.14/site-packages/polars/lazyframe/opt_flags.py:343, in forward_old_opt_flags..decorate..wrapper(*args, **kwargs)
340 optflags = cb(optflags, kwargs.pop(key)) # type: ignore[no-untyped-call,unused-ignore]
342 kwargs["optimizations"] = optflags
--> 343 return function(*args, **kwargs)

File /miniforge3/envs/cudf-dev/lib/python3.14/site-packages/polars/lazyframe/frame.py:2630, in LazyFrame.collect(self, type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, cluster_with_columns, collapse_joins, no_optimization, engine, background, optimizations, **_kwargs)
2626 return InProcessQuery(ldf.collect_concurrently())
2627
2628 # Only for testing purposes
2629 callback = _kwargs.get("post_opt_callback", callback)
-> 2630 return wrap_df(ldf.collect(engine, callback))

File /cudf/python/cudf_polars/cudf_polars/callback.py:320, in _callback(ir, with_columns, pyarrow_predicate, n_rows, should_time, memory_resource, config_options, timer)
313 msg = textwrap.dedent("""\
314 LazyFrame.profile() is not supported with the streaming executor.
315 To profile execution with the streaming executor, use NVIDIA
316 NSight Systems with the 'streaming' scheduler.
317 """)
318 raise NotImplementedError(msg)
--> 320 return evaluate_streaming(ir, config_options)
321 assert_never(config_options.executor)

File /cudf/python/cudf_polars/cudf_polars/streaming/parallel.py:140, in evaluate_streaming(ir, config_options)
136 _clear_source_info_cache()
138 from cudf_polars.streaming.actor_graph.core import evaluate_logical_plan
--> 140 result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False)
141 return result

File /cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py:104, in evaluate_logical_plan(ir, config_options, collect_metadata)
99 case "spmd" | "default_singleton":
100 from cudf_polars.engine.spmd import (
101 evaluate_pipeline_spmd_mode,
102 )
--> 104 result, metadata_collector = evaluate_pipeline_spmd_mode(
105 ir,
106 config_options,
107 collect_metadata=collect_metadata,
108 query_id=query_id,
109 )
110 case "ray":
111 from cudf_polars.engine.ray import (
112 evaluate_pipeline_ray_mode,
113 )

File /cudf/python/cudf_polars/cudf_polars/engine/spmd.py:113, in evaluate_pipeline_spmd_mode(ir, config_options, collect_metadata, query_id)
110 context = config_options.executor.spmd_context.context
111 py_executor = config_options.executor.spmd_context.py_executor
--> 113 df, metadata = evaluate_on_rank(
114 context,
115 comm,
116 py_executor,
117 ir,
118 config_options,
119 query_id=query_id,
120 )
121 return df, metadata if collect_metadata else None

File /cudf/python/cudf_polars/cudf_polars/engine/core.py:689, in evaluate_on_rank(ctx, comm, py_executor, ir, config_options, query_id)
655 """
656 Evaluate a polars IR plan on a single rank.
657
(...) 686 Collected channel metadata.
687 """
688 stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor)
--> 689 ir, partition_info = lower_ir_graph(
690 ir, config_options, stats, rank=comm.rank, nranks=comm.nranks
691 )
693 if comm.rank == 0:
694 # At least for now, the query plan is identical on all ranks,
695 # so we only log it once.
696 log_query_plan(ir, config_options)

File /cudf/python/cudf_polars/cudf_polars/streaming/parallel.py:114, in lower_ir_graph(ir, config_options, stats, rank, nranks)
107 state: State = {
108 "config_options": config_options,
109 "stats": stats,
110 "rank": rank,
111 "nranks": nranks,
112 }
113 mapper: LowerIRTransformer = CachingVisitor(lower_ir_node, state=state)
--> 114 return mapper(ir)

File /cudf/python/cudf_polars/cudf_polars/dsl/traversal.py:224, in CachingVisitor.__call__(self, value)
222 return self.cache[value]
223 except KeyError:
--> 224 return self.cache.setdefault(value, self.fn(value, self))

File /miniforge3/envs/cudf-dev/lib/python3.14/functools.py:982, in singledispatch..wrapper(*args, **kw)
979 if not args:
980 raise TypeError(f'{funcname} requires at least '
981 '1 positional argument')
--> 982 return dispatch(args[0].__class__)(*args, **kw)

File /cudf/python/cudf_polars/cudf_polars/streaming/select.py:397, in _(ir, rec)
389 elif all(isinstance(e.value, Col) and e.name == e.value.name for e in ir.exprs):
390 # Fast path: Equivalent to Projection.
391 # We can only do this if the child is NOT HStack.
392 return lower_ir_node(
393 Projection(ir.schema, ir.children[0]),
394 rec,
395 )
--> 397 child, partition_info = rec(ir.children[0])
398 pi = partition_info[child]
400 config_options = rec.state["config_options"]

File /cudf/python/cudf_polars/cudf_polars/dsl/traversal.py:224, in CachingVisitor.__call__(self, value)
222 return self.cache[value]
223 except KeyError:
--> 224 return self.cache.setdefault(value, self.fn(value, self))

File /miniforge3/envs/cudf-dev/lib/python3.14/functools.py:982, in singledispatch..wrapper(*args, **kw)
979 if not args:
980 raise TypeError(f'{funcname} requires at least '
981 '1 positional argument')
--> 982 return dispatch(args[0].__class__)(*args, **kw)

File /cudf/python/cudf_polars/cudf_polars/streaming/select.py:95, in _(ir, rec)
92 # Partition-wise default - Import here to avoid circular import
93 from cudf_polars.streaming.parallel import _lower_ir_pwise
---> 95 return _lower_ir_pwise(ir, rec, preserve_partitioning=True)

File /cudf/python/cudf_polars/cudf_polars/streaming/parallel.py:199, in _lower_ir_pwise(ir, rec, preserve_partitioning)
193 def _lower_ir_pwise(
194 ir: IR, rec: LowerIRTransformer, *, preserve_partitioning: bool = False
195 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
196 # Lower a partition-wise (i.e. embarrassingly-parallel) IR node
197
198 # Lower children
--> 199 children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
200 partition_info = reduce(operator.or_, _partition_info)
201 counts = {partition_info[c].count for c in children}

File /cudf/python/cudf_polars/cudf_polars/streaming/parallel.py:199, in (.0)
193 def _lower_ir_pwise(
194 ir: IR, rec: LowerIRTransformer, *, preserve_partitioning: bool = False
195 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
196 # Lower a partition-wise (i.e. embarrassingly-parallel) IR node
197
198 # Lower children
--> 199 children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
200 partition_info = reduce(operator.or_, _partition_info)
201 counts = {partition_info[c].count for c in children}

File /cudf/python/cudf_polars/cudf_polars/dsl/traversal.py:224, in CachingVisitor.__call__(self, value)
222 return self.cache[value]
223 except KeyError:
--> 224 return self.cache.setdefault(value, self.fn(value, self))

File /miniforge3/envs/cudf-dev/lib/python3.14/functools.py:982, in singledispatch..wrapper(*args, **kw)
979 if not args:
980 raise TypeError(f'{funcname} requires at least '
981 '1 positional argument')
--> 982 return dispatch(args[0].__class__)(*args, **kw)

File /cudf/python/cudf_polars/cudf_polars/streaming/parallel.py:63, in _(ir, rec)
58 @lower_ir_node.register(IR)
59 def _(
60 ir: IR, rec: LowerIRTransformer
61 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: # pragma: no cover
62 # Default logic - Requires single partition
---> 63 return _lower_ir_fallback(
64 ir, rec, msg=f"Class {type(ir)} does not support multiple partitions."
65 )

File /cudf/python/cudf_polars/cudf_polars/streaming/utils.py:95, in _lower_ir_fallback(ir, rec, msg)
92 _fallback_inform(msg, rec.state["config_options"])
94 # Reconstruct and return
---> 95 new_node = ir.reconstruct(children)
96 partition_info[new_node] = PartitionInfo(count=1)
97 return new_node, partition_info

File /cudf/python/cudf_polars/cudf_polars/dsl/nodebase.py:79, in Node.reconstruct(self, children)
66 def reconstruct(self, children: Sequence[T]) -> Self:
67 """
68 Rebuild this node with new children.
69
(...) 77 New node with new children. Non-child data is shared with the input.
78 """
---> 79 return type(self)(*self._ctor_arguments(children))

File /cudf/python/cudf_polars/cudf_polars/dsl/ir.py:2962, in MergeSorted.__init__(self, schema, key, left, right)
2960 left_sort_child = left if isinstance(left, Sort) else left.children[0]
2961 right_sort_child = right if isinstance(right, Sort) else right.children[0]
-> 2962 assert isinstance(left_sort_child, Sort)
2963 assert isinstance(right_sort_child, Sort)
2964 assert left_sort_child.order == right_sort_child.order

AssertionError:
```

Discovered in https://github.com/rapidsai/cudf/pull/23015

Contributor guide

Open the contributing guide

Research direction

Reproduce the reported LazyFrame.merge_sorted query with the streaming executor and max_rows_per_partition set to 2. Start in cudf_polars/streaming/select.py around the child lowering path, then inspect cudf_polars/dsl/traversal.py and cudf_polars/streaming/utils.py. Done means the query completes without the KeyError or AssertionError while preserving the expected merged and selected result.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.