Track local row ordering in streaming cudf-polars metadata

cpp/libcudf_streaming/include/cudf_streaming/channel_metadata.hpp

@@ -60,7 +60,7 @@ struct order_key {
 };
 
 /**
- * @brief A valid ordering description for sorted/range-partitioned data.
+ * @brief A valid ordering description for order-partitioned data.
  *
  * Data is partitioned by value ranges based on predetermined boundaries.
  * For N partitions, there are N-1 boundary rows:
@@ -77,12 +77,18 @@ struct order_key {
  * `strict_boundaries`: when true, every row in a chunk belongs to a single partition's
  * half-open key range (partition keys do not straddle chunk interiors). When false,
  * a chunk may contain keys spanning multiple partitions.
+ *
+ * `locally_ordered`: when true, rows within each partition at this level are ordered by
+ * `keys`. When false, partitions are still ordered across partition boundaries, but row
+ * order within a partition is not guaranteed.
  */
 struct ordering {
   std::vector<order_key> keys;              ///< Sort keys (column, order, null_order per entry).
   std::shared_ptr<table_chunk> boundaries;  ///< N-1 boundary rows for N partitions.
   /// See struct-level note on `strict_boundaries` semantics.
   bool strict_boundaries{false};
+  /// See struct-level note on `locally_ordered` semantics.
+  bool locally_ordered{true};
 
   /// @brief Default constructor. Produces an invalid (empty) ordering.
   ordering() = default;
@@ -94,12 +100,14 @@ struct ordering {
    * @param boundaries Non-null, device-resident boundary table (N-1 rows for N
    * partitions). Accepts a `unique_ptr<table_chunk>` via implicit conversion.
    * @param strict_boundaries See struct-level doc. Defaults to false.
+   * @param locally_ordered See struct-level doc. Defaults to true.
    * @throws std::invalid_argument if `keys` is empty, `boundaries` is null or not
    * device-resident, or `keys.size() != boundaries->shape().second`.
    */
   ordering(std::vector<order_key> keys,
            std::shared_ptr<table_chunk> boundaries,
-           bool strict_boundaries = false);
+           bool strict_boundaries = false,
+           bool locally_ordered   = true);
 
   /**
    * @brief Return a new ordering with updated key column indices, sharing
@@ -108,32 +116,39 @@ struct ordering {
    * @param new_keys Replacement sort keys; size must equal
    * `boundaries->shape().second`.
    * @return A new ordering with `new_keys` and the same boundaries and
-   * strictness.
+   * ordering flags.
    * @throws std::invalid_argument if `new_keys` is empty or size mismatches
    * boundaries.
    */
   [[nodiscard]] ordering with_keys(std::vector<order_key> new_keys) const;
 
+  /**
+   * @brief Return a new ordering with updated local row-order metadata.
+   *
+   * @param locally_ordered Whether rows within each partition are ordered by `keys`.
+   * @return A new ordering with the same keys, boundaries, and strictness.
+   */
+  [[nodiscard]] ordering with_locally_ordered(bool locally_ordered) const;
+
   /**
    * @brief Check whether boundary values are aligned with another ordering.
    *
    * @param other The ordering to compare against.
    * @param br Buffer resource used for temporary allocations during comparison.
    * @return True when both orderings have matching boundary values and
    * strict_boundaries attributes, and are otherwise compatible (same order and
-   * null_order).
+   * null_order). This comparison intentionally ignores `locally_ordered`.
    */
   [[nodiscard]] bool boundaries_aligned_with(ordering const& other,
                                              rapidsmpf::BufferResource& br) const;
 };
 
 /**
- * @brief Order-based partitioning scheme for sorted/range-partitioned data.
+ * @brief Order-based partitioning scheme for order-partitioned data.
  *
- * An order_scheme advertises that the same stream is sorted/range-partitioned
- * with respect to any individual ordering it contains. Consumers are
- * responsible for selecting the ordering that is relevant to a particular
- * operation.
+ * An order_scheme advertises that the same stream is order-partitioned with
+ * respect to any individual ordering it contains. Consumers are responsible for
+ * selecting the ordering that is relevant to a particular operation.
  */
 struct order_scheme {
   std::vector<ordering> orderings;  ///< Ordering descriptions valid for the stream.
@@ -148,7 +163,8 @@ struct order_scheme {
    */
   order_scheme(std::vector<order_key> keys,
                std::shared_ptr<table_chunk> boundaries,
-               bool strict_boundaries = false);
+               bool strict_boundaries = false,
+               bool locally_ordered   = true);
 
   /**
    * @brief Construct a validated multi-ordering order_scheme.

cpp/libcudf_streaming/src/channel_metadata.cpp

@@ -42,15 +42,24 @@ void validate_ordering(ordering const& ordering)
 
 ordering::ordering(std::vector<order_key> keys,
                    std::shared_ptr<table_chunk> boundaries,
-                   bool strict_boundaries)
-  : keys{std::move(keys)}, boundaries{std::move(boundaries)}, strict_boundaries{strict_boundaries}
+                   bool strict_boundaries,
+                   bool locally_ordered)
+  : keys{std::move(keys)},
+    boundaries{std::move(boundaries)},
+    strict_boundaries{strict_boundaries},
+    locally_ordered{locally_ordered}
 {
   validate_ordering(*this);
 }
 
 ordering ordering::with_keys(std::vector<order_key> new_keys) const
 {
-  return ordering{std::move(new_keys), boundaries, strict_boundaries};
+  return ordering{std::move(new_keys), boundaries, strict_boundaries, locally_ordered};
+}
+
+ordering ordering::with_locally_ordered(bool locally_ordered) const
+{
+  return ordering{keys, boundaries, strict_boundaries, locally_ordered};
 }
 
 bool ordering::boundaries_aligned_with(ordering const& other, rapidsmpf::BufferResource& br) const
@@ -90,9 +99,10 @@ bool ordering::boundaries_aligned_with(ordering const& other, rapidsmpf::BufferR
 
 order_scheme::order_scheme(std::vector<order_key> keys,
                            std::shared_ptr<table_chunk> boundaries,
-                           bool strict_boundaries)
-  : order_scheme(
-      std::vector<ordering>{ordering{std::move(keys), std::move(boundaries), strict_boundaries}})
+                           bool strict_boundaries,
+                           bool locally_ordered)
+  : order_scheme(std::vector<ordering>{
+      ordering{std::move(keys), std::move(boundaries), strict_boundaries, locally_ordered}})
 {
 }
 

cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp

@@ -192,13 +192,28 @@ TEST_F(StreamingChannelMetadataGPU, OrderingReplaceKeys)
   EXPECT_EQ(o2.keys[0].column_index, 5);
   EXPECT_EQ(o2.keys[0].order, cudf::order::DESCENDING);
   EXPECT_EQ(o2.strict_boundaries, o1.strict_boundaries);
+  EXPECT_EQ(o2.locally_ordered, o1.locally_ordered);
   EXPECT_EQ(o2.boundaries->shape(), o1.boundaries->shape());
   EXPECT_EQ(o2.boundaries.get(), b.get());
   EXPECT_NE(o1.keys[0].column_index, o2.keys[0].column_index);
 
   EXPECT_THROW(static_cast<void>(o1.with_keys({k0, k5})), std::invalid_argument);
 }
 
+TEST_F(StreamingChannelMetadataGPU, OrderingReplaceLocallyOrdered)
+{
+  order_key k0{0, cudf::order::ASCENDING, cudf::null_order::BEFORE};
+
+  auto b      = make_chunk({100, 200});
+  ordering o1 = ordering({k0}, b, /*strict_boundaries=*/true);
+  auto o2     = o1.with_locally_ordered(false);
+
+  EXPECT_EQ(o2.keys[0], o1.keys[0]);
+  EXPECT_EQ(o2.strict_boundaries, o1.strict_boundaries);
+  EXPECT_FALSE(o2.locally_ordered);
+  EXPECT_EQ(o2.boundaries.get(), b.get());
+}
+
 TEST_F(StreamingChannelMetadataGPU, OrderSchemeMultipleOrderings)
 {
   order_key k0{0, cudf::order::ASCENDING, cudf::null_order::BEFORE};
@@ -212,6 +227,7 @@ TEST_F(StreamingChannelMetadataGPU, OrderSchemeMultipleOrderings)
   EXPECT_EQ(o.orderings[0].keys[0], k0);
   EXPECT_EQ(o.orderings[0].boundaries.get(), b0.get());
   EXPECT_TRUE(o.orderings[0].strict_boundaries);
+  EXPECT_TRUE(o.orderings[0].locally_ordered);
   EXPECT_EQ(o.orderings[1].keys[0], k2);
   EXPECT_EQ(o.orderings[1].boundaries.get(), b1.get());
   EXPECT_FALSE(o.orderings[1].strict_boundaries);
@@ -232,6 +248,12 @@ TEST_F(StreamingChannelMetadataGPU, OrderingBoundariesAlignedWith)
   ordering o_strict({k0}, make_chunk({100, 200}), /*strict_boundaries=*/true);
   EXPECT_FALSE(o1.boundaries_aligned_with(o_strict, *br));
 
+  ordering o_unordered({k0},
+                       make_chunk({100, 200}),
+                       /*strict_boundaries=*/false,
+                       /*locally_ordered=*/false);
+  EXPECT_TRUE(o1.boundaries_aligned_with(o_unordered, *br));
+
   ordering o_diff({k0}, make_chunk({100, 300}));
   EXPECT_FALSE(o1.boundaries_aligned_with(o_diff, *br));
 }

python/cudf_polars/cudf_polars/dsl/expressions/base.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
 # TODO: remove need for this
 # ruff: noqa: D101
@@ -14,6 +14,7 @@
 
 from cudf_polars.containers import Column
 from cudf_polars.dsl.nodebase import Node
+from cudf_polars.dsl.traversal import traversal
 
 if TYPE_CHECKING:
     from typing import Self
@@ -109,6 +110,10 @@ def evaluate(
         """
         return self.do_evaluate(df, context=context)
 
+    def all_pointwise(self) -> bool:
+        """Return True when this expression and all descendants are pointwise."""
+        return all(e.is_pointwise for e in traversal([self]))
+
     @property
     def agg_request(self) -> plc.aggregation.Aggregation:
         """
@@ -199,6 +204,10 @@ def evaluate(
         """
         return self.value.evaluate(df, context=context).rename(self.name)
 
+    def all_pointwise(self) -> bool:
+        """Return True when the underlying expression tree is pointwise."""
+        return self.value.all_pointwise()
+
     def reconstruct(self, expr: Expr) -> Self:
         """
         Rebuild with a new `Expr` value.

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -250,9 +250,22 @@ class IR(Node["IR"]):
     _non_child_args: tuple[Any, ...]
     # The number of non-child arguments to pass to do_evaluate.
     _n_non_child_args: ClassVar[int]
+    # Class-level opt-in for :attr:`preserves_output_order`.
+    _preserves_output_order: ClassVar[bool] = False
     schema: Schema
     """Mapping from column names to their data types."""
 
+    @property
+    def preserves_output_order(self) -> bool:
+        """
+        Whether output rows appear in the same relative order as this node's input.
+
+        Only meaningful for nodes with a single input. Multi-input nodes
+        (``Join``, ``Union``) need per-child reasoning and are handled by
+        their streaming actors.
+        """
+        return self._preserves_output_order
+
     def get_hashable(self) -> Hashable:
         """
         Hashable representation of node, treating schema dictionary.
@@ -1560,6 +1573,7 @@ class Cache(IR):
     Used for CSE at the plan level.
     """
 
+    _preserves_output_order: ClassVar[bool] = True
     __slots__ = ("key", "refcount")
     _non_child = ("schema", "key", "refcount")
     _n_non_child_args = 2
@@ -1784,6 +1798,11 @@ def __init__(
         ):  # pragma: no cover
             raise NotImplementedError(f"Unsupported scan type: {df.typ}")
 
+    @property
+    def preserves_output_order(self) -> bool:
+        """Whether the selected expressions keep input appearance order."""
+        return all(e.all_pointwise() for e in self.exprs)
+
     @staticmethod
     def _is_len_expr(exprs: tuple[expr.NamedExpr, ...]) -> bool:  # pragma: no cover
         if len(exprs) == 1:
@@ -1913,6 +1932,7 @@ def do_evaluate(
 class Rolling(IR):
     """Perform a (possibly grouped) rolling aggregation."""
 
+    _preserves_output_order: ClassVar[bool] = True
     __slots__ = (
         "agg_requests",
         "closed_window",
@@ -2161,6 +2181,11 @@ def __init__(
             self.zlice,
         )
 
+    @property
+    def preserves_output_order(self) -> bool:
+        """Whether grouped rows keep input appearance order."""
+        return self.maintain_order
+
     @classmethod
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="GroupBy")
@@ -3158,6 +3183,11 @@ def __init__(
         self._non_child_args = (self.columns, self.should_broadcast)
         self.children = (df,)
 
+    @property
+    def preserves_output_order(self) -> bool:
+        """Whether the stacked expressions keep input appearance order."""
+        return all(e.all_pointwise() for e in self.columns)
+
     @classmethod
     @log_do_evaluate
     @nvtx_annotate_cudf_polars(message="HStack")
@@ -3222,6 +3252,11 @@ def __init__(
         self._non_child_args = (keep, subset, zlice, stable)
         self.children = (df,)
 
+    @property
+    def preserves_output_order(self) -> bool:
+        """Whether distinct rows keep input appearance order."""
+        return self.stable
+
     _KEEP_MAP: ClassVar[dict[str, plc.stream_compaction.DuplicateKeepOption]] = {
         "first": plc.stream_compaction.DuplicateKeepOption.KEEP_FIRST,
         "last": plc.stream_compaction.DuplicateKeepOption.KEEP_LAST,
@@ -3370,6 +3405,7 @@ def do_evaluate(
 class Slice(IR):
     """Slice a dataframe."""
 
+    _preserves_output_order: ClassVar[bool] = True
     __slots__ = ("length", "offset")
     _non_child = ("schema", "offset", "length")
     _n_non_child_args = 2
@@ -3398,6 +3434,7 @@ def do_evaluate(
 class Filter(IR):
     """Filter a dataframe with a boolean mask."""
 
+    _preserves_output_order: ClassVar[bool] = True
     __slots__ = ("mask",)
     _non_child = ("schema", "mask")
     _n_non_child_args = 1
@@ -3423,6 +3460,7 @@ def do_evaluate(
 class Projection(IR):
     """Select a subset of columns from a dataframe."""
 
+    _preserves_output_order: ClassVar[bool] = True
     __slots__ = ()
     _non_child = ("schema",)
     _n_non_child_args = 1
@@ -3611,6 +3649,11 @@ def __init__(self, schema: Schema, name: str, options: Any, df: IR):
                 )
         self._non_child_args = (schema, name, self.options)
 
+    @property
+    def preserves_output_order(self) -> bool:
+        """Whether this map keeps input appearance order."""
+        return self.name in {"rechunk", "rename", "row_index", "hint_sorted"}
+
     def get_hashable(self) -> Hashable:
         """
         Hashable representation of the node.