feat(starrocks): translate EXCHANGE_NODE as a stream read of the engine's exchange view

experimental/starrocks/crates/starrocks-plan-translator/src/descriptor_table.rs

@@ -238,23 +238,31 @@ impl DescriptorTable {
 
     /// Builds the Substrait schema for a StarRocks tuple.
     pub fn named_struct(&self, tuple_id: i32) -> Result<NamedStruct> {
-        let (names, types): (Vec<_>, Vec<_>) = self
-            .materialized_slot_ids(tuple_id)?
-            .into_iter()
-            .map(|slot_id| {
+        self.named_struct_for_tuples(&[tuple_id])
+    }
+
+    /// Builds the Substrait schema for a row layout spanning one or more tuples.
+    ///
+    /// Names and types both come from the descriptor table. An exchange read renames the
+    /// columns afterwards: they are the sender's, bound positionally (see
+    /// `node_translator::translate_exchange`).
+    pub(crate) fn named_struct_for_tuples(&self, tuple_ids: &[i32]) -> Result<NamedStruct> {
+        let mut names = Vec::new();
+        let mut types = Vec::new();
+        for &tuple_id in tuple_ids {
+            for slot_id in self.materialized_slot_ids(tuple_id)? {
                 let slot = self.slot(tuple_id, slot_id)?;
-                let substrait_type =
+                names.push(slot.output_name());
+                types.push(
                     slot.substrait_type
                         .clone()
                         .ok_or(TranslateError::MissingField {
                             context: "materialized TSlotDescriptor",
                             field: "slotType",
-                        })?;
-                Ok((slot.output_name(), substrait_type))
-            })
-            .collect::<Result<Vec<_>>>()?
-            .into_iter()
-            .unzip();
+                        })?,
+                );
+            }
+        }
 
         Ok(NamedStruct {
             names,
@@ -269,29 +277,51 @@ impl DescriptorTable {
     /// Resolves a StarRocks `(tuple id, slot id)` to its zero-based field index in `row_tuples`.
     ///
     /// The tuple id comes from the `TSlotRef` and disambiguates slots: ids are unique only within
-    /// a tuple, so the same slot id can name different columns in different tuples.
+    /// a tuple, so the same slot id can name different columns in different tuples. It is NOT
+    /// authoritative, though. StarRocks' BE resolves slot refs by slot id alone (the tuple id is
+    /// only consulted by `is_bound`), and the FE relies on that: `buildAggregateTuple` never
+    /// rebinds `colRefToExpr` for grouping columns, so nodes above a multi-stage aggregation keep
+    /// naming them by the tuple they had below it (the TPC-H q16 shape). New-optimizer slot ids
+    /// are query-wide `ColumnRefOperator` ids, so the same id in another of the row's tuples is
+    /// the same column.
+    ///
+    /// Resolution rule: an exact `(tuple, slot)` match wins; otherwise, if exactly one tuple in
+    /// `row_tuples` carries the slot id, resolve to it (reproducing BE semantics). Zero
+    /// candidates is a descriptor error; more than one is too, since a fallback two tuples could
+    /// satisfy is a guess, not a resolution.
     pub fn slot_global_index(
         &self,
         tuple_id: i32,
         slot_id: i32,
         row_tuples: &[i32],
     ) -> Result<usize> {
         let mut offset = 0;
+        let mut by_slot_id = Vec::new();
         for &candidate_tuple in row_tuples {
             let materialized = self.materialized_slot_ids(candidate_tuple)?;
-            if candidate_tuple == tuple_id
-                && let Some(index) = materialized
-                    .iter()
-                    .position(|candidate| *candidate == slot_id)
+            if let Some(index) = materialized
+                .iter()
+                .position(|candidate| *candidate == slot_id)
             {
-                return Ok(offset + index);
+                if candidate_tuple == tuple_id {
+                    return Ok(offset + index);
+                }
+                by_slot_id.push(offset + index);
             }
             offset += materialized.len();
         }
 
-        Err(TranslateError::descriptor(format!(
-            "slot {slot_id} (tuple {tuple_id}) is not part of row_tuples {row_tuples:?}"
-        )))
+        match by_slot_id.as_slice() {
+            [index] => Ok(*index),
+            [] => Err(TranslateError::descriptor(format!(
+                "slot {slot_id} (tuple {tuple_id}) is not part of row_tuples {row_tuples:?}"
+            ))),
+            _ => Err(TranslateError::descriptor(format!(
+                "slot {slot_id} (tuple {tuple_id}) is not part of row_tuples {row_tuples:?}, and {} \
+                 of those tuples carry a slot {slot_id} to fall back on",
+                by_slot_id.len()
+            ))),
+        }
     }
 
     /// Returns the Substrait named-table path for a tuple's backing table.
@@ -401,6 +431,36 @@ mod tests {
         assert!(matches!(err, TranslateError::Descriptor(_)));
     }
 
+    /// A slot ref naming a tuple absent from the row still resolves when exactly one row tuple
+    /// carries its slot id: the TPC-H q16 shape, where `buildAggregateTuple` leaves grouping
+    /// columns bound to the tuple below the aggregation.
+    #[test]
+    fn slot_global_index_falls_back_to_the_slot_id_across_tuples() {
+        let desc = two_tuple_desc();
+        // Refs stale-bound to tuple 0 resolve against a row of [1] by slot id alone.
+        assert_eq!(desc.slot_global_index(0, 3, &[1]).unwrap(), 0);
+        assert_eq!(desc.slot_global_index(0, 4, &[1]).unwrap(), 1);
+    }
+
+    /// A fallback two tuples could satisfy is a guess, not a resolution. It must fail loudly.
+    #[test]
+    fn slot_global_index_refuses_an_ambiguous_slot_id_fallback() {
+        // Tuples 0 and 1 both carry a slot 1.
+        let desc_tbl = TDescriptorTable::new(
+            Some(vec![slot(1, 0, 0, "a"), slot(1, 1, 0, "x")]),
+            vec![
+                TTupleDescriptor::new(Some(0), None, None, None, None),
+                TTupleDescriptor::new(Some(1), None, None, None, None),
+            ],
+            None,
+            None,
+        );
+        let desc = DescriptorTable::try_from(&desc_tbl).unwrap();
+
+        let err = desc.slot_global_index(2, 1, &[0, 1]).unwrap_err();
+        assert!(matches!(err, TranslateError::Descriptor(_)));
+    }
+
     /// Verifies an unknown slot id surfaces a descriptor error.
     #[test]
     fn slot_global_index_reports_unknown_slot() {

experimental/starrocks/crates/starrocks-plan-translator/src/lib.rs

@@ -6,9 +6,9 @@
 //! invariants over breadth: it translates one fragment at a time, and everything
 //! outside the supported surface returns a structured [`TranslateError`] that
 //! names the offending node/type — so the next contributor knows exactly what to
-//! implement next. In particular `EXCHANGE_NODE` is rejected: a fragment is
-//! translated in isolation, and multi-fragment plans (every exchange is a
-//! fragment boundary) are a later milestone.
+//! implement next. An `EXCHANGE_NODE` is a fragment boundary; it translates only
+//! when the compute node supplies, per exchange node, the engine stream view its
+//! senders' batches arrive on (see [`ExchangeInput`]).
 //!
 //! # Wire format: flat preorder
 //!
@@ -35,11 +35,12 @@
 //! | `SORT_NODE`          | `ProjectRel` (sort tuple) + `SortRel` (global row-number top-N only) |
 //! | `HASH_JOIN_NODE`      | `JoinRel` (inner/outer/left-semi; left/right anti as outer join + `is_null` filter, null-aware left anti as mark join + `not`) |
 //! | `NESTLOOP_JOIN_NODE` | `JoinRel` (constant-key inner) + optional `FilterRel`, inner/cross only |
+//! | `EXCHANGE_NODE`      | `ReadRel` (named table = the engine's `sirius_stream_<node_id>` view); a merging exchange adds a `SortRel` |
 //!
 //! Node-level `conjuncts` (scan/filter predicates, HAVING, post-join filters) become a
 //! `FilterRel` over the node's output on every supported node.
 //!
-//! Any node's non-negative `limit` (plus a sort offset) becomes a `FetchRel` on top of its
+//! Any node's non-negative `limit` (plus a sort or exchange offset) becomes a `FetchRel` on top of its
 //! relation.
 //!
 //! | Expression node   | Substrait expression |
@@ -89,6 +90,7 @@ use std::fmt;
 use prost::Message;
 use starrocks_thrift::exprs::{TExpr, TExprNodeType};
 use starrocks_thrift::internal_service::TExecPlanFragmentParams;
+use starrocks_thrift::partitions::TPartitionType;
 use substrait::proto::extensions::simple_extension_declaration;
 use substrait::proto::extensions::{SimpleExtensionDeclaration, SimpleExtensionUrn};
 use substrait::proto::{Plan, PlanRel, RelRoot, plan_rel};
@@ -128,6 +130,47 @@ pub struct TranslatedPlan {
     pub plan: Plan,
     /// Root output names as emitted in the Substrait plan.
     pub output_names: Vec<String>,
+    /// For a fragment whose data-stream sink is HASH_PARTITIONED: the output column index of
+    /// each partition key, in the sink's partition-expression order. Every sender instance of
+    /// one exchange derives the same indices from the same FE thrift, which is one leg of the
+    /// cross-sender hash-parity contract (the others are one hash function and one destination
+    /// count). `None` for UNPARTITIONED / result sinks.
+    pub output_partition_columns: Option<Vec<usize>>,
+    /// One entry per exchange node lowered to a stream read. The caller declares these on the
+    /// engine before handing it the plan: a stream has no file to infer a schema from.
+    pub stream_inputs: Vec<StreamInputSchema>,
+}
+
+/// The input stream bound to one StarRocks exchange node.
+#[derive(Clone, Debug)]
+pub struct ExchangeInput {
+    /// Receiver `EXCHANGE_NODE` id.
+    pub node_id: i32,
+    /// Name of the engine view this exchange's input stream is read through.
+    pub stream_view: String,
+    /// Sender output names, which become the stream's column names. Bound by position: one per
+    /// column of the exchange's `input_row_tuples`, in row order.
+    pub names: Vec<String>,
+}
+
+/// The schema one exchange's input stream must be declared with, as the plan reads it.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct StreamInputSchema {
+    /// Receiver `EXCHANGE_NODE` id, which is also the engine-side stream id.
+    pub node_id: i32,
+    /// Name of the engine view the plan reads this stream through.
+    pub stream_view: String,
+    /// Columns in plan order.
+    pub columns: Vec<StreamInputColumn>,
+}
+
+/// One column of a declared input stream.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct StreamInputColumn {
+    /// Column name, matching the read's base schema.
+    pub name: String,
+    /// DuckDB type name the engine parses when declaring the stream.
+    pub ty: String,
 }
 
 impl TranslatedPlan {
@@ -147,6 +190,8 @@ impl fmt::Debug for TranslatedPlan {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("TranslatedPlan")
             .field("output_names", &self.output_names)
+            .field("output_partition_columns", &self.output_partition_columns)
+            .field("stream_inputs", &self.stream_inputs)
             .field("plan", &self.explain())
             .finish()
     }
@@ -200,7 +245,19 @@ impl PlanTranslator {
     }
 
     /// Translates a StarRocks execution fragment into a Substrait plan.
+    ///
+    /// Equivalent to [`Self::translate_fragment_with_exchange_inputs`] with no bound streams, so
+    /// a fragment containing an `EXCHANGE_NODE` is refused.
     pub fn translate_fragment(&self, params: &TExecPlanFragmentParams) -> Result<TranslatedPlan> {
+        self.translate_fragment_with_exchange_inputs(params, &[])
+    }
+
+    /// Translates a fragment whose exchange nodes read the given input streams.
+    pub fn translate_fragment_with_exchange_inputs(
+        &self,
+        params: &TExecPlanFragmentParams,
+        exchange_inputs: &[ExchangeInput],
+    ) -> Result<TranslatedPlan> {
         let fragment = params
             .fragment
             .as_ref()
@@ -222,9 +279,21 @@ impl PlanTranslator {
 
         let desc = DescriptorTable::try_from(desc_tbl)?;
         let scan_paths = ScanFilePaths::from_fragment(params, &desc)?;
+        let exchange_inputs = exchange_inputs
+            .iter()
+            .map(|input| (input.node_id, input))
+            .collect::<HashMap<_, _>>();
         let mut registry = ExtensionRegistry::new();
-        let mut translated =
-            node_translator::translate_plan(plan, &desc, &scan_paths, &mut registry)?;
+        let node_translator::TranslatedFragment {
+            root: mut translated,
+            stream_inputs,
+        } = node_translator::translate_plan(
+            plan,
+            &desc,
+            &scan_paths,
+            &exchange_inputs,
+            &mut registry,
+        )?;
 
         let output_names = if let Some(output_exprs) = fragment
             .output_exprs
@@ -246,6 +315,69 @@ impl PlanTranslator {
         };
 
         let output_names = unique_names(output_names).collect::<Vec<_>>();
+
+        // One loud guard against any width/name drift at the root: the names are what the
+        // receiver (or the client) reads the row by, so a mismatch is a wrong column read
+        // waiting to happen. Never ship it.
+        if output_names.len() != translated.output_width {
+            return Err(TranslateError::malformed(format!(
+                "fragment root emits {} columns but derived {} output names",
+                translated.output_width,
+                output_names.len()
+            )));
+        }
+
+        // Resolve a hash-partitioned sink's keys to output column indices while the row layout
+        // is still in hand. Bare SLOT_REFs only: any transform would make this sender hash a
+        // value its peers do not, silently splitting equal keys across destinations.
+        let output_partition_columns = match fragment
+            .output_sink
+            .as_ref()
+            .and_then(|sink| sink.stream_sink.as_ref())
+            .map(|stream_sink| &stream_sink.output_partition)
+        {
+            Some(partition) if partition.type_ == TPartitionType::HASH_PARTITIONED => {
+                if fragment
+                    .output_exprs
+                    .as_ref()
+                    .is_some_and(|exprs| !exprs.is_empty())
+                {
+                    return Err(TranslateError::malformed(
+                        "a hash-partitioned stream sink with output_exprs cannot map its \
+                         partition keys onto the sink row (never emitted by the FE)",
+                    ));
+                }
+                let exprs = partition.partition_exprs.as_deref().unwrap_or_default();
+                if exprs.is_empty() {
+                    return Err(TranslateError::malformed(
+                        "a hash-partitioned stream sink carries no partition expressions",
+                    ));
+                }
+                let mut columns = Vec::with_capacity(exprs.len());
+                for expr in exprs {
+                    let slot_ref = match expr.nodes.as_slice() {
+                        [node] if node.node_type == TExprNodeType::SLOT_REF => {
+                            node.slot_ref.as_ref()
+                        }
+                        _ => None,
+                    };
+                    let Some(slot_ref) = slot_ref else {
+                        return Err(TranslateError::malformed(
+                            "a hash-partition key is not a bare slot reference; hashing a \
+                             transformed key would silently split equal keys across senders",
+                        ));
+                    };
+                    columns.push(desc.slot_global_index(
+                        slot_ref.tuple_id,
+                        slot_ref.slot_id,
+                        &translated.row_tuples,
+                    )?);
+                }
+                Some(columns)
+            }
+            _ => None,
+        };
+
         let (extension_urns, extensions) = registry.into_extensions();
         let substrait_plan = Plan {
             // Source the spec version from the `substrait` crate so it tracks the
@@ -268,6 +400,8 @@ impl PlanTranslator {
         Ok(TranslatedPlan {
             plan: substrait_plan,
             output_names,
+            output_partition_columns,
+            stream_inputs,
         })
     }
 }