From 54fecb99bea5ed7a019c45764b9a0b9176fdad31 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 14 Aug 2026 08:17:03 -0700 Subject: [PATCH 1/6] basic in-memory support --- python/cudf_polars/cudf_polars/dsl/ir.py | 32 +++++++++++-- .../cudf_polars/streaming/parallel.py | 5 ++ .../tests/expressions/test_sort.py | 11 ++--- python/cudf_polars/tests/test_mapfunction.py | 48 +++++++++++++++++-- 4 files changed, 81 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 559777f43c7d..76ebec87fee7 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -55,7 +55,7 @@ offsets_to_windows, range_window_bounds, ) -from cudf_polars.utils import dtypes +from cudf_polars.utils import dtypes, sorting from cudf_polars.utils.cuda_stream import ( get_cuda_stream, stream_ordered_after, @@ -3399,10 +3399,14 @@ def __init__(self, schema: Schema, name: str, options: Any, df: IR): raise NotImplementedError( "Fast count unsupported for CSV scans" ) # pragma: no cover - elif ( - self.name == "hint_sorted" - ): # pragma: no cover; polars prunes hints in some cases - raise NotImplementedError("Hint sorted unsupported") + elif self.name == "hint_sorted": + (sorted_info,) = options + self.options = ( + tuple( + (name, bool(descending), bool(nulls_last)) + for name, descending, nulls_last in sorted_info + ), + ) self._non_child_args = (schema, name, self.options) def get_hashable(self) -> Hashable: @@ -3512,6 +3516,24 @@ def do_evaluate( dtype=dtype, ) return DataFrame([index_col, *df.columns], stream=df.stream) + elif name == "hint_sorted": + (sorted_info,) = options + column_names, descending, nulls_last = zip(*sorted_info, strict=True) + orders, null_orders = sorting.sort_order( + descending, + nulls_last=nulls_last, + num_keys=len(column_names), + ) + result = DataFrame([col.copy() for col in df.columns], stream=df.stream) + for column_name, order, null_order in zip( + column_names, orders, null_orders, strict=True + ): + result.column_map[column_name].set_sorted( + is_sorted=plc.types.Sorted.YES, + order=order, + null_order=null_order, + ) + return result else: raise AssertionError("Should never be reached") # pragma: no cover diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 2ac5c8c2eef7..794ea08b6acb 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -302,6 +302,11 @@ def _( def _( ir: MapFunction, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: + if ir.name == "hint_sorted": + raise NotImplementedError( + "hint_sorted is not supported by the streaming executor." + ) + # Allow pointwise operations if ir.name in ("rename", "explode"): return _lower_ir_pwise(ir, rec) diff --git a/python/cudf_polars/tests/expressions/test_sort.py b/python/cudf_polars/tests/expressions/test_sort.py index ccc3049e5df3..0840568a718a 100644 --- a/python/cudf_polars/tests/expressions/test_sort.py +++ b/python/cudf_polars/tests/expressions/test_sort.py @@ -11,7 +11,7 @@ from cudf_polars.testing.asserts import ( assert_gpu_result_equal, ) -from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_140 +from cudf_polars.utils.versions import POLARS_VERSION_LT_136 @pytest.mark.parametrize("descending", [False, True]) @@ -69,11 +69,10 @@ def test_setsorted(engine: pl.GPUEngine, request, descending, nulls_last, with_n "fixed in https://github.com/pola-rs/polars/pull/25250" ) ) - elif not POLARS_VERSION_LT_140: - # polars >= 1.40 keeps the hint_sorted node in the optimized plan for a - # bare set_sorted; we do not support it, so it raises. 1.36-1.39 pruned - # it during optimization and passed. - request.applymarker(pytest.mark.xfail(reason="Hint sorted unsupported")) + elif engine.config.get("executor") != "in-memory": + request.applymarker( + pytest.mark.xfail(reason="Streaming hint_sorted unsupported") + ) sorted_values = sorted([1, 2, 3, 4, 5, 6, -2], reverse=descending) values: list[int | None] = [*sorted_values] if with_nulls == "nulls": diff --git a/python/cudf_polars/tests/test_mapfunction.py b/python/cudf_polars/tests/test_mapfunction.py index 3e476afc9b0e..b3f09482136d 100644 --- a/python/cudf_polars/tests/test_mapfunction.py +++ b/python/cudf_polars/tests/test_mapfunction.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -6,8 +6,10 @@ import polars as pl +import pylibcudf as plc + from cudf_polars.containers import DataType -from cudf_polars.dsl.ir import DataFrameScan, MapFunction +from cudf_polars.dsl.ir import DataFrameScan, IRExecutionContext, MapFunction from cudf_polars.dsl.translate import Translator from cudf_polars.testing.asserts import ( assert_gpu_result_equal, @@ -110,8 +112,13 @@ def test_unique_hash(): assert hash(ir_a) != hash(ir_b) -@pytest.mark.xfail(reason="HintIR not supported") -def test_set_sorted_then_inner_join(engine: pl.GPUEngine): +def test_set_sorted_then_inner_join( + engine: pl.GPUEngine, request: pytest.FixtureRequest +): + if engine.config.get("executor") != "in-memory": + request.applymarker( + pytest.mark.xfail(reason="Streaming hint_sorted unsupported") + ) df = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}) q = df.set_sorted("a").join( @@ -120,6 +127,39 @@ def test_set_sorted_then_inner_join(engine: pl.GPUEngine): assert_gpu_result_equal(q, engine=engine) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("nulls_last", [False, True]) +def test_hint_sorted_marks_column_metadata(descending, nulls_last) -> None: + schema = { + "a": DataType(pl.Int64()), + "b": DataType(pl.Int64()), + } + child = DataFrameScan( + schema, + pl.DataFrame({"a": [1, 2, 3], "b": [3, 2, 1]})._df, + None, + ) + node = MapFunction( + schema, + "hint_sorted", + [[("a", descending, nulls_last)]], + child, + ) + + result = node.evaluate(cache={}, timer=None, context=IRExecutionContext()) + + order = plc.types.Order.DESCENDING if descending else plc.types.Order.ASCENDING + null_order = ( + plc.types.NullOrder.AFTER + if descending != nulls_last + else plc.types.NullOrder.BEFORE + ) + assert result.column_map["a"].is_sorted == plc.types.Sorted.YES + assert result.column_map["a"].order == order + assert result.column_map["a"].null_order == null_order + assert result.column_map["b"].is_sorted == plc.types.Sorted.NO + + def test_explode_single_legacy_options(): # Cover the branch: POLARS_VERSION_LT_136 or len(self.options) == 1 # On polars >= 1.36 this branch is only reachable by direct construction From db7443c50bb6247ec8812e7e405f0f9cddac491e Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 14 Aug 2026 08:31:51 -0700 Subject: [PATCH 2/6] fix xfail check --- python/cudf_polars/tests/expressions/test_sort.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/expressions/test_sort.py b/python/cudf_polars/tests/expressions/test_sort.py index 0840568a718a..a4db1d2dca8c 100644 --- a/python/cudf_polars/tests/expressions/test_sort.py +++ b/python/cudf_polars/tests/expressions/test_sort.py @@ -11,7 +11,7 @@ from cudf_polars.testing.asserts import ( assert_gpu_result_equal, ) -from cudf_polars.utils.versions import POLARS_VERSION_LT_136 +from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_140 @pytest.mark.parametrize("descending", [False, True]) @@ -69,7 +69,7 @@ def test_setsorted(engine: pl.GPUEngine, request, descending, nulls_last, with_n "fixed in https://github.com/pola-rs/polars/pull/25250" ) ) - elif engine.config.get("executor") != "in-memory": + elif not POLARS_VERSION_LT_140 and engine.config.get("executor") != "in-memory": request.applymarker( pytest.mark.xfail(reason="Streaming hint_sorted unsupported") ) From 47915fd69cd3560788ad19b548d937bbd73590af Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 14 Aug 2026 12:38:42 -0700 Subject: [PATCH 3/6] adress suggestions and CI xpass --- python/cudf_polars/cudf_polars/dsl/ir.py | 17 ++++-- .../cudf_polars/tests/expressions/test_agg.py | 10 +++- .../tests/expressions/test_sort.py | 3 +- python/cudf_polars/tests/test_mapfunction.py | 55 +++++++++++++++++-- 4 files changed, 71 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 9d161a412734..d542b8ef9cf5 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -3603,11 +3603,17 @@ def __init__(self, schema: Schema, name: str, options: Any, df: IR): ) # pragma: no cover elif self.name == "hint_sorted": (sorted_info,) = options + column_names = [] + descending = [] + nulls_last = [] + for column_name, is_descending, is_nulls_last in sorted_info: + column_names.append(column_name) + descending.append(bool(is_descending)) + nulls_last.append(bool(is_nulls_last)) self.options = ( - tuple( - (name, bool(descending), bool(nulls_last)) - for name, descending, nulls_last in sorted_info - ), + tuple(column_names), + tuple(descending), + tuple(nulls_last), ) self._non_child_args = (schema, name, self.options) @@ -3719,8 +3725,7 @@ def do_evaluate( ) return DataFrame([index_col, *df.columns], stream=df.stream) elif name == "hint_sorted": - (sorted_info,) = options - column_names, descending, nulls_last = zip(*sorted_info, strict=True) + column_names, descending, nulls_last = options orders, null_orders = sorting.sort_order( descending, nulls_last=nulls_last, diff --git a/python/cudf_polars/tests/expressions/test_agg.py b/python/cudf_polars/tests/expressions/test_agg.py index b363619e192f..53a3851a2a3e 100644 --- a/python/cudf_polars/tests/expressions/test_agg.py +++ b/python/cudf_polars/tests/expressions/test_agg.py @@ -14,8 +14,10 @@ assert_gpu_result_equal, assert_ir_translation_raises, ) +from cudf_polars.testing.engine_utils import is_streaming_engine from cudf_polars.utils.versions import ( POLARS_VERSION_LT_136, + POLARS_VERSION_LT_140, ) @@ -56,12 +58,16 @@ def is_sorted(request): @pytest.fixture -def xfail_if_sorted(is_sorted, request): +def xfail_if_sorted(engine, is_sorted, request): # See https://github.com/rapidsai/cudf/pull/20791#issuecomment-3750528419 - if is_sorted: + if is_sorted and POLARS_VERSION_LT_136: request.applymarker( pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/24981") ) + elif is_sorted and not POLARS_VERSION_LT_140 and is_streaming_engine(engine): + request.applymarker( + pytest.mark.xfail(reason="Streaming hint_sorted unsupported") + ) @pytest.fixture diff --git a/python/cudf_polars/tests/expressions/test_sort.py b/python/cudf_polars/tests/expressions/test_sort.py index a4db1d2dca8c..d2145c4dad64 100644 --- a/python/cudf_polars/tests/expressions/test_sort.py +++ b/python/cudf_polars/tests/expressions/test_sort.py @@ -11,6 +11,7 @@ from cudf_polars.testing.asserts import ( assert_gpu_result_equal, ) +from cudf_polars.testing.engine_utils import is_streaming_engine from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_140 @@ -69,7 +70,7 @@ def test_setsorted(engine: pl.GPUEngine, request, descending, nulls_last, with_n "fixed in https://github.com/pola-rs/polars/pull/25250" ) ) - elif not POLARS_VERSION_LT_140 and engine.config.get("executor") != "in-memory": + elif not POLARS_VERSION_LT_140 and is_streaming_engine(engine): request.applymarker( pytest.mark.xfail(reason="Streaming hint_sorted unsupported") ) diff --git a/python/cudf_polars/tests/test_mapfunction.py b/python/cudf_polars/tests/test_mapfunction.py index b3f09482136d..7e10739d9fe5 100644 --- a/python/cudf_polars/tests/test_mapfunction.py +++ b/python/cudf_polars/tests/test_mapfunction.py @@ -15,6 +15,7 @@ assert_gpu_result_equal, assert_ir_translation_raises, ) +from cudf_polars.testing.engine_utils import is_streaming_engine def test_explode_multiple_raises(engine: pl.GPUEngine): @@ -115,7 +116,7 @@ def test_unique_hash(): def test_set_sorted_then_inner_join( engine: pl.GPUEngine, request: pytest.FixtureRequest ): - if engine.config.get("executor") != "in-memory": + if is_streaming_engine(engine): request.applymarker( pytest.mark.xfail(reason="Streaming hint_sorted unsupported") ) @@ -136,7 +137,12 @@ def test_hint_sorted_marks_column_metadata(descending, nulls_last) -> None: } child = DataFrameScan( schema, - pl.DataFrame({"a": [1, 2, 3], "b": [3, 2, 1]})._df, + pl.DataFrame( + { + "a": [2, None, 1], + "b": [3, 1, 2], + } + )._df, None, ) node = MapFunction( @@ -154,12 +160,51 @@ def test_hint_sorted_marks_column_metadata(descending, nulls_last) -> None: if descending != nulls_last else plc.types.NullOrder.BEFORE ) - assert result.column_map["a"].is_sorted == plc.types.Sorted.YES - assert result.column_map["a"].order == order - assert result.column_map["a"].null_order == null_order + assert result.column_map["a"].check_sorted( + order=order, null_order=null_order, stream=result.stream + ) assert result.column_map["b"].is_sorted == plc.types.Sorted.NO +def test_hint_sorted_marks_multiple_column_metadata() -> None: + schema = { + "a": DataType(pl.Int64()), + "b": DataType(pl.Int64()), + "c": DataType(pl.Int64()), + } + child = DataFrameScan( + schema, + pl.DataFrame( + { + "a": [2, None, 1], + "b": [2, None, 3], + "c": [3, 1, 2], + } + )._df, + None, + ) + node = MapFunction( + schema, + "hint_sorted", + [[("a", False, False), ("b", True, False)]], + child, + ) + + result = node.evaluate(cache={}, timer=None, context=IRExecutionContext()) + + assert result.column_map["a"].check_sorted( + order=plc.types.Order.ASCENDING, + null_order=plc.types.NullOrder.BEFORE, + stream=result.stream, + ) + assert result.column_map["b"].check_sorted( + order=plc.types.Order.DESCENDING, + null_order=plc.types.NullOrder.AFTER, + stream=result.stream, + ) + assert result.column_map["c"].is_sorted == plc.types.Sorted.NO + + def test_explode_single_legacy_options(): # Cover the branch: POLARS_VERSION_LT_136 or len(self.options) == 1 # On polars >= 1.36 this branch is only reachable by direct construction From 70dd69b73f9bc8b56273964d065b143c744cb1d3 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 14 Aug 2026 13:42:21 -0700 Subject: [PATCH 4/6] another CI fix --- python/cudf_polars/tests/test_mapfunction.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/test_mapfunction.py b/python/cudf_polars/tests/test_mapfunction.py index 7e10739d9fe5..738588420d41 100644 --- a/python/cudf_polars/tests/test_mapfunction.py +++ b/python/cudf_polars/tests/test_mapfunction.py @@ -16,6 +16,7 @@ assert_ir_translation_raises, ) from cudf_polars.testing.engine_utils import is_streaming_engine +from cudf_polars.utils.versions import POLARS_VERSION_LT_140 def test_explode_multiple_raises(engine: pl.GPUEngine): @@ -116,7 +117,11 @@ def test_unique_hash(): def test_set_sorted_then_inner_join( engine: pl.GPUEngine, request: pytest.FixtureRequest ): - if is_streaming_engine(engine): + if POLARS_VERSION_LT_140: + request.applymarker( + pytest.mark.xfail(reason="set_sorted lowers to unsupported hint ir") + ) + elif is_streaming_engine(engine): request.applymarker( pytest.mark.xfail(reason="Streaming hint_sorted unsupported") ) From 033a82a679c7008251af289065f203b4a628cefc Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 14 Aug 2026 15:33:56 -0700 Subject: [PATCH 5/6] Tolerate hint_sorted in streaming lowering Signed-off-by: rjzamora --- python/cudf_polars/cudf_polars/dsl/ir.py | 34 +++++++++----- .../cudf_polars/streaming/parallel.py | 4 +- .../cudf_polars/tests/expressions/test_agg.py | 8 +--- .../tests/expressions/test_sort.py | 7 +-- python/cudf_polars/tests/test_mapfunction.py | 47 +++++++++++++++++-- 5 files changed, 66 insertions(+), 34 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index d542b8ef9cf5..dc10af430bac 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -3602,19 +3602,27 @@ def __init__(self, schema: Schema, name: str, options: Any, df: IR): "Fast count unsupported for CSV scans" ) # pragma: no cover elif self.name == "hint_sorted": - (sorted_info,) = options - column_names = [] - descending = [] - nulls_last = [] - for column_name, is_descending, is_nulls_last in sorted_info: - column_names.append(column_name) - descending.append(bool(is_descending)) - nulls_last.append(bool(is_nulls_last)) - self.options = ( - tuple(column_names), - tuple(descending), - tuple(nulls_last), - ) + if len(options) == 3: + column_names, descending, nulls_last = options + self.options = ( + tuple(column_names), + tuple(bool(value) for value in descending), + tuple(bool(value) for value in nulls_last), + ) + else: + (sorted_info,) = options + column_names = [] + descending = [] + nulls_last = [] + for column_name, is_descending, is_nulls_last in sorted_info: + column_names.append(column_name) + descending.append(bool(is_descending)) + nulls_last.append(bool(is_nulls_last)) + self.options = ( + tuple(column_names), + tuple(descending), + tuple(nulls_last), + ) self._non_child_args = (schema, name, self.options) def get_hashable(self) -> Hashable: diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 794ea08b6acb..61a011e49cdb 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -303,9 +303,7 @@ def _( ir: MapFunction, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: if ir.name == "hint_sorted": - raise NotImplementedError( - "hint_sorted is not supported by the streaming executor." - ) + return _lower_ir_pwise(ir, rec, preserve_partitioning=True) # Allow pointwise operations if ir.name in ("rename", "explode"): diff --git a/python/cudf_polars/tests/expressions/test_agg.py b/python/cudf_polars/tests/expressions/test_agg.py index 53a3851a2a3e..1f7942765b09 100644 --- a/python/cudf_polars/tests/expressions/test_agg.py +++ b/python/cudf_polars/tests/expressions/test_agg.py @@ -14,10 +14,8 @@ assert_gpu_result_equal, assert_ir_translation_raises, ) -from cudf_polars.testing.engine_utils import is_streaming_engine from cudf_polars.utils.versions import ( POLARS_VERSION_LT_136, - POLARS_VERSION_LT_140, ) @@ -58,16 +56,12 @@ def is_sorted(request): @pytest.fixture -def xfail_if_sorted(engine, is_sorted, request): +def xfail_if_sorted(is_sorted, request): # See https://github.com/rapidsai/cudf/pull/20791#issuecomment-3750528419 if is_sorted and POLARS_VERSION_LT_136: request.applymarker( pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/24981") ) - elif is_sorted and not POLARS_VERSION_LT_140 and is_streaming_engine(engine): - request.applymarker( - pytest.mark.xfail(reason="Streaming hint_sorted unsupported") - ) @pytest.fixture diff --git a/python/cudf_polars/tests/expressions/test_sort.py b/python/cudf_polars/tests/expressions/test_sort.py index d2145c4dad64..5644f0e54908 100644 --- a/python/cudf_polars/tests/expressions/test_sort.py +++ b/python/cudf_polars/tests/expressions/test_sort.py @@ -11,8 +11,7 @@ from cudf_polars.testing.asserts import ( assert_gpu_result_equal, ) -from cudf_polars.testing.engine_utils import is_streaming_engine -from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_140 +from cudf_polars.utils.versions import POLARS_VERSION_LT_136 @pytest.mark.parametrize("descending", [False, True]) @@ -70,10 +69,6 @@ def test_setsorted(engine: pl.GPUEngine, request, descending, nulls_last, with_n "fixed in https://github.com/pola-rs/polars/pull/25250" ) ) - elif not POLARS_VERSION_LT_140 and is_streaming_engine(engine): - request.applymarker( - pytest.mark.xfail(reason="Streaming hint_sorted unsupported") - ) sorted_values = sorted([1, 2, 3, 4, 5, 6, -2], reverse=descending) values: list[int | None] = [*sorted_values] if with_nulls == "nulls": diff --git a/python/cudf_polars/tests/test_mapfunction.py b/python/cudf_polars/tests/test_mapfunction.py index 738588420d41..10f67fb94560 100644 --- a/python/cudf_polars/tests/test_mapfunction.py +++ b/python/cudf_polars/tests/test_mapfunction.py @@ -2,22 +2,32 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +from typing import TYPE_CHECKING + import pytest import polars as pl import pylibcudf as plc +import cudf_polars.streaming.parallel # noqa: F401 from cudf_polars.containers import DataType from cudf_polars.dsl.ir import DataFrameScan, IRExecutionContext, MapFunction from cudf_polars.dsl.translate import Translator +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.dispatch import lower_ir_node from cudf_polars.testing.asserts import ( assert_gpu_result_equal, assert_ir_translation_raises, ) -from cudf_polars.testing.engine_utils import is_streaming_engine from cudf_polars.utils.versions import POLARS_VERSION_LT_140 +if TYPE_CHECKING: + from collections.abc import MutableMapping + + from cudf_polars.dsl.ir import IR + from cudf_polars.streaming.dispatch import State + def test_explode_multiple_raises(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [[1, 2], [3, 4]], "b": [[5, 6], [7, 8]]}) @@ -121,10 +131,6 @@ def test_set_sorted_then_inner_join( request.applymarker( pytest.mark.xfail(reason="set_sorted lowers to unsupported hint ir") ) - elif is_streaming_engine(engine): - request.applymarker( - pytest.mark.xfail(reason="Streaming hint_sorted unsupported") - ) df = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}) q = df.set_sorted("a").join( @@ -210,6 +216,37 @@ def test_hint_sorted_marks_multiple_column_metadata() -> None: assert result.column_map["c"].is_sorted == plc.types.Sorted.NO +def test_hint_sorted_normalized_options_roundtrip() -> None: + schema = {"a": DataType(pl.Int64())} + child = DataFrameScan(schema, pl.DataFrame({"a": [1]})._df, None) + node = MapFunction(schema, "hint_sorted", [[("a", False, False)]], child) + reconstructed = MapFunction(schema, "hint_sorted", node.options, child) + + assert reconstructed.options == node.options + + +def test_hint_sorted_streaming_lowering_preserves_partitioning() -> None: + schema = {"a": DataType(pl.Int64())} + child = DataFrameScan(schema, pl.DataFrame({"a": [1, 2, 3]})._df, None) + node = MapFunction(schema, "hint_sorted", [[("a", False, False)]], child) + child_partition = PartitionInfo(count=3) + + class Rec: + @property + def state(self) -> State: + raise AssertionError("state is not used by hint_sorted lowering") + + def __call__(self, ir: IR) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: + assert ir is child + return ir, {ir: child_partition} + + lowered, partition_info = lower_ir_node(node, Rec()) + + assert isinstance(lowered, MapFunction) + assert lowered.name == "hint_sorted" + assert partition_info[lowered] is child_partition + + def test_explode_single_legacy_options(): # Cover the branch: POLARS_VERSION_LT_136 or len(self.options) == 1 # On polars >= 1.36 this branch is only reachable by direct construction From a89874d17e19728cd74c8de5cc96f4feb90ce907 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 17 Aug 2026 11:54:34 -0700 Subject: [PATCH 6/6] trigger CI