feat(sql): ROLLUP, CUBE, GROUPING SETS, grouping() and grouping_id() (P2-T053b)

crates/krishiv-common/src/stack.rs

@@ -7,7 +7,7 @@
 //! never needs more than `STEP` frames of the caller's stack.
 
 /// Levels a recursive pass may take on one stack segment.
-pub const STEP: usize = 64;
+pub const STEP: usize = 32;
 
 /// Runs `f` on a new thread with a 64 MiB stack (16 MiB in release
 /// builds), waiting for it; `None` when no thread can be started. A panic

crates/krishiv-exec/src/lib.rs

@@ -115,10 +115,11 @@ pub use operator::{
     FlushOutput, GlobalState, LocalState, Operator, OperatorOutput, Sink, SinkStatus, Source,
 };
 pub use operators::{
-    DEFAULT_MAX_MERGE_FAN_IN, FileScan, FileScanMetrics, FileScanSource, FileSelections, FileWrite,
-    FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, MemoryScan,
-    Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RowStageSettings, RuntimeFilters,
-    ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, SortMetrics, TopN, UnionAlign,
+    DEFAULT_MAX_MERGE_FAN_IN, Expand, FileScan, FileScanMetrics, FileScanSource, FileSelections,
+    FileWrite, FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN,
+    MemoryScan, NullOf, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows,
+    RowStageSettings, RuntimeFilters, ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey,
+    SortMetrics, TopN, UnionAlign,
 };
 pub use pipeline::{ExecNode, PipelineBuilder};
 pub use runner::{QueryRunner, Scheduler};

crates/krishiv-exec/src/operators/expand.rs

@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: Apache-2.0
+//! Each row once per projection: the first step of `ROLLUP`, `CUBE` and
+//! `GROUPING SETS` (Spark's `Expand`).
+
+use std::sync::Arc;
+
+use arrow_array::{ArrayRef, RecordBatch, new_null_array};
+use arrow_schema::{DataType, SchemaRef};
+use krishiv_types::batch::Batch;
+
+use super::{Project, ScalarExpr, local_mut, no_state};
+use crate::context::TaskContext;
+use crate::error::{Result, internal};
+use crate::operator::{LocalState, Operator, OperatorOutput};
+
+/// Emits, for every input batch, one batch per projection: the batch
+/// projected by it. The projections share one output schema. Rows keep their
+/// weights in every copy. Output is produced one projection per
+/// [`Operator::execute`] call (`HaveMoreOutput` until the last), so only
+/// one copy of an input batch exists at a time, and the plain columns of
+/// each copy share the input's buffers.
+#[derive(Debug)]
+pub struct Expand {
+    projects: Vec<Arc<Project>>,
+    schema: SchemaRef,
+}
+
+impl Expand {
+    /// Expands batches by `projects`.
+    ///
+    /// # Errors
+    /// `INTERNAL_ERROR` if there is no projection, or two of them output
+    /// different columns.
+    pub fn try_new(projects: Vec<Arc<Project>>) -> Result<Arc<Self>> {
+        let first = projects
+            .first()
+            .ok_or_else(|| internal("an expansion without projections".to_owned()))?;
+        let schema = first.schema();
+        for p in &projects {
+            let (a, b) = (p.schema(), &schema);
+            let same = a.fields().len() == b.fields().len()
+                && a.fields().iter().zip(b.fields()).all(|(x, y)| {
+                    x.name() == y.name()
+                        && x.data_type() == y.data_type()
+                        && x.is_nullable() == y.is_nullable()
+                });
+            if !same {
+                return Err(internal(format!(
+                    "the projections of an expansion differ: {:?} and {:?}",
+                    a.fields(),
+                    b.fields()
+                )));
+            }
+        }
+        Ok(Arc::new(Self { projects, schema }))
+    }
+}
+
+/// The next projection to emit for the current input batch.
+struct Cursor(usize);
+
+impl Operator for Expand {
+    fn name(&self) -> &'static str {
+        "Expand"
+    }
+
+    fn schema(&self) -> SchemaRef {
+        Arc::clone(&self.schema)
+    }
+
+    fn local_state(&self, _ctx: &TaskContext) -> Result<LocalState> {
+        Ok(Box::new(Cursor(0)))
+    }
+
+    fn execute(&self, input: &Batch, local: &mut LocalState) -> Result<OperatorOutput> {
+        let cursor = local_mut::<Cursor>(local, "Expand")?;
+        if input.num_rows() == 0 {
+            cursor.0 = 0;
+            return Ok(OperatorOutput::NeedMoreInput(None));
+        }
+        let project = self
+            .projects
+            .get(cursor.0)
+            .ok_or_else(|| internal("an expansion past its last projection".to_owned()))?;
+        let OperatorOutput::NeedMoreInput(Some(out)) = project.execute(input, &mut no_state())?
+        else {
+            return Err(internal("a projection did not output a batch".to_owned()));
+        };
+        cursor.0 += 1;
+        if cursor.0 < self.projects.len() {
+            return Ok(OperatorOutput::HaveMoreOutput(out));
+        }
+        cursor.0 = 0;
+        Ok(OperatorOutput::NeedMoreInput(Some(out)))
+    }
+}
+
+/// A column of nulls, in the Arrow type of an input column. An expansion
+/// nulls a grouping column in the carrier its other copies have: strings
+/// reach an operator as `Utf8` from a registered table and as `Utf8View`
+/// from a file.
+#[derive(Debug)]
+pub struct NullOf {
+    index: usize,
+    data_type: DataType,
+}
+
+impl NullOf {
+    /// Nulls of input column `index` of batches of `input`.
+    ///
+    /// # Errors
+    /// `INTERNAL_ERROR` if the column is not in `input`.
+    pub fn new(input: &SchemaRef, index: usize) -> Result<Self> {
+        let field = input
+            .fields()
+            .get(index)
+            .ok_or_else(|| internal(format!("null-of index {index} out of range")))?;
+        Ok(Self {
+            index,
+            data_type: field.data_type().clone(),
+        })
+    }
+}
+
+impl ScalarExpr for NullOf {
+    fn data_type(&self) -> DataType {
+        self.data_type.clone()
+    }
+
+    fn nullable(&self) -> bool {
+        true
+    }
+
+    fn evaluate(&self, batch: &RecordBatch) -> Result<ArrayRef> {
+        let column = batch
+            .columns()
+            .get(self.index)
+            .ok_or_else(|| internal(format!("null-of index {} out of range", self.index)))?;
+        Ok(new_null_array(column.data_type(), batch.num_rows()))
+    }
+}

crates/krishiv-exec/src/operators/mod.rs

@@ -25,6 +25,7 @@
 //! - **`UNION`.** See [`union`] for how `UNION ALL` fan-in and the `UNION`
 //!   (distinct) breaker are split between tasks.
 
+mod expand;
 mod expr;
 mod file_scan;
 mod file_write;
@@ -37,6 +38,7 @@ mod scan;
 mod sort;
 mod union;
 
+pub use expand::{Expand, NullOf};
 pub use expr::{FnPredicate, Predicate, ScalarExpr};
 pub use file_scan::{FileScan, FileScanMetrics, FileScanSource, METADATA_COLUMN};
 pub use file_write::{FileWrite, FileWriteSink};

crates/krishiv-exec/src/operators/tests.rs

@@ -841,3 +841,127 @@ fn replicate_rows_keeps_weights_and_rejects_non_bigint_counts() {
         .unwrap_or_else(|| unreachable!("an error"));
     assert_eq!(err.error_class(), ErrorClass::InternalError);
 }
+
+/// A projection of `schema()`: `a` (or nulls of it), then `s` (or nulls).
+fn expand_projection(null_a: bool, null_s: bool) -> Arc<Project> {
+    let item = |index: usize, name: &str, null: bool| {
+        if null {
+            let nulls = NullOf::new(&schema(), index).unwrap_or_else(|e| unreachable!("{e}"));
+            ProjectItem::expr(name, Arc::new(nulls))
+        } else {
+            ProjectItem::column(index)
+        }
+    };
+    Project::try_new(&schema(), vec![item(0, "a", null_a), item(1, "s", null_s)])
+        .unwrap_or_else(|e| unreachable!("{e}"))
+}
+
+#[test]
+fn expand_streams_one_copy_per_projection_and_keeps_weights() {
+    let (_root, ctx) = ctx();
+    let op = Expand::try_new(vec![
+        expand_projection(false, false),
+        expand_projection(false, true),
+        expand_projection(true, true),
+    ])
+    .unwrap_or_else(|e| unreachable!("{e}"));
+    assert_eq!(op.name(), "Expand");
+    assert_eq!(op.schema(), schema());
+    let mut local = op.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}"));
+    let input = weighted(&[1, 2], &[3, -4]);
+    // One copy per call: `HaveMoreOutput` until the last projection.
+    let mut copies = Vec::new();
+    for _ in 0..2 {
+        let OperatorOutput::HaveMoreOutput(b) = op
+            .execute(&input, &mut local)
+            .unwrap_or_else(|e| unreachable!("{e}"))
+        else {
+            unreachable!("a copy and more to come")
+        };
+        copies.push(b);
+    }
+    copies.extend(out_batch(
+        op.execute(&input, &mut local)
+            .unwrap_or_else(|e| unreachable!("{e}")),
+    ));
+    assert_eq!(copies.len(), 3);
+    assert_eq!(col_a(&copies[0]), [1, 2].map(Some));
+    assert_eq!(col_a(&copies[1]), [1, 2].map(Some));
+    assert_eq!(col_a(&copies[2]), [None, None]);
+    let strings = |b: &Batch| -> Vec<Option<String>> {
+        b.data()
+            .column(1)
+            .as_string::<i32>()
+            .iter()
+            .map(|v| v.map(str::to_owned))
+            .collect()
+    };
+    assert_eq!(
+        strings(&copies[0]),
+        [Some("s1".to_owned()), Some("s2".to_owned())]
+    );
+    assert_eq!(strings(&copies[1]), [None, None]);
+    assert!(copies.iter().all(|b| b.weights() == input.weights()));
+    // The cursor starts again for the next input; an empty batch is skipped.
+    assert!(
+        out_batch(
+            op.execute(&ints(&[]), &mut local)
+                .unwrap_or_else(|e| unreachable!("{e}"))
+        )
+        .is_none()
+    );
+    assert!(matches!(
+        op.execute(&ints(&[7]), &mut local),
+        Ok(OperatorOutput::HaveMoreOutput(_))
+    ));
+}
+
+#[test]
+fn expand_rejects_bad_plans() {
+    let class = |r: Result<Arc<Expand>, KrishivError>| {
+        r.err()
+            .unwrap_or_else(|| unreachable!("an error"))
+            .error_class()
+    };
+    assert_eq!(class(Expand::try_new(vec![])), ErrorClass::InternalError);
+    // A projection that outputs other columns.
+    let narrow = Project::try_new(&schema(), vec![ProjectItem::column(0)])
+        .unwrap_or_else(|e| unreachable!("{e}"));
+    assert_eq!(
+        class(Expand::try_new(vec![
+            expand_projection(false, false),
+            narrow
+        ])),
+        ErrorClass::InternalError
+    );
+    let renamed = Project::try_new(
+        &schema(),
+        vec![ProjectItem::aliased(0, "x"), ProjectItem::column(1)],
+    )
+    .unwrap_or_else(|e| unreachable!("{e}"));
+    assert_eq!(
+        class(Expand::try_new(vec![
+            expand_projection(false, false),
+            renamed
+        ])),
+        ErrorClass::InternalError
+    );
+    // Nulls of a column the input has not; and of a batch that lost it.
+    let err = NullOf::new(&schema(), 2)
+        .err()
+        .unwrap_or_else(|| unreachable!("an error"));
+    assert_eq!(err.error_class(), ErrorClass::InternalError);
+    let nulls = NullOf::new(&schema(), 1).unwrap_or_else(|e| unreachable!("{e}"));
+    assert_eq!(nulls.data_type(), DataType::Utf8);
+    assert!(nulls.nullable());
+    let one_column = RecordBatch::try_new(
+        Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])),
+        vec![Arc::new(Int64Array::from(vec![1_i64]))],
+    )
+    .unwrap_or_else(|e| unreachable!("{e}"));
+    let err = nulls
+        .evaluate(&one_column)
+        .err()
+        .unwrap_or_else(|| unreachable!("an error"));
+    assert_eq!(err.error_class(), ErrorClass::InternalError);
+}

crates/krishiv-exec/src/physical/instantiate.rs

@@ -48,8 +48,8 @@ use crate::join::{
 };
 use crate::operator::{Operator, Sink, Source};
 use crate::operators::{
-    FileWrite, FileWriteSink, Filter, FilterPushdown, Predicate, Project, ProjectItem,
-    PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN,
+    Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project,
+    ProjectItem, PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN,
 };
 use crate::pipeline::{ExecNode, PipelineBuilder};
 use crate::planner::TableProviderResolver;
@@ -296,6 +296,15 @@ impl Instantiation<'_> {
             } => self.with_subqueries(input, subqueries, |input, v| {
                 replicate_op(input, layout, count, v)
             }),
+            PhysNode::Expand {
+                input,
+                layout,
+                projections,
+                nullable,
+                subqueries,
+            } => self.with_subqueries(input, subqueries, |input, v| {
+                expand_op(input, layout, projections, nullable, v)
+            }),
             PhysNode::Limit {
                 input,
                 offset,
@@ -586,6 +595,9 @@ pub(crate) fn project_op(
                     name.clone(),
                     Arc::new(Compiled(compile_one(expr, layout, schema, v)?)),
                 ),
+                ProjectItemPlan::NullOf { index, name } => {
+                    ProjectItem::expr(name.clone(), Arc::new(NullOf::new(schema, *index)?))
+                }
             })
         })
         .collect::<Result<Vec<_>>>()?;
@@ -595,6 +607,20 @@ pub(crate) fn project_op(
     }
 }
 
+pub(crate) fn expand_op(
+    schema: &SchemaRef,
+    layout: &[AttrId],
+    projections: &[Vec<ProjectItemPlan>],
+    nullable: &[bool],
+    v: &Values<'_>,
+) -> Result<Arc<dyn Operator>> {
+    let projects = projections
+        .iter()
+        .map(|items| project_op(schema, layout, items, Some(nullable), v))
+        .collect::<Result<Vec<_>>>()?;
+    Ok(Expand::try_new(projects)?)
+}
+
 pub(crate) fn replicate_op(
     schema: &SchemaRef,
     layout: &[AttrId],