diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index cd0b61c63c1f..8c7841cd3eb8 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -3592,10 +3592,28 @@ 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": + 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: @@ -3705,6 +3723,23 @@ def do_evaluate( dtype=dtype, ) return DataFrame([index_col, *df.columns], stream=df.stream) + elif name == "hint_sorted": + column_names, descending, nulls_last = options + 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..61a011e49cdb 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -302,6 +302,9 @@ def _( def _( ir: MapFunction, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: + if ir.name == "hint_sorted": + return _lower_ir_pwise(ir, rec, preserve_partitioning=True) + # Allow pointwise operations if ir.name in ("rename", "explode"): return _lower_ir_pwise(ir, rec) diff --git a/python/cudf_polars/tests/expressions/test_agg.py b/python/cudf_polars/tests/expressions/test_agg.py index b363619e192f..1f7942765b09 100644 --- a/python/cudf_polars/tests/expressions/test_agg.py +++ b/python/cudf_polars/tests/expressions/test_agg.py @@ -58,7 +58,7 @@ def is_sorted(request): @pytest.fixture def xfail_if_sorted(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") ) diff --git a/python/cudf_polars/tests/expressions/test_sort.py b/python/cudf_polars/tests/expressions/test_sort.py index ccc3049e5df3..5644f0e54908 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,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: - # 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")) 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..10f67fb94560 100644 --- a/python/cudf_polars/tests/test_mapfunction.py +++ b/python/cudf_polars/tests/test_mapfunction.py @@ -1,18 +1,32 @@ -# 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 +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, MapFunction +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.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): @@ -110,8 +124,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 POLARS_VERSION_LT_140: + request.applymarker( + pytest.mark.xfail(reason="set_sorted lowers to unsupported hint ir") + ) df = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}) q = df.set_sorted("a").join( @@ -120,6 +139,114 @@ 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": [2, None, 1], + "b": [3, 1, 2], + } + )._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"].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_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