In-memory support for set_sorted/MapFunction("hint_sorted")

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
 

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)

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")
         )

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":

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