From 127fcdd1768e7d68369df70d4462f57a174fb96c Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:20:12 +0200 Subject: [PATCH 1/8] wip: rollup/cube/grouping sets Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/lib.rs | 4 +- crates/krishiv-exec/src/operators/expand.rs | 144 ++++ crates/krishiv-exec/src/operators/mod.rs | 2 + .../krishiv-exec/src/physical/instantiate.rs | 28 +- crates/krishiv-exec/src/physical/mod.rs | 33 +- crates/krishiv-exec/src/planner/build.rs | 98 ++- crates/krishiv-optimizer/src/prune.rs | 46 ++ crates/krishiv-optimizer/src/stats.rs | 8 + crates/krishiv-plan/src/analyzer/analytics.rs | 692 ++++++++++++++++++ crates/krishiv-plan/src/analyzer/bind.rs | 4 + crates/krishiv-plan/src/analyzer/exprs.rs | 10 +- crates/krishiv-plan/src/analyzer/mod.rs | 1 + crates/krishiv-plan/src/analyzer/ops.rs | 127 +++- crates/krishiv-plan/src/analyzer/resolve.rs | 11 +- crates/krishiv-plan/src/analyzer/setops.rs | 2 +- crates/krishiv-plan/src/display.rs | 20 + crates/krishiv-plan/src/encode.rs | 11 + crates/krishiv-plan/src/lib.rs | 3 +- crates/krishiv-plan/src/plan.rs | 20 + crates/krishiv-plan/src/tests/analyzer.rs | 7 +- crates/krishiv-plan/src/tests/inventory.rs | 12 +- crates/krishiv-plan/src/tests/props.rs | 23 +- crates/krishiv-plan/src/walk.rs | 7 + crates/krishiv-session/src/dataframe.rs | 95 ++- .../krishiv-session/tests/scratch_rollup.rs | 84 +++ crates/krishiv-sql/src/lower.rs | 147 +++- 26 files changed, 1575 insertions(+), 64 deletions(-) create mode 100644 crates/krishiv-exec/src/operators/expand.rs create mode 100644 crates/krishiv-plan/src/analyzer/analytics.rs create mode 100644 crates/krishiv-session/tests/scratch_rollup.rs diff --git a/crates/krishiv-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index f6efa4c04..b9081da93 100644 --- a/crates/krishiv-exec/src/lib.rs +++ b/crates/krishiv-exec/src/lib.rs @@ -113,8 +113,8 @@ 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, + DEFAULT_MAX_MERGE_FAN_IN, Expand, NullOf, 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, }; diff --git a/crates/krishiv-exec/src/operators/expand.rs b/crates/krishiv-exec/src/operators/expand.rs new file mode 100644 index 000000000..5f0a5c233 --- /dev/null +++ b/crates/krishiv-exec/src/operators/expand.rs @@ -0,0 +1,144 @@ +// 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>, + 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>) -> Result> { + 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 { + Ok(Box::new(Cursor(0))) + } + + fn execute(&self, input: &Batch, local: &mut LocalState) -> Result { + let cursor = local_mut::(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 { + 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 { + 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())) + } +} diff --git a/crates/krishiv-exec/src/operators/mod.rs b/crates/krishiv-exec/src/operators/mod.rs index f8cf3b0ad..da85c64c6 100644 --- a/crates/krishiv-exec/src/operators/mod.rs +++ b/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}; diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index 266642cbd..728028a00 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -30,7 +30,7 @@ use crate::join::{ }; use crate::operator::{Operator, Sink, Source}; use crate::operators::{ - FileWrite, FileWriteSink, Filter, FilterPushdown, Predicate, Project, ProjectItem, + Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, }; use crate::pipeline::{ExecNode, PipelineBuilder}; @@ -278,6 +278,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, limit } => { let input = self.node(input)?; let op = crate::operators::Limit::new(Arc::clone(&input.schema), *limit); @@ -547,6 +556,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::>>()?; @@ -556,6 +568,20 @@ pub(crate) fn project_op( } } +pub(crate) fn expand_op( + schema: &SchemaRef, + layout: &[AttrId], + projections: &[Vec], + nullable: &[bool], + v: &Values<'_>, +) -> Result> { + let projects = projections + .iter() + .map(|items| project_op(schema, layout, items, Some(nullable), v)) + .collect::>>()?; + Ok(Expand::try_new(projects)?) +} + pub(crate) fn replicate_op( schema: &SchemaRef, layout: &[AttrId], diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index f4f1ce667..2ca8958d9 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -51,7 +51,7 @@ pub(crate) use instantiate::node_ids; pub use instantiate::{SpillScopes, WorkerContext, instantiate}; pub(crate) use instantiate::{ Unbuilt, Values, aggregate_node, aggregate_spec, filter_op, hash_join, nested_loop_join, - project_op, replicate_op, scan_filter, sort_op, top_n_op, + expand_op, project_op, replicate_op, scan_filter, sort_op, top_n_op, }; use crate::aggregate::{AggregateExpr, AggregateOutput}; @@ -144,6 +144,20 @@ pub enum PhysNode { /// The subqueries `count` reads. subqueries: Vec, }, + /// Each row of `input` once per projection (`ROLLUP`, `CUBE`, + /// `GROUPING SETS`). + Expand { + /// The input. + input: Box, + /// The attribute of each input column. + layout: Vec, + /// The projections: the same output columns, each computed one way. + projections: Vec>, + /// Whether each output column is nullable. + nullable: Vec, + /// The subqueries the projections read. + subqueries: Vec, + }, /// `input` sorted by `keys` (columns of `input`). Sort { /// The input. @@ -252,6 +266,14 @@ pub enum ProjectItemPlan { /// The value. expr: Arc, }, + /// A column of nulls in the Arrow type of an input column (the cell an + /// expansion nulls). + NullOf { + /// The input column whose type it takes. + index: usize, + /// The output name. + name: String, + }, } /// Scalar or `EXISTS`. @@ -346,6 +368,7 @@ plan_data!(enum PhysNode { Project { input, layout, items, nullable, subqueries }, Limit { input, limit }, ReplicateRows { input, layout, count, subqueries }, + Expand { input, layout, projections, nullable, subqueries }, Sort { input, keys }, TopN { input, keys, limit }, Aggregate { input, aggregate }, @@ -358,7 +381,7 @@ plan_data!(struct ScanNode { id, scan, projection, filter, runtime_filters, row_ plan_data!(struct ScanFilter { layout, predicate, exact }); plan_data!(struct ScanRuntimeFilter { id, column, key_type }); plan_data!(struct JoinFilterPlan { id, key, options }); -plan_data!(enum ProjectItemPlan { Column { index, name }, Expr { name, expr } }); +plan_data!(enum ProjectItemPlan { Column { index, name }, Expr { name, expr }, NullOf { index, name } }); plan_data!( enum SubqueryValueKind { Scalar, @@ -386,6 +409,9 @@ impl PhysNode { | Self::Project { input, subqueries, .. } + | Self::Expand { + input, subqueries, .. + } | Self::ReplicateRows { input, subqueries, .. } => std::iter::once(input.as_ref()) @@ -412,6 +438,9 @@ impl PhysNode { | Self::Project { input, subqueries, .. } + | Self::Expand { + input, subqueries, .. + } | Self::ReplicateRows { input, subqueries, .. } => std::iter::once(input.as_mut()) diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index b7aadefb5..e9fcc721e 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -40,8 +40,8 @@ use crate::operator::{Operator, Source}; use crate::operators::{FilterPushdown, Limit, Project, ProjectItem, ScanPushdowns, SortKey, TopN}; use crate::physical::{ PhysNode, PhysicalPlan, ProjectItemPlan, ScanFilter, ScanId, ScanNode, ScanRuntimeFilter, - SubqueryId, SubqueryPlan, SubqueryValueKind, Values, filter_op, project_op, replicate_op, - scan_filter, sort_op, top_n_op, + SubqueryId, SubqueryPlan, SubqueryValueKind, Values, expand_op, filter_op, project_op, + replicate_op, scan_filter, sort_op, top_n_op, }; /// How far above the planner's first fresh attribute subquery ids start: @@ -348,6 +348,14 @@ impl<'a> Planner<'a> { let input = self.plan(input)?; self.replicate(input, count) } + PlanNode::Expand { + input, + projections, + output, + } => { + let input = self.plan(input)?; + self.expand(input, projections, output) + } PlanNode::OneRowRelation => Ok(one_row()), _ => { let text = p.to_string(); @@ -696,6 +704,92 @@ impl<'a> Planner<'a> { Ok(planned) } + /// Each row once per projection (`ROLLUP`, `CUBE`, `GROUPING SETS`): + /// plain columns pass through without a copy. + fn expand( + &mut self, + input: Planned, + projections: &[Vec>], + output: &[BoundColumn], + ) -> Result { + if projections.iter().any(|row| row.len() != output.len()) { + return Err(internal("an expansion row of the wrong width")); + } + // A column an expansion nulls keeps the Arrow type of the column + // another projection copies. + let copied: Vec> = (0..output.len()) + .map(|j| { + projections + .iter() + .find_map(|row| row.get(j).and_then(|e| column_of(e))) + .and_then(|c| input.column(c).ok()) + }) + .collect(); + let null_of = |j: usize, e: &Expr| { + let null = matches!(&e.node, ExprNode::Literal(Literal::Null(Some(_)))); + copied.get(j).copied().flatten().filter(|_| null) + }; + let computed: Vec<&Expr> = projections + .iter() + .flat_map(|row| row.iter().enumerate()) + .filter(|(j, e)| column_of(e).is_none() && null_of(*j, e).is_none()) + .map(|(_, e)| e.as_ref()) + .collect(); + let prepared = self.prepare(&computed)?; + let mut values = prepared.exprs.iter(); + let mut rows = Vec::with_capacity(projections.len()); + for row in projections { + let mut items = Vec::with_capacity(row.len()); + for (j, (e, out)) in row.iter().zip(output).enumerate() { + let name = out.name.clone(); + if let Some(c) = column_of(e) { + items.push(ProjectItemPlan::Column { + index: input.column(c)?, + name: Some(name), + }); + } else if let Some(index) = null_of(j, e) { + items.push(ProjectItemPlan::NullOf { index, name }); + } else { + let expr = values.next().ok_or_else(|| internal("expansion arity"))?; + items.push(ProjectItemPlan::Expr { + name, + expr: Arc::clone(expr), + }); + } + } + rows.push(items); + } + let nullable: Vec = output.iter().map(|c| c.nullable).collect(); + let layout: Vec = output.iter().map(|c| c.id).collect(); + let op = expand_op( + &input.schema, + &input.layout, + &rows, + &nullable, + &self.values(&prepared), + )?; + let names: Vec<&str> = output.iter().map(|c| c.name.as_str()).collect(); + let text = names.join(", "); + let mut planned = Self::attach( + input, + prepared, + op.schema(), + layout, + ("Expand", format!("{} projections [{text}]", projections.len())), + |input, layout, subqueries| PhysNode::Expand { + input, + layout, + projections: rows, + nullable, + subqueries, + }, + ); + // The rows are repeated with columns nulled: no runtime filter may + // reach a scan through the expansion. + planned.sites.clear(); + Ok(planned) + } + /// `SELECT exprs`: plain columns pass through without a copy. pub(crate) fn project(&mut self, input: Planned, exprs: &[Arc]) -> Result { let computed: Vec<&Expr> = exprs diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index cf56e0ef2..7fb5f23ab 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -76,6 +76,20 @@ impl Graph { grouping_refs(p, &mut need); child(input, Some(need)); } + PlanNode::Expand { + input, + projections, + output, + } => { + let mut need = Attrs::new(); + for j in expand_keep(output, req) { + projections + .iter() + .filter_map(|row| row.get(j)) + .for_each(|e| references_into(e, &mut need)); + } + child(input, Some(need)); + } PlanNode::SubqueryAlias { input, columns, .. } => { // A renaming alias keeps its input whole. child(input, columns.is_empty().then(|| req.clone())); @@ -197,6 +211,26 @@ impl Rewrite<'_> { replace(&rebuilt, PlanNode::Project { input, exprs: kept }) } } + PlanNode::Expand { + input, + projections, + output, + } => { + let kept = expand_keep(output, req); + if kept.len() == output.len() { + rebuilt + } else { + let pick = |row: &Vec>| -> Vec> { + kept.iter().filter_map(|&j| row.get(j).cloned()).collect() + }; + let node = PlanNode::Expand { + input: Arc::clone(input), + projections: projections.iter().map(pick).collect(), + output: kept.iter().filter_map(|&j| output.get(j).cloned()).collect(), + }; + replace(&rebuilt, node) + } + } PlanNode::Aggregate { input, grouping, @@ -304,6 +338,18 @@ fn keep(exprs: &[Arc], req: &Attrs) -> Vec> { kept } +/// The output columns of an expansion that are read (by position), keeping +/// at least one so the row count survives. +fn expand_keep(output: &[BoundColumn], req: &Attrs) -> Vec { + let mut kept: Vec = (0..output.len()) + .filter(|&j| req.contains(&output[j].id)) + .collect(); + if kept.is_empty() && !output.is_empty() { + kept.push(0); + } + kept +} + fn refs(exprs: &[Arc]) -> Attrs { let mut out = Attrs::new(); exprs.iter().for_each(|e| references_into(e, &mut out)); diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index fd00375b1..d96d5abf1 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -239,6 +239,14 @@ impl<'a> Estimator<'a> { // Each distinct row of an `INTERSECT ALL` / `EXCEPT ALL` count: // no count is known, the input's rows bound it from below. | PlanNode::ReplicateRows { input, .. } => Est::clone(&self.plan(input)), + // Every input row once per projection. + PlanNode::Expand { + input, projections, .. + } => { + let est = Est::clone(&self.plan(input)); + let rows = est.rows * to_f64(u64::try_from(projections.len()).unwrap_or(u64::MAX)); + est.with_rows(rows) + } PlanNode::Union { inputs, .. } => { let rows = inputs.iter().map(|i| self.plan(i).rows).sum(); Est::unknown(rows, ids()) diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs new file mode 100644 index 000000000..c6d34c79e --- /dev/null +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -0,0 +1,692 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Grouping analytics: `GROUP BY ROLLUP`, `CUBE` and `GROUPING SETS`, and the +//! `grouping` and `grouping_id` functions. +//! +//! The rewrite is Spark's (the `ResolveGroupingAnalytics` rule and the +//! `Expand` operator): every distinct grouping expression gets a nullable +//! copy `g`, `Expand` repeats each input row once per grouping set (`g` is +//! the expression for a set that has it and `NULL` for a set that has not) +//! with the set's `spark_grouping_id`, and an ordinary aggregate groups by +//! the copies and the id. Bit `n - 1 - j` of the id is 1 when the `j`-th +//! distinct grouping expression (in order of first appearance) is absent +//! from the set. Select items, `HAVING` and `ORDER BY` read the copies where +//! they name a grouping expression (aggregate calls keep reading the +//! input's own columns, which stay in the expansion), `grouping(e)` is bit +//! `e` of the id (a `TINYINT`) and `grouping_id(..)` is the id (a `BIGINT`). +//! +//! Sets that are equal as sets stay separate rows (`GROUPING SETS ((a), (a))` +//! has two rows per group): the expansion then also numbers its sets +//! (`_gen_grouping_pos`) and the aggregate groups by that too. A single +//! empty set over no grouping expression is a grand total, an aggregate +//! without `GROUP BY` (one row over empty input) in which `grouping_id()` is 0. +//! The virtual column `grouping__id` names the id. +//! +//! Exact behaviour, Spark 4.2: `tests/spark-golden/v4.2.0` (`group-analytics`, +//! `grouping_set`, `postgreSQL/groupingsets`). + +use std::collections::HashMap; +use std::sync::Arc; + +use krishiv_common::error::{ + GroupingColumnMismatch, GroupingIdColumnMismatch, GroupingSizeLimitExceeded, + UnsupportedGroupingExpression, WrongNumArgsWithoutSuggestion, +}; +use krishiv_types::DataType; + +use super::bind::unsupported; +use super::exprs::{Bound, Cx, pretty, semantic_key}; +use super::ops::col_expr; +use super::resolve::{Attr, Env, Rel, Resolver}; +use super::{AnalysisError, ArgType}; +use crate::{ + BoundColumn, CastMode, Expr, ExprNode, FunctionCall, FunctionKind, Grouping, Literal, + LogicalPlan, Origin, PlanNode, ResolvedExpr, +}; + +/// The most grouping expressions one `GROUP BY` takes: the grouping id is a +/// `BIGINT` with one bit each. +const MAX_GROUPING: usize = 63; + +/// `grouping` or `grouping_id`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum GroupingFn { + Grouping, + GroupingId, +} + +impl GroupingFn { + pub(super) fn from_name(name: &str) -> Option { + if name.eq_ignore_ascii_case("grouping") { + Some(Self::Grouping) + } else if name.eq_ignore_ascii_case("grouping_id") { + Some(Self::GroupingId) + } else { + None + } + } + + const fn name(self) -> &'static str { + match self { + Self::Grouping => "grouping", + Self::GroupingId => "grouping_id", + } + } + + const fn data_type(self) -> DataType { + match self { + Self::Grouping => DataType::Byte, + Self::GroupingId => DataType::Long, + } + } +} + +/// One distinct grouping expression. +pub(super) struct GroupExpr { + /// Its semantic key. + key: u128, + /// Its attribute in the expansion: a nullable copy. + pub(super) attr: Attr, +} + +/// The expansion an aggregate reads. +pub(super) struct ExpandPlan { + pub(super) plan: Arc, + /// The span of attribute ids the expansion defines. + pub(super) ids: Option<(u64, u64)>, +} + +/// An aggregate's grouping sets. +pub(super) struct Analytics { + /// The distinct grouping expressions, in order of first appearance. + pub(super) exprs: Vec, + /// The grouping id as the query names it (`grouping__id`); `None` for a + /// grand total, where it is 0. + gid: Option, + /// The expansion; `None` for a grand total. + pub(super) expand: Option, +} + +impl Analytics { + /// The attributes `HAVING` and `ORDER BY` may name: the grouping + /// expressions' copies and `grouping__id`. + pub(super) fn keys(&self) -> Vec { + self.exprs + .iter() + .map(|g| g.attr.clone()) + .chain(self.gid.iter().cloned()) + .collect() + } + + /// The position of `e` among the grouping expressions: a copy of one, or + /// an expression equal to one. + fn position(&self, e: &Expr, memo: &mut HashMap) -> Option { + if let ExprNode::Resolved(ResolvedExpr::Column(c)) = &e.node + && let Some(j) = self.exprs.iter().position(|g| g.attr.col.id == c.id) + { + return Some(j); + } + let key = semantic_key(e, memo); + self.exprs.iter().position(|g| g.key == key) + } +} + +fn literal(l: Literal) -> Arc { + Arc::new(Expr::new(ExprNode::Literal(l))) +} + +fn arg_type(data_type: DataType, nullable: bool) -> ArgType { + ArgType { + data_type: Some(data_type), + nullable, + } +} + +fn bound_of(col: &BoundColumn) -> Bound { + Bound { + expr: col_expr(col, Origin::UNKNOWN), + ty: arg_type(col.data_type.clone(), col.nullable), + agg: None, + outer: None, + lateral: Vec::new(), + shown: None, + } +} + +/// Spark's text for an expression in a grouping error: a column carries its +/// attribute id. +fn with_id(e: &Expr) -> String { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Column(c)) => format!("{}#{}", c.name, c.id.0), + _ => pretty(e), + } +} + +/// Whether a subquery is anywhere in `e`. +fn has_subquery(e: &Arc) -> bool { + let mut stack = vec![Arc::clone(e)]; + while let Some(x) = stack.pop() { + if matches!(x.node, ExprNode::Subquery { .. }) { + return true; + } + x.for_each_child(&mut |c| stack.push(Arc::clone(c))); + } + false +} + +/// The shape of a grouping: all its expressions in order, and each grouping +/// set as positions into them. +fn shape(grouping: &Grouping) -> (Vec>, Vec>) { + match grouping { + Grouping::GroupBy(v) => (v.clone(), vec![(0..v.len()).collect()]), + Grouping::Rollup(v) => ( + v.clone(), + (0..=v.len()).rev().map(|n| (0..n).collect()).collect(), + ), + Grouping::Cube(v) => { + // The full set first, then down to the empty one. + let n = v.len(); + let sets = (0..1_usize << n) + .rev() + .map(|m| (0..n).filter(|i| m >> (n - 1 - i) & 1 == 1).collect()) + .collect(); + (v.clone(), sets) + } + Grouping::GroupingSets { exprs, sets } => { + let mut flat = exprs.clone(); + let mut out = Vec::with_capacity(sets.len()); + for set in sets { + out.push((flat.len()..flat.len() + set.len()).collect()); + flat.extend(set.iter().cloned()); + } + (flat, out) + } + } +} + +impl Resolver<'_> { + /// `grouping(e)` or `grouping_id(..)` in a select list, `HAVING` or + /// `ORDER BY`: a call the aggregate resolves once it knows its grouping + /// sets ([`Self::resolve_grouping`]). Anywhere else it is + /// `UNSUPPORTED_GROUPING_EXPRESSION`. + #[inline(never)] + pub(super) fn grouping_call( + &mut self, + which: GroupingFn, + e: &Arc, + kids: &[Bound], + cx: &Cx<'_>, + ) -> Result { + // Spark has no source position for the error in `HAVING` and + // `ORDER BY`. + let origin = if matches!(cx.op, "Filter" | "Sort") { + Origin::UNKNOWN + } else { + e.origin + }; + if !cx.grouping { + return Err(AnalysisError::new(UnsupportedGroupingExpression {}, origin)); + } + if which == GroupingFn::Grouping && kids.len() != 1 { + return Err(AnalysisError::new( + WrongNumArgsWithoutSuggestion { + function_name: "`grouping`".to_owned(), + expected_num: "1".to_owned(), + actual_num: kids.len().to_string(), + docroot: super::exprs::SPARK_DOC_ROOT.to_owned(), + }, + e.origin, + )); + } + let call = FunctionCall { + name: which.name().to_owned(), + args: kids.iter().map(|k| Arc::clone(&k.expr)).collect(), + distinct: false, + filter: None, + }; + Ok(Bound { + expr: Arc::new(Expr::new(ExprNode::Function(call)).with_origin(e.origin)), + ty: arg_type(which.data_type(), false), + agg: kids.iter().find_map(|k| k.agg), + outer: kids.iter().filter_map(|k| k.outer).min(), + lateral: kids.iter().flat_map(|k| k.lateral.clone()).collect(), + shown: None, + }) + } + + /// A resolved `Expand` read again: its expressions bound over its + /// input, its output as it states it. + #[inline(never)] + pub(super) fn expand_node( + &mut self, + origin: Origin, + input: &Arc, + projections: &[Vec>], + output: &[BoundColumn], + env: &Arc, + ) -> Result { + let c = self.child(input, env)?; + let rels = [&c]; + let cx = Self::cx(&c.out, &rels, env, "Expand"); + let mut bound = Vec::with_capacity(projections.len()); + for row in projections { + let mut exprs = Vec::with_capacity(row.len()); + for e in row { + exprs.push(self.bind(e, &cx)?.expr); + } + bound.push(exprs); + } + let out = output + .iter() + .map(|col| Attr { + col: col.clone(), + qual: Arc::from(Vec::new()), + source: None, + }) + .collect(); + let node = PlanNode::Expand { + input: Arc::clone(&c.plan), + projections: bound, + output: output.to_vec(), + }; + Ok(self.make(origin, node, out, &[&c])) + } + + /// `e` with its `grouping` and `grouping_id` calls resolved against + /// `an`. + #[inline(never)] + pub(super) fn resolve_grouping( + &mut self, + e: &Arc, + an: &Analytics, + env: &Arc, + ) -> Result, AnalysisError> { + self.rewrite(e, an, false, env, &mut HashMap::new()) + } + + /// An aggregate over `ROLLUP`, `CUBE` or `GROUPING SETS`. + #[inline(never)] + pub(super) fn grouping_sets( + &mut self, + p: &LogicalPlan, + input: &Arc, + grouping: &Grouping, + aggregates: &[Arc], + env: &Arc, + ) -> Result { + let c = self.child(input, env)?; + let (flat, sets) = shape(grouping); + self.check_star_ordinals(&flat, aggregates)?; + // A single empty set over nothing is a grand total. + let grand = flat.is_empty() && sets.len() == 1; + let gid_col = (!grand).then(|| { + self.attr( + "spark_grouping_id", + &arg_type(DataType::Long, false), + ) + }); + let gid = gid_col.as_ref().map(|col| Attr { + col: BoundColumn { + name: "grouping__id".to_owned(), + ..col.clone() + }, + qual: Arc::from(Vec::new()), + source: None, + }); + let keys: Vec = gid.iter().cloned().collect(); + let items = self.select_items_in(&c, aggregates, env, "Aggregate", Some(&keys))?; + let bound = self.group_keys(&c, &flat, aggregates, &items, env)?; + let mut memo = HashMap::new(); + // The distinct expressions, and the sets as positions among them. + let mut distinct: Vec = Vec::new(); + let mut keys_of: Vec = Vec::new(); + let mut at = Vec::with_capacity(bound.len()); + for (i, b) in bound.iter().enumerate() { + let key = semantic_key(&b.expr, &mut memo); + let j = match keys_of.iter().position(|k| *k == key) { + Some(j) => j, + None => { + keys_of.push(key); + distinct.push(i); + keys_of.len() - 1 + } + }; + at.push(j); + } + let n = distinct.len(); + if n > MAX_GROUPING { + return Err(AnalysisError::new( + GroupingSizeLimitExceeded { + max_size: MAX_GROUPING.to_string(), + }, + p.origin, + )); + } + let sets: Vec> = sets + .into_iter() + .map(|s| { + let mut s: Vec = s.into_iter().map(|i| at[i]).collect(); + s.sort_unstable(); + s.dedup(); + s + }) + .collect(); + if grand { + let an = Analytics { + exprs: Vec::new(), + gid: None, + expand: None, + }; + let items = self.rewrite_items(items, &an, env)?; + return self.build_aggregate(p.origin, &c, Vec::new(), &items, true, Some(Arc::new(an))); + } + let Some(gid_col) = gid_col else { + return Err(unsupported("grouping sets without an id", p.origin)); + }; + let numbered = (0..sets.len()).any(|i| sets[..i].contains(&sets[i])); + // The expansion's input: the child's columns, and the grouping + // expressions that are not one of them. + let mut passed: Vec> = c.out.iter().map(|a| col_expr(&a.col, Origin::UNKNOWN)).collect(); + let mut passed_out: Vec = c.out.to_vec(); + let mut value: Vec = Vec::with_capacity(n); + let mut copies: Vec = Vec::with_capacity(n); + let mut computed = false; + for (j, &i) in distinct.iter().enumerate() { + let b = &bound[i]; + let column = match &b.expr.node { + ExprNode::Resolved(ResolvedExpr::Column(col)) => { + c.out.iter().find(|a| a.col.id == col.id).cloned() + } + _ => None, + }; + let name = match &column { + Some(a) => a.col.name.clone(), + None => pretty(&b.shown()), + }; + let v = match &column { + Some(a) => a.col.clone(), + None => { + let v = self.attr(&name, &b.ty); + passed.push(Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { + expr: Arc::clone(&b.expr), + output: v.clone(), + })))); + passed_out.push(Attr { + col: v.clone(), + qual: Arc::from(Vec::new()), + source: None, + }); + computed = true; + v + } + }; + let ty = arg_type(v.data_type.clone(), true); + let copy = self.attr(&name, &ty); + copies.push(GroupExpr { + key: keys_of[j], + attr: Attr { + col: copy, + qual: column.as_ref().map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), + source: column.and_then(|a| a.source), + }, + }); + value.push(v); + } + let below = if computed { + let out = passed_out.clone(); + self.projection(Origin::UNKNOWN, &c, passed, out) + } else { + c.clone() + }; + let pos_col = numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); + let mut output: Vec = passed_out.iter().map(|a| a.col.clone()).collect(); + output.extend(copies.iter().map(|g| g.attr.col.clone())); + output.push(gid_col.clone()); + output.extend(pos_col.iter().cloned()); + let mut projections = Vec::with_capacity(sets.len()); + for (si, set) in sets.iter().enumerate() { + let mut row: Vec> = passed_out + .iter() + .map(|a| col_expr(&a.col, Origin::UNKNOWN)) + .collect(); + let mut mask: i64 = 0; + for (j, v) in value.iter().enumerate() { + if set.contains(&j) { + row.push(col_expr(v, Origin::UNKNOWN)); + } else { + row.push(literal(Literal::Null(Some(v.data_type.clone())))); + mask |= 1 << (n - 1 - j); + } + } + row.push(literal(Literal::Long(mask))); + if pos_col.is_some() { + row.push(literal(Literal::Integer( + i32::try_from(si).map_err(|_| unsupported("that many grouping sets", p.origin))?, + ))); + } + projections.push(row); + } + let out_attrs: Vec = passed_out + .iter() + .cloned() + .chain(copies.iter().map(|g| g.attr.clone())) + .chain(std::iter::once(Attr { + col: gid_col.clone(), + qual: Arc::from(Vec::new()), + source: None, + })) + .chain(pos_col.iter().map(|col| Attr { + col: col.clone(), + qual: Arc::from(Vec::new()), + source: None, + })) + .collect(); + let node = PlanNode::Expand { + input: Arc::clone(&below.plan), + projections, + output, + }; + let expand = self.make(Origin::UNKNOWN, node, out_attrs, &[&below]); + let mut group: Vec = copies.iter().map(|g| bound_of(&g.attr.col)).collect(); + group.push(bound_of(&gid_col)); + group.extend(pos_col.iter().map(bound_of)); + let an = Analytics { + exprs: copies, + gid, + expand: Some(ExpandPlan { + plan: Arc::clone(&expand.plan), + ids: expand.ids, + }), + }; + let items = self.rewrite_items(items, &an, env)?; + self.build_aggregate(p.origin, &c, group, &items, false, Some(Arc::new(an))) + } + + /// The select `items` over the grouping sets: grouping expressions read + /// their copies, `grouping` and `grouping_id` are resolved. + #[inline(never)] + fn rewrite_items( + &mut self, + items: Vec, + an: &Analytics, + env: &Arc, + ) -> Result, AnalysisError> { + let mut memo = HashMap::new(); + let mut out = Vec::with_capacity(items.len()); + for b in items { + let expr = self.rewrite(&b.expr, an, true, env, &mut memo)?; + let ty = if Arc::ptr_eq(&expr, &b.expr) { + b.ty.clone() + } else { + self.type_of(&expr, env)? + }; + out.push(Bound { expr, ty, ..b }); + } + Ok(out) + } + + /// The type of the resolved `e`. + fn type_of(&mut self, e: &Arc, env: &Arc) -> Result { + let inner = match &e.node { + ExprNode::Resolved(ResolvedExpr::Alias { expr, .. }) => expr, + _ => e, + }; + Ok(self.bind(inner, &Self::cx(&[], &[], env, "Aggregate"))?.ty) + } + + /// `e` over the grouping sets: with `replace`, a subtree that is a + /// grouping expression reads its copy (aggregate calls and subqueries + /// are left as they are); `grouping` and `grouping_id` are resolved. + /// Nodes above a change are typed again, as the copies are nullable. + #[inline(never)] + fn rewrite( + &mut self, + e: &Arc, + an: &Analytics, + replace: bool, + env: &Arc, + memo: &mut HashMap, + ) -> Result, AnalysisError> { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) if f.kind == FunctionKind::Aggregate => { + return Ok(Arc::clone(e)); + } + ExprNode::Subquery { .. } => return Ok(Arc::clone(e)), + ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some() => { + return self.grouping_value(f, an, env, memo); + } + ExprNode::Resolved(ResolvedExpr::Alias { expr, output }) => { + let inner = self.rewrite(expr, an, replace, env, memo)?; + if Arc::ptr_eq(&inner, expr) { + return Ok(Arc::clone(e)); + } + let ty = self.type_of(&inner, env)?; + let output = BoundColumn { + nullable: ty.nullable, + ..output.clone() + }; + return Ok(Arc::new( + Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { + expr: inner, + output, + })) + .with_origin(e.origin), + )); + } + _ => {} + } + if replace && let Some(j) = an.position(e, memo) { + return Ok(col_expr(&an.exprs[j].attr.col, e.origin)); + } + let rebuilt = Arc::clone(e).map_children(|c| self.rewrite(&c, an, replace, env, memo))?; + if Arc::ptr_eq(&rebuilt, e) { + return Ok(rebuilt); + } + self.retype(rebuilt, env) + } + + /// `e`, whose children changed, typed again. + fn retype(&mut self, e: Arc, env: &Arc) -> Result, AnalysisError> { + let origin = e.origin; + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) if !has_subquery(&e) => { + self.call(f.id.name(), f.args.clone(), env, origin) + } + ExprNode::Resolved(ResolvedExpr::Function(f)) => { + // A call over a subquery keeps its type, nullable. + let mut f = f.clone(); + f.nullable = true; + Ok(Arc::new( + Expr::new(ExprNode::Resolved(ResolvedExpr::Function(f))).with_origin(origin), + )) + } + ExprNode::Cast { .. } if !has_subquery(&e) => { + Ok(self.bind(&e, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + } + _ => Ok(e), + } + } + + /// The value of a `grouping` or `grouping_id` call. + fn grouping_value( + &mut self, + f: &FunctionCall, + an: &Analytics, + env: &Arc, + memo: &mut HashMap, + ) -> Result, AnalysisError> { + let n = an.exprs.len(); + let id = match &an.gid { + Some(g) => col_expr(&g.col, Origin::UNKNOWN), + None => literal(Literal::Long(0)), + }; + let positions: Vec> = + f.args.iter().map(|a| an.position(a, memo)).collect(); + match GroupingFn::from_name(&f.name) { + Some(GroupingFn::GroupingId) => { + let exact = positions.len() == n + && positions.iter().enumerate().all(|(i, p)| *p == Some(i)); + if !f.args.is_empty() && !exact { + return Err(AnalysisError::new( + GroupingIdColumnMismatch { + grouping_id_column: f + .args + .iter() + .map(|a| with_id(a)) + .collect::>() + .join(","), + group_by_columns: an + .exprs + .iter() + .map(|g| format!("{}#{}", g.attr.col.name, g.attr.col.id.0)) + .collect::>() + .join(","), + }, + Origin::UNKNOWN, + )); + } + Ok(id) + } + _ => { + let (Some(arg), Some(Some(j))) = (f.args.first(), positions.first()) else { + return Err(AnalysisError::new( + GroupingColumnMismatch { + grouping: f.args.first().map(|a| with_id(a)).unwrap_or_default(), + grouping_columns: an + .exprs + .iter() + .map(|g| format!("{}#{}", g.attr.col.name, g.attr.col.id.0)) + .collect::>() + .join(","), + }, + Origin::UNKNOWN, + )); + }; + let _ = arg; + // Bit `n - 1 - j` of the id, as a `TINYINT`. + let shift = n - 1 - j; + let high = if shift == 0 { + id + } else { + self.call( + "div", + vec![id, literal(Literal::Long(1_i64 << shift))], + env, + Origin::UNKNOWN, + )? + }; + let bit = self.call( + "%", + vec![high, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )?; + let cast = Arc::new(Expr::new(ExprNode::Cast { + expr: bit, + to: DataType::Byte, + mode: CastMode::Default, + })); + Ok(self.bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + } + } + } +} diff --git a/crates/krishiv-plan/src/analyzer/bind.rs b/crates/krishiv-plan/src/analyzer/bind.rs index 3a8dcc69b..02f44f54f 100644 --- a/crates/krishiv-plan/src/analyzer/bind.rs +++ b/crates/krishiv-plan/src/analyzer/bind.rs @@ -227,6 +227,10 @@ fn check_function(ctx: &AnalysisContext, name: &str, origin: Origin) -> Result<( if crate::Builtin::from_name(name).is_some() { return Ok(()); } + // `grouping` and `grouping_id` are the analyzer's too (grouping sets). + if super::analytics::GroupingFn::from_name(name).is_some() { + return Ok(()); + } match ctx .functions .lookup(std::slice::from_ref(&name.to_owned()), ctx.profile) diff --git a/crates/krishiv-plan/src/analyzer/exprs.rs b/crates/krishiv-plan/src/analyzer/exprs.rs index 6eb5d5a27..b16b69299 100644 --- a/crates/krishiv-plan/src/analyzer/exprs.rs +++ b/crates/krishiv-plan/src/analyzer/exprs.rs @@ -21,6 +21,7 @@ use krishiv_common::error::{ use krishiv_types::DataType; use twox_hash::XxHash3_128; +use super::analytics::GroupingFn; use super::bind::{addr, same, unsupported}; use super::catalog::quoted; use super::coerce::{self, Arg, BadCast, Mismatch, ordinal, quoted_type}; @@ -50,6 +51,10 @@ pub(super) struct Cx<'s> { /// Over an aggregate: its input, where aggregate-function arguments /// resolve. pub(super) agg_scope: Option<&'s [Attr]>, + /// Whether `grouping(..)` and `grouping_id(..)` may appear: in the + /// select list, `HAVING` and `ORDER BY` of an aggregate over grouping + /// sets (see `analytics`). + pub(super) grouping: bool, /// A name `scope` holds several times is looked up further rather /// than ambiguous (Spark's `ORDER BY` over a select list). pub(super) lenient: bool, @@ -204,7 +209,7 @@ fn arg_of(b: &Bound) -> Arg<'_> { } /// Spark's documentation root, as its errors cite it. -const SPARK_DOC_ROOT: &str = "https://spark.apache.org/docs/latest"; +pub(super) const SPARK_DOC_ROOT: &str = "https://spark.apache.org/docs/latest"; fn column(c: &BoundColumn) -> ArgType { arg_type(c.data_type.clone(), c.nullable) @@ -444,6 +449,9 @@ impl Resolver<'_> { e.origin, )); } + if let Some(which) = GroupingFn::from_name(&f.name) { + return self.grouping_call(which, e, &kids[..n], cx); + } if let Some(op) = Builtin::from_name(&f.name) { return self.operator_call(op, e.origin, &kids[..n]); } diff --git a/crates/krishiv-plan/src/analyzer/mod.rs b/crates/krishiv-plan/src/analyzer/mod.rs index 44bca0668..560064ff2 100644 --- a/crates/krishiv-plan/src/analyzer/mod.rs +++ b/crates/krishiv-plan/src/analyzer/mod.rs @@ -62,6 +62,7 @@ //! provider's text. Catalog functions are not policy-checked yet (no //! catalog functions exist before P3). +mod analytics; mod bind; mod builtins; mod catalog; diff --git a/crates/krishiv-plan/src/analyzer/ops.rs b/crates/krishiv-plan/src/analyzer/ops.rs index 1e932bd85..8238aec0e 100644 --- a/crates/krishiv-plan/src/analyzer/ops.rs +++ b/crates/krishiv-plan/src/analyzer/ops.rs @@ -14,6 +14,7 @@ use krishiv_common::error::{ UnsupportedFeatureLateralColumnAliasInGroupBy, }; +use super::analytics::Analytics; use super::bind::{BoundTable, same, unsupported}; use super::catalog::quoted; use super::coerce::quoted_type; @@ -244,7 +245,7 @@ impl Resolver<'_> { ) -> Result { let items = self.select_items(c, exprs, env, "Aggregate")?; if items.iter().any(|b| b.agg.is_some()) { - return self.build_aggregate(p.origin, c, Vec::new(), &items, true); + return self.build_aggregate(p.origin, c, Vec::new(), &items, true, None); } Ok(self.lateral_project(p.origin, c, &items, &[])) } @@ -258,6 +259,20 @@ impl Resolver<'_> { exprs: &[Arc], env: &Arc, op: &'static str, + ) -> Result, AnalysisError> { + self.select_items_in(c, exprs, env, op, None) + } + + /// [`Self::select_items`] of an aggregate over grouping sets, which + /// may call `grouping` and `grouping_id` and name `grouping__id` + /// (`extra`, searched after the input). + pub(super) fn select_items_in( + &mut self, + c: &Rel, + exprs: &[Arc], + env: &Arc, + op: &'static str, + extra: Option<&[Attr]>, ) -> Result, AnalysisError> { let rels = [c]; let mut items = Vec::new(); @@ -274,6 +289,8 @@ impl Resolver<'_> { let b = { let mut cx = Self::cx(&c.out, &rels, env, op); cx.lateral = lca.then_some(&lateral[..]); + cx.extra = extra; + cx.grouping = extra.is_some(); self.bind(e, &cx)? }; let b = self.named(b)?; @@ -660,10 +677,16 @@ impl Resolver<'_> { let mut cx = Self::cx(&agg.out, &rels, env, op); cx.extra = Some(&info.keys); cx.agg_scope = Some(&info.input.out); + cx.grouping = info.analytics.is_some(); let mut bounds = Vec::new(); for e in exprs { bounds.push(self.bind(e, &cx)?); } + if let Some(an) = &info.analytics { + for b in &mut bounds { + b.expr = self.resolve_grouping(&b.expr, an, env)?; + } + } let mut memo = HashMap::new(); let mut known: Vec<(u128, Attr)> = info .aggregates @@ -754,8 +777,14 @@ impl Resolver<'_> { /// aliases if it has any; `pass` (attributes of the aggregate) are /// output after the select items. fn aggregate_rel(&self, origin: Origin, info: AggInfo, pass: &[Attr]) -> Rel { + // Over grouping sets, the aggregate reads their expansion; the + // attributes it defines are part of the aggregate's id span. + let (input, defined) = match info.analytics.as_deref().and_then(|a| a.expand.as_ref()) { + Some(e) => (Arc::clone(&e.plan), e.ids), + None => (Arc::clone(&info.input.plan), None), + }; let node = PlanNode::Aggregate { - input: Arc::clone(&info.input.plan), + input, grouping: Grouping::GroupBy(info.grouping.clone()), aggregates: info.aggregates.clone(), }; @@ -766,6 +795,7 @@ impl Resolver<'_> { self.lateral_project(origin, &a, top, pass) } }; + r.ids = super::resolve::hull(r.ids, defined); let info = Arc::new(info); r.widen = Some(Arc::new(Widen::Aggregate(Arc::clone(&info)))); r.agg = Some(info); @@ -781,38 +811,66 @@ impl Resolver<'_> { env: &Arc, ) -> Result { let Grouping::GroupBy(keys) = grouping else { - return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", p.origin)); + return self.grouping_sets(p, input, grouping, aggregates, env); }; let c = self.child(input, env)?; - let conf = self.ctx.conf; - let ordinal = |k: &Expr| match &k.node { - ExprNode::Literal(Literal::Integer(n)) if conf.group_by_ordinal => Some(*n), - _ => None, - }; + self.check_star_ordinals(keys, aggregates)?; + // Spark resolves the select list first, then ordinals, then the + // other keys. + let items = self.select_items(&c, aggregates, env, "Aggregate")?; + let group = self.group_keys(&c, keys, aggregates, &items, env)?; + self.build_aggregate(p.origin, &c, group, &items, false, None) + } + + /// `SELECT *` with `GROUP BY n` is an error (Spark's + /// `STAR_GROUP_BY_POS`). + pub(super) fn check_star_ordinals( + &self, + keys: &[Arc], + aggregates: &[Arc], + ) -> Result<(), AnalysisError> { let star = aggregates .iter() .any(|e| matches!(e.node, ExprNode::Star(_))); - if star && keys.iter().any(|k| ordinal(k).is_some()) { + if star && keys.iter().any(|k| self.ordinal(k).is_some()) { return Err(AnalysisError::new(StarGroupByPos {}, Origin::UNKNOWN)); } - // Spark resolves the select list first, then ordinals, then the - // other keys. - let items = self.select_items(&c, aggregates, env, "Aggregate")?; + Ok(()) + } + + /// `GROUP BY n` with ordinals on: `n`. + fn ordinal(&self, k: &Expr) -> Option { + match &k.node { + ExprNode::Literal(Literal::Integer(n)) if self.ctx.conf.group_by_ordinal => Some(*n), + _ => None, + } + } + + /// The grouping expressions `keys` over `c`: an ordinal names a select + /// item, other names may be select-list aliases. + pub(super) fn group_keys( + &mut self, + c: &Rel, + keys: &[Arc], + aggregates: &[Arc], + items: &[Bound], + env: &Arc, + ) -> Result, AnalysisError> { let mut group: Vec> = vec![None; keys.len()]; for (slot, k) in group.iter_mut().zip(keys) { - if let Some(n) = ordinal(k) { - *slot = Some(self.group_ordinal(n, k.origin, &items, &c)?); + if let Some(n) = self.ordinal(k) { + *slot = Some(self.group_ordinal(n, k.origin, items, c)?); } } let aliases: Vec<(String, Arc)> = aggregates .iter() - .filter(|_| conf.group_by_aliases) + .filter(|_| self.ctx.conf.group_by_aliases) .filter_map(|e| match &e.node { ExprNode::Alias { expr, name } => Some((name.clone(), Arc::clone(expr))), _ => None, }) .collect(); - let rels = [&c]; + let rels = [c]; for (slot, k) in group.iter_mut().zip(keys) { if slot.is_none() { let mut cx = Self::cx(&c.out, &rels, env, "Aggregate"); @@ -820,8 +878,7 @@ impl Resolver<'_> { *slot = Some(self.bind(k, &cx)?); } } - let group = group.into_iter().flatten().collect(); - self.build_aggregate(p.origin, &c, group, &items, false) + Ok(group.into_iter().flatten().collect()) } /// `GROUP BY n`: the `n`-th select item, unaliased. @@ -885,6 +942,7 @@ impl Resolver<'_> { group: Vec, items: &[Bound], global: bool, + analytics: Option>, ) -> Result { let mut memo = HashMap::new(); let gkeys: Vec<(u128, ArgType)> = group @@ -898,22 +956,33 @@ impl Resolver<'_> { .cloned() .collect(); self.check_aggregation(&base, &key_set, &mut memo, origin, global)?; - let keys: Vec = group - .iter() - .filter_map(|b| match &b.expr.node { - ExprNode::Resolved(ResolvedExpr::Column(col)) => { - c.out.iter().find(|a| a.col.id == col.id).cloned() - } - _ => None, - }) - .collect(); + // Over grouping sets, the grouping attributes are the expansion's. + let (keys, scope): (Vec, Vec) = match &analytics { + Some(a) => ( + a.keys(), + c.out.iter().chain(a.exprs.iter().map(|g| &g.attr)).cloned().collect(), + ), + None => ( + group + .iter() + .filter_map(|b| match &b.expr.node { + ExprNode::Resolved(ResolvedExpr::Column(col)) => { + c.out.iter().find(|a| a.col.id == col.id).cloned() + } + _ => None, + }) + .collect(), + c.out.to_vec(), + ), + }; let mut info = AggInfo { input: c.clone(), grouping: group.into_iter().map(|b| b.expr).collect(), aggregates: base.iter().map(|b| Arc::clone(&b.expr)).collect(), - out: base.iter().map(|b| out_attr(b, &c.out)).collect(), + out: base.iter().map(|b| out_attr(b, &scope)).collect(), keys, top: None, + analytics, }; if base.len() == items.len() { return Ok(self.aggregate_rel(origin, info, &[])); @@ -925,7 +994,7 @@ impl Resolver<'_> { let mut top = Vec::new(); for b in items { if b.lateral.is_empty() { - let a = out_attr(b, &c.out); + let a = out_attr(b, &scope); top.push(Bound { expr: col_expr(&a.col, b.expr.origin), ..b.clone() diff --git a/crates/krishiv-plan/src/analyzer/resolve.rs b/crates/krishiv-plan/src/analyzer/resolve.rs index da70c8c97..0c6f4f564 100644 --- a/crates/krishiv-plan/src/analyzer/resolve.rs +++ b/crates/krishiv-plan/src/analyzer/resolve.rs @@ -39,6 +39,9 @@ pub(super) struct AggInfo { /// above the aggregate: every select item, over the aggregate's output /// (`None` when no item does). pub(super) top: Option>, + /// The grouping sets, for an aggregate over `ROLLUP`, `CUBE` or + /// `GROUPING SETS`: its input is then their expansion. + pub(super) analytics: Option>, } /// How a relation gets back columns it dropped: Spark adds the attributes @@ -193,7 +196,7 @@ fn first_free_id(root: &LogicalPlan) -> Result { Ok(next) } -fn hull(a: Option<(u64, u64)>, b: Option<(u64, u64)>) -> Option<(u64, u64)> { +pub(super) fn hull(a: Option<(u64, u64)>, b: Option<(u64, u64)>) -> Option<(u64, u64)> { match (a, b) { (Some(x), Some(y)) => Some((x.0.min(y.0), x.1.max(y.1))), (x, None) | (None, x) => x, @@ -446,6 +449,7 @@ impl Resolver<'_> { aliases: &[], extra: None, agg_scope: None, + grouping: false, lenient: false, outer: true, env, @@ -624,6 +628,11 @@ impl Resolver<'_> { }; Ok(self.make(o, node, c.out.to_vec(), &[&c])) } + PlanNode::Expand { + input, + projections, + output, + } => self.expand_node(o, input, projections, output, env), PlanNode::Distinct { input } => self.distinct(p, input, env), PlanNode::Hint { input, diff --git a/crates/krishiv-plan/src/analyzer/setops.rs b/crates/krishiv-plan/src/analyzer/setops.rs index 9e25f448f..0ac2ae780 100644 --- a/crates/krishiv-plan/src/analyzer/setops.rs +++ b/crates/krishiv-plan/src/analyzer/setops.rs @@ -387,7 +387,7 @@ impl Resolver<'_> { exprs.push(alias(sum, name)); } let items = self.select_items(&u, &exprs, env, "Aggregate")?; - let agg = self.build_aggregate(Origin::UNKNOWN, &u, group, &items, false)?; + let agg = self.build_aggregate(Origin::UNKNOWN, &u, group, &items, false, None)?; let col = |i: usize| col_expr(&agg.out[i].col, Origin::UNKNOWN); let (lc, rc) = (col(n), col(n + 1)); let (predicate, count) = match op { diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index 7483abc98..092b6eeed 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -288,6 +288,26 @@ impl Display for NodeLine<'_> { PlanNode::ReplicateRows { count, .. } => { write!(f, "Generate replicaterows({})", Text(count)) } + // Spark's `Expand [[a, null, 1], [a, b, 0]], [a, b, spark_grouping_id]`. + PlanNode::Expand { + projections, + output, + .. + } => { + f.write_str("Expand [")?; + for (i, p) in projections.iter().enumerate() { + write!(f, "{}[", sep(i))?; + for (j, e) in p.iter().enumerate() { + write!(f, "{}{}", sep(j), Text(e))?; + } + f.write_str("]")?; + } + f.write_str("], [")?; + for (i, c) in output.iter().enumerate() { + write!(f, "{}{}#{}", sep(i), c.name, c.id.0)?; + } + f.write_str("]") + } PlanNode::WithCte { ctes, .. } => { f.write_str("WithCte [")?; for (i, c) in ctes.iter().enumerate() { diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index 2957c5fc4..976a4f499 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -791,6 +791,17 @@ impl<'k, 'a> Encoder<'k, 'a> { self.u8(21); self.expr(count); } + PlanNode::Expand { + projections, + output, + .. + } => { + self.u8(22); + self.len(projections.len()); + projections.iter().for_each(|p| self.exprs(p)); + self.len(output.len()); + output.iter().for_each(|c| self.bound(c)); + } } } diff --git a/crates/krishiv-plan/src/lib.rs b/crates/krishiv-plan/src/lib.rs index 35ddaf2e1..79607784f 100644 --- a/crates/krishiv-plan/src/lib.rs +++ b/crates/krishiv-plan/src/lib.rs @@ -67,7 +67,7 @@ //! | `PlanNode::Project` | U, R | Q2, Q4, Q6-Q9, Q11, Q14-Q17, Q19-Q22 | `select` | //! | `PlanNode::WithColumns` | U | - | `withColumn` | //! | `PlanNode::Filter` | U, R | Q1-Q12, Q14-Q22 | `filter`, `where` | -//! | `PlanNode::Aggregate` | U, R | Q1, Q3-Q5, Q7-Q13, Q15, Q16, Q18, Q21, Q22 | `groupBy.agg` | +//! | `PlanNode::Aggregate` | U, R | Q1, Q3-Q5, Q7-Q13, Q15, Q16, Q18, Q21, Q22 | `groupBy.agg`, `rollup.agg`, `cube.agg` | //! | `PlanNode::Join` | U, R | Q2, Q3, Q5, Q7-Q21 | `join` | //! | `PlanNode::Sort` | U, R | Q1-Q5, Q7-Q13, Q15, Q16, Q18, Q20-Q22 | `orderBy`, `sort` | //! | `PlanNode::Limit` | U, R | Q2, Q3, Q10, Q18, Q21 | `limit` | @@ -81,6 +81,7 @@ //! | `PlanNode::Distinct` | U | - | - | //! | `PlanNode::SetOp` | U | - | - | //! | `PlanNode::ReplicateRows` | R | - | - | +//! | `PlanNode::Expand` | R | - | - | //! | `ExprNode::Column` | U | all | `col`, `df["a"]` | //! | `ExprNode::Literal` | U, R | all | `lit` | //! | `ExprNode::Star` | U | Q1, Q4, Q13, Q21, Q22 | `select("*")` | diff --git a/crates/krishiv-plan/src/plan.rs b/crates/krishiv-plan/src/plan.rs index 59df152d5..d0e791be4 100644 --- a/crates/krishiv-plan/src/plan.rs +++ b/crates/krishiv-plan/src/plan.rs @@ -243,6 +243,19 @@ pub enum PlanNode { /// A `BIGINT` expression over the input: the copies of each row. count: Arc, }, + /// Each input row once per projection (Spark's `Expand`, the first step + /// of `ROLLUP`, `CUBE` and `GROUPING SETS`): row `i` of the output for + /// an input row is the row evaluated by `projections[i]`. Resolved + /// plans only: the analyzer builds it, with one expression per output + /// column in every projection. + Expand { + /// Input. + input: Arc, + /// The projections, in output order; at least one. + projections: Vec>>, + /// The output attributes, aligned with each projection. + output: Vec, + }, } /// The operation of a [`PlanNode::SetOp`]. @@ -554,6 +567,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } + | PlanNode::Expand { input, .. } | PlanNode::ReplicateRows { input, .. } => f(input), PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { f(left); @@ -585,6 +599,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } + | PlanNode::Expand { input, .. } | PlanNode::ReplicateRows { input, .. } => f(input), PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { f(left); @@ -644,6 +659,7 @@ impl LogicalPlan { JoinCondition::None | JoinCondition::Using(_) => {} }, PlanNode::Sort { order, .. } => order.iter().for_each(|o| f(&o.expr)), + PlanNode::Expand { projections, .. } => projections.iter().flatten().for_each(f), } } @@ -694,6 +710,9 @@ impl LogicalPlan { JoinCondition::None | JoinCondition::Using(_) => {} }, PlanNode::Sort { order, .. } => order.iter_mut().for_each(|o| f(&mut o.expr)), + PlanNode::Expand { projections, .. } => { + projections.iter_mut().flatten().for_each(f); + } } } @@ -748,6 +767,7 @@ impl LogicalPlan { ); } PlanNode::OneRowRelation => return Some(Vec::new()), + PlanNode::Expand { output, .. } => return Some(output.clone()), PlanNode::Scan(s) => return Some(s.output.clone()), PlanNode::CteRef(r) => return Some(r.output.clone()), PlanNode::Shift { input, offset } => { diff --git a/crates/krishiv-plan/src/tests/analyzer.rs b/crates/krishiv-plan/src/tests/analyzer.rs index 1e280cebe..c05b3ba6b 100644 --- a/crates/krishiv-plan/src/tests/analyzer.rs +++ b/crates/krishiv-plan/src/tests/analyzer.rs @@ -1423,13 +1423,8 @@ fn unsupported_inputs_fail_explicitly() { paths: vec![krishiv_common::secret::OptionValue::from_user("path", "/d")], options: std::collections::BTreeMap::new(), })); - let rollup = LogicalPlan::new(PlanNode::Aggregate { - input: Arc::new(LogicalPlan::table(["t"])), - grouping: Grouping::Rollup(vec![Arc::new(c("a"))]), - aggregates: vec![Arc::new(c("a"))], - }); let sha = LogicalPlan::table(["t"]).project(vec![Expr::function("factorial", vec![c("b")])]); - for plan in [files, rollup, sha] { + for plan in [files, sha] { assert_eq!(class(&run(plan, &ctx)), "KRISHIV_UNSUPPORTED_FEATURE"); } // An untyped NULL is the void type; an interval is not supported yet. diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index 011e3a136..fcd401999 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -12,7 +12,7 @@ use krishiv_types::DataType; use super::dsl::{c, int}; use super::tpch::QUERIES; use crate::{ - Builtin, ColumnRef, Expr, ExprNode, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, + Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, }; @@ -41,6 +41,7 @@ pub(super) fn kinds(plan: &LogicalPlan, out: &mut BTreeSet<&'static str>) { PlanNode::Distinct { .. } => "PlanNode::Distinct", PlanNode::SetOp { .. } => "PlanNode::SetOp", PlanNode::ReplicateRows { .. } => "PlanNode::ReplicateRows", + PlanNode::Expand { .. } => "PlanNode::Expand", }); plan.for_each_expr(&mut |e| expr_kinds(e, out)); plan.for_each_child(&mut |c| kinds(c, out)); @@ -153,6 +154,13 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { columns: vec![("b".to_owned(), Arc::new(c("a")))], }); let sum = Expr::builtin(Builtin::Add, vec![c("a"), int(1)]); + let grouped = |grouping| { + LogicalPlan::new(PlanNode::Aggregate { + input: Arc::new(df(1)), + grouping, + aggregates: vec![Arc::new(c("a"))], + }) + }; vec![ ("read.table", df(1)), ("read.parquet", parquet), @@ -162,6 +170,8 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { ("filter", df(1).filter(c("a"))), ("where", df(1).filter(c("a"))), ("groupBy.agg", df(1).aggregate(vec![c("a")], vec![c("a")])), + ("rollup.agg", grouped(Grouping::Rollup(vec![Arc::new(c("a"))]))), + ("cube.agg", grouped(Grouping::Cube(vec![Arc::new(c("a"))]))), ( "join", df(1).join(df(2), JoinKind::Inner, JoinCondition::None), diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index 715382a29..ec013bb5d 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -288,7 +288,7 @@ impl Gen<'_> { fn plan(&mut self, depth: u8) -> LogicalPlan { let origin = self.origin(); - let kinds = if depth == 0 { 4 } else { 21 }; + let kinds = if depth == 0 { 4 } else { 22 }; let d = depth.saturating_sub(1); let node = match self.pick(kinds) { 0 => PlanNode::Read(self.source()), @@ -424,6 +424,11 @@ impl Gen<'_> { input: self.input(d), count: Arc::new(self.expr(1)), }, + 21 => PlanNode::Expand { + input: self.input(d), + projections: self.many(1, 2, |g| g.exprs()), + output: self.many(0, 2, Gen::bound), + }, _ => PlanNode::WithCte { input: self.input(d), ctes: self.many(1, 2, |g| Cte { @@ -813,6 +818,22 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { } | PlanNode::Limit { input, limit: e } | PlanNode::ReplicateRows { input, count: e } => m_expr(e, w) || m_plan(input, w), + PlanNode::Expand { + input, + projections, + output, + } => { + site!(w, projections.push(Vec::new())); + site!(w, output.push(output.first().cloned().unwrap_or(BoundColumn { + id: AttrId(0), + name: "x".to_owned(), + data_type: DataType::Integer, + nullable: false, + }))); + projections.iter_mut().any(|p| m_exprs(p, w)) + || output.iter_mut().any(|c| m_bound(c, w)) + || m_plan(input, w) + } PlanNode::SetOp { left, right, diff --git a/crates/krishiv-plan/src/walk.rs b/crates/krishiv-plan/src/walk.rs index f79c41542..ee0d9ed3b 100644 --- a/crates/krishiv-plan/src/walk.rs +++ b/crates/krishiv-plan/src/walk.rs @@ -155,6 +155,7 @@ pub(crate) fn all_resolved(root: Node<'_>) -> bool { | PlanNode::SubqueryAlias { .. } | PlanNode::WithCte { .. } | PlanNode::JoinHint { .. } + | PlanNode::Expand { .. } | PlanNode::ReplicateRows { .. } => true, }, Node::Expr(e) => match &e.node { @@ -257,6 +258,12 @@ impl Teardown { self.plans.push(input); self.exprs.push(e); } + PlanNode::Expand { + input, projections, .. + } => { + self.plans.push(input); + self.exprs.extend(projections.into_iter().flatten()); + } PlanNode::Aggregate { input, grouping, diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index de0b9a94b..5c499c5bd 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -255,9 +255,73 @@ impl DataFrame { /// ``` #[must_use] pub fn group_by(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::GroupBy) + } + + /// Groups rows by `keys` and by each prefix of them, down to the grand + /// total, for [`GroupedData::agg`] (PySpark's `rollup`): a key a row is + /// rolled up over reads `NULL`. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s.table("t").rollup([col("a"), col("b")]).agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn rollup(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::Rollup) + } + + /// Groups rows by every subset of `keys`, for [`GroupedData::agg`] + /// (PySpark's `cube`): a key a row is rolled up over reads `NULL`. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s.table("t").cube([col("a"), col("b")]).agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn cube(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::Cube) + } + + /// Groups rows by each of `sets` (grouping expressions among `keys`), + /// for [`GroupedData::agg`] (PySpark's `groupingSets`): the output has + /// the columns of `keys`, `NULL` where a set does not group by them. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s + /// .table("t") + /// .grouping_sets([vec![col("a")], vec![col("b")], vec![]], [col("a"), col("b")]) + /// .agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn grouping_sets( + &self, + sets: impl IntoIterator>, + keys: impl IntoIterator>, + ) -> GroupedData { + let sets = sets + .into_iter() + .map(|set| set.into_iter().map(|c| c.expr).collect()) + .collect(); + self.grouped(keys, Sets::Explicit(sets)) + } + + fn grouped( + &self, + keys: impl IntoIterator>, + sets: Sets, + ) -> GroupedData { GroupedData { df: self.clone(), keys: keys.into_iter().map(|k| k.into().expr).collect(), + sets, } } @@ -749,11 +813,26 @@ impl DataFrame { } } -/// Rows grouped by keys ([`DataFrame::group_by`]). +/// Which groups of its keys a [`GroupedData`] aggregates. +#[derive(Debug, Clone)] +enum Sets { + /// All the keys together. + GroupBy, + /// Each prefix of the keys. + Rollup, + /// Each subset of the keys. + Cube, + /// These sets. + Explicit(Vec>), +} + +/// Rows grouped by keys ([`DataFrame::group_by`], [`DataFrame::rollup`], +/// [`DataFrame::cube`], [`DataFrame::grouping_sets`]). #[derive(Debug, Clone)] pub struct GroupedData { df: DataFrame, keys: Vec, + sets: Sets, } impl GroupedData { @@ -771,9 +850,21 @@ impl GroupedData { let keys: Vec> = self.keys.iter().cloned().map(Arc::new).collect(); let mut output = keys.clone(); output.extend(aggregates.into_iter().map(|c| Arc::new(c.expr))); + let grouping = match &self.sets { + Sets::GroupBy => Grouping::GroupBy(keys), + Sets::Rollup => Grouping::Rollup(keys), + Sets::Cube => Grouping::Cube(keys), + Sets::Explicit(sets) => Grouping::GroupingSets { + exprs: keys, + sets: sets + .iter() + .map(|set| set.iter().cloned().map(Arc::new).collect()) + .collect(), + }, + }; self.df.derive(PlanNode::Aggregate { input: Arc::clone(&self.df.plan), - grouping: Grouping::GroupBy(keys), + grouping, aggregates: output, }) } diff --git a/crates/krishiv-session/tests/scratch_rollup.rs b/crates/krishiv-session/tests/scratch_rollup.rs new file mode 100644 index 000000000..f12c60530 --- /dev/null +++ b/crates/krishiv-session/tests/scratch_rollup.rs @@ -0,0 +1,84 @@ +//! scratch +#![cfg(test)] +use std::sync::Arc; + +use arrow_array::{Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType as T, Field, Schema}; +use futures::executor::block_on; +use krishiv_session::Session; + +#[test] +fn scratch() { + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", T::Int32, false), + Field::new("b", T::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2, 3, 3])), + Arc::new(Int32Array::from(vec![1, 2, 1, 2, 1, 2])), + ], + ) + .unwrap(); + s.register_batches("td", schema, vec![batch]).unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("course", T::Utf8, false), + Field::new("year", T::Int32, false), + Field::new("earnings", T::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["dotNET", "Java", "dotNET", "dotNET", "Java"])), + Arc::new(Int32Array::from(vec![2012, 2012, 2012, 2013, 2013])), + Arc::new(Int32Array::from(vec![10000, 20000, 5000, 48000, 30000])), + ], + ) + .unwrap(); + s.register_batches("cs", schema, vec![batch]).unwrap(); + for q in std::fs::read_to_string("/home/krishiv/krishiv-data/tmp/s-rollup/q.sql") + .unwrap() + .split(";\n") + .filter(|q| !q.trim().is_empty()) + { + match block_on(s.sql(q)) { + Err(e) => println!("ERR {q}\n {e}"), + Ok(df) => match block_on(df.collect()) { + Err(e) => println!("RUNERR {q}\n {e}"), + Ok(b) => { + println!("OK {q}"); + for b in &b { + println!("{}", arrow_cast_fmt(b)); + } + } + }, + } + } +} + +fn arrow_cast_fmt(b: &RecordBatch) -> String { + let mut out = format!("{:?}\n", b.schema().fields().iter().map(|f| format!("{}:{}{}", f.name(), f.data_type(), if f.is_nullable() {"?"} else {""})).collect::>()); + for r in 0..b.num_rows() { + let cells: Vec = b + .columns() + .iter() + .map(|c| { + if c.is_null(r) { + "NULL".to_owned() + } else { + let d = arrow_array::cast::AsArray::as_primitive_opt::(c) + .map(|a| a.value(r).to_string()) + .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) + .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) + .or_else(|| arrow_array::cast::AsArray::as_string_opt::(c).map(|a| a.value(r).to_owned())); + d.unwrap_or_else(|| "?".to_owned()) + } + }) + .collect(); + out.push_str(&cells.join("\t")); + out.push('\n'); + } + out +} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index 6a4e5462c..8daa1db39 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -27,7 +27,8 @@ use krishiv_plan::{ use krishiv_types::DecimalType; use crate::ast::{ - Alias, BinaryOp, Expr, ExprKind, FromClause, GroupItemKind, JoinConstraint, JoinKind, Literal, + Alias, BinaryOp, Expr, ExprKind, FromClause, GroupBy, GroupItem, GroupItemKind, GroupWith, + GroupingSet, JoinConstraint, JoinKind, Literal, NullTreatment, NumKind, Query, Relation, RelationExt, Select, SetExpr, SetExprKind, SetOperator, SetQuantifier, SortItem, Span, TableFactor, TableFactorKind, Truth, UnaryOp, }; @@ -36,6 +37,9 @@ use crate::{ParseError, ParserOptions, parse_query}; type R = Result; +/// The most grouping sets one `GROUP BY` expands to. +const MAX_GROUPING_SETS: usize = 1 << 16; + /// Parses and lowers one query; spans index `source`. /// /// # Errors @@ -389,22 +393,14 @@ impl Lowerer { s.span, )); } - let mut keys = Vec::new(); - if let Some(g) = &s.group_by { - if g.with.is_some() { - return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", g.span)); - } - for item in &g.items { - match &item.kind { - GroupItemKind::Expr(e) => keys.push(arc(self.expr(e)?)), - _ => return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", item.span)), - } - } - } + let grouping = match &s.group_by { + Some(g) => self.grouping(g)?, + None => Grouping::GroupBy(Vec::new()), + }; plan = self.node( PlanNode::Aggregate { input: arc(plan), - grouping: Grouping::GroupBy(keys), + grouping, aggregates: items, }, s.span, @@ -437,6 +433,129 @@ impl Lowerer { Ok(plan) } + /// `GROUP BY`, as Spark's `AstBuilder` shapes it: plain keys, `WITH + /// ROLLUP` / `WITH CUBE` over them, `GROUPING SETS (..)` after them (the + /// keys only fix the grouping expressions), or items that are + /// `ROLLUP`, `CUBE` and `GROUPING SETS` among keys, whose grouping sets + /// are the cross product of the items' (a key is a set of itself). + #[inline(never)] + fn grouping(&mut self, g: &GroupBy) -> R { + let mut plain = Vec::with_capacity(g.items.len()); + for item in &g.items { + if let GroupItemKind::Expr(e) = &item.kind { + plain.push(arc(self.expr(e)?)); + } + } + let all_plain = plain.len() == g.items.len(); + match (&g.with, all_plain) { + (None, true) => return Ok(Grouping::GroupBy(plain)), + (Some(GroupWith::Rollup), true) => return Ok(Grouping::Rollup(plain)), + (Some(GroupWith::Cube), true) => return Ok(Grouping::Cube(plain)), + (Some(GroupWith::GroupingSets(sets)), true) => { + let sets = sets + .iter() + .map(|s| self.set(s)) + .collect::>>()?; + return Ok(Grouping::GroupingSets { exprs: plain, sets }); + } + (Some(_), false) => { + return Err(unsupported( + "ROLLUP, CUBE and GROUPING SETS items with WITH ROLLUP, WITH CUBE or GROUPING SETS", + g.span, + )); + } + (None, false) => {} + } + let mut exprs = Vec::new(); + let mut product: Vec>> = vec![Vec::new()]; + for item in &g.items { + let sets = self.item_sets(item, &mut exprs)?; + if product.len().saturating_mul(sets.len()) > MAX_GROUPING_SETS { + return Err(unsupported( + &format!("more than {MAX_GROUPING_SETS} grouping sets"), + g.span, + )); + } + product = product + .iter() + .flat_map(|head| { + sets.iter().map(move |tail| { + head.iter().chain(tail).cloned().collect::>() + }) + }) + .collect(); + } + Ok(Grouping::GroupingSets { + exprs, + sets: product, + }) + } + + /// The grouping sets one `GROUP BY` item stands for; the expressions it + /// names are appended to `written` in order. + fn item_sets( + &mut self, + item: &GroupItem, + written: &mut Vec>, + ) -> R>>> { + Ok(match &item.kind { + GroupItemKind::Expr(e) => { + let e = arc(self.expr(e)?); + written.push(Arc::clone(&e)); + vec![vec![e]] + } + GroupItemKind::Set(s) => { + let set = self.set(s)?; + written.extend(set.iter().cloned()); + vec![set] + } + GroupItemKind::GroupingSets(items) => { + let mut sets = Vec::new(); + for item in items { + sets.extend(self.item_sets(item, written)?); + } + sets + } + GroupItemKind::Rollup(elements) | GroupItemKind::Cube(elements) => { + let rollup = matches!(item.kind, GroupItemKind::Rollup(_)); + let mut parts = Vec::with_capacity(elements.len()); + for element in elements { + let part = self.set(element)?; + written.extend(part.iter().cloned()); + parts.push(part); + } + // `ROLLUP (a, b)`: (a, b), (a), (); `CUBE (a, b)`: every + // subset of the elements. + let n = parts.len(); + if !rollup && n >= usize::BITS as usize - 1 + || !rollup && (1_usize << n) > MAX_GROUPING_SETS + { + return Err(unsupported( + &format!("more than {MAX_GROUPING_SETS} grouping sets"), + item.span, + )); + } + let chosen: Vec> = if rollup { + (0..=n).rev().map(|k| (0..k).collect()).collect() + } else { + (0..1_usize << n) + .rev() + .map(|m| (0..n).filter(|i| m >> (n - 1 - i) & 1 == 1).collect()) + .collect() + }; + chosen + .into_iter() + .map(|c| c.into_iter().flat_map(|i| parts[i].iter().cloned()).collect()) + .collect() + } + }) + } + + /// The expressions of one grouping set. + fn set(&mut self, s: &GroupingSet) -> R>> { + s.exprs.iter().map(|e| Ok(arc(self.expr(e)?))).collect() + } + fn from(&mut self, f: &FromClause) -> R { if f.pivot.is_some() || f.unpivot.is_some() { return Err(unsupported("PIVOT and UNPIVOT", f.span)); From 7bb422483cacecf444a1a884e5a4e967cae8dfc6 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:59:20 +0200 Subject: [PATCH 2/8] feat(sql): ROLLUP, CUBE, GROUPING SETS, grouping() and grouping_id() via Expand Spark's ResolveGroupingAnalytics shape: an Expand plan node (plan data, streamed one copy per projection) feeds an ordinary aggregate. Adds the DataFrame rollup/cube/grouping_sets, Spark golden tests, optimizer cases and commutative semantic keys. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/lib.rs | 9 +- crates/krishiv-exec/src/operators/expand.rs | 13 +- crates/krishiv-exec/src/operators/tests.rs | 124 ++++++ .../krishiv-exec/src/physical/instantiate.rs | 4 +- crates/krishiv-exec/src/physical/mod.rs | 4 +- crates/krishiv-exec/src/planner/build.rs | 5 +- crates/krishiv-optimizer/src/prune.rs | 5 +- crates/krishiv-optimizer/tests/it/rules.rs | 39 ++ .../it__rules__expand_cube_pruned.snap | 7 + .../it__rules__expand_duplicate_sets.snap | 7 + .../it__rules__expand_grouping_sets.snap | 8 + .../snapshots/it__rules__expand_rollup.snap | 7 + .../tests/it/support/eval.rs | 24 ++ crates/krishiv-plan/src/analyzer/analytics.rs | 127 ++++-- crates/krishiv-plan/src/analyzer/exprs.rs | 38 +- crates/krishiv-plan/src/analyzer/ops.rs | 12 +- crates/krishiv-plan/src/tests/analyzer.rs | 4 +- crates/krishiv-plan/src/tests/inventory.rs | 9 +- crates/krishiv-plan/src/tests/props.rs | 17 +- .../tests/grouping_analytics.rs | 326 +++++++++++++++ .../tests/grouping_analytics_api.rs | 381 ++++++++++++++++++ .../krishiv-session/tests/scratch_rollup.rs | 84 ---- crates/krishiv-sql/src/lower.rs | 22 +- ...cts_lower_or_are_named_as_unsupported.snap | 4 +- 24 files changed, 1126 insertions(+), 154 deletions(-) create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap create mode 100644 crates/krishiv-session/tests/grouping_analytics.rs create mode 100644 crates/krishiv-session/tests/grouping_analytics_api.rs delete mode 100644 crates/krishiv-session/tests/scratch_rollup.rs diff --git a/crates/krishiv-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index b9081da93..d189d9e1c 100644 --- a/crates/krishiv-exec/src/lib.rs +++ b/crates/krishiv-exec/src/lib.rs @@ -113,10 +113,11 @@ pub use operator::{ FlushOutput, GlobalState, LocalState, Operator, OperatorOutput, Sink, SinkStatus, Source, }; pub use operators::{ - DEFAULT_MAX_MERGE_FAN_IN, Expand, NullOf, 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}; diff --git a/crates/krishiv-exec/src/operators/expand.rs b/crates/krishiv-exec/src/operators/expand.rs index 5f0a5c233..aedb7a718 100644 --- a/crates/krishiv-exec/src/operators/expand.rs +++ b/crates/krishiv-exec/src/operators/expand.rs @@ -39,14 +39,11 @@ impl Expand { 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() - }); + && 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 {:?}", diff --git a/crates/krishiv-exec/src/operators/tests.rs b/crates/krishiv-exec/src/operators/tests.rs index 2f344affe..d628fb75a 100644 --- a/crates/krishiv-exec/src/operators/tests.rs +++ b/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 { + 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> { + b.data() + .column(1) + .as_string::() + .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, 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); +} diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index 728028a00..58d5971fe 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -30,8 +30,8 @@ use crate::join::{ }; use crate::operator::{Operator, Sink, Source}; use crate::operators::{ - Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, 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; diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index 2ca8958d9..9d4833bf2 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -50,8 +50,8 @@ pub use data::{PlanData, assert_plan_data}; pub(crate) use instantiate::node_ids; pub use instantiate::{SpillScopes, WorkerContext, instantiate}; pub(crate) use instantiate::{ - Unbuilt, Values, aggregate_node, aggregate_spec, filter_op, hash_join, nested_loop_join, - expand_op, project_op, replicate_op, scan_filter, sort_op, top_n_op, + Unbuilt, Values, aggregate_node, aggregate_spec, expand_op, filter_op, hash_join, + nested_loop_join, project_op, replicate_op, scan_filter, sort_op, top_n_op, }; use crate::aggregate::{AggregateExpr, AggregateOutput}; diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index e9fcc721e..a004d4315 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -775,7 +775,10 @@ impl<'a> Planner<'a> { prepared, op.schema(), layout, - ("Expand", format!("{} projections [{text}]", projections.len())), + ( + "Expand", + format!("{} projections [{text}]", projections.len()), + ), |input, layout, subqueries| PhysNode::Expand { input, layout, diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index 7fb5f23ab..332c99b11 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -226,7 +226,10 @@ impl Rewrite<'_> { let node = PlanNode::Expand { input: Arc::clone(input), projections: projections.iter().map(pick).collect(), - output: kept.iter().filter_map(|&j| output.get(j).cloned()).collect(), + output: kept + .iter() + .filter_map(|&j| output.get(j).cloned()) + .collect(), }; replace(&rebuilt, node) } diff --git a/crates/krishiv-optimizer/tests/it/rules.rs b/crates/krishiv-optimizer/tests/it/rules.rs index b69d424bf..f797e7a46 100644 --- a/crates/krishiv-optimizer/tests/it/rules.rs +++ b/crates/krishiv-optimizer/tests/it/rules.rs @@ -398,6 +398,23 @@ const CASES: &[(&str, &str)] = &[ "prune_exists", "SELECT a FROM t WHERE EXISTS (SELECT * FROM u WHERE u.a = t.a)", ), + // Grouping sets: the expansion keeps what the aggregate reads. + ( + "expand_rollup", + "SELECT a, sum(c) FROM t GROUP BY ROLLUP(a, b)", + ), + ( + "expand_cube_pruned", + "SELECT count(*) FROM t GROUP BY CUBE(a, b)", + ), + ( + "expand_grouping_sets", + "SELECT a, grouping(a), grouping_id(a, b), sum(c) FROM t WHERE c > 1 GROUP BY GROUPING SETS ((a), (b), ())", + ), + ( + "expand_duplicate_sets", + "SELECT a, sum(c) FROM t GROUP BY GROUPING SETS ((a), (a))", + ), ]; #[test] @@ -443,6 +460,28 @@ fn rules() { } } +#[test] +fn an_expansion_has_a_copy_of_its_rows_per_projection() { + let ctx = support::ctx(); + let rows = |sql: &str| { + let plan = support::resolve(sql, &ctx); + let optimized = optimize(&plan).unwrap(); + let mut node = Arc::clone(optimized.root()); + while !matches!(node.node, PlanNode::Expand { .. }) { + let mut next = None; + node.for_each_child(&mut |c| next = Some(Arc::clone(c))); + node = next.expect("an expansion below"); + } + krishiv_optimizer::estimated_rows(&node, &krishiv_optimizer::NoStatistics) + }; + let one_set = rows("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))"); + let four_sets = rows("SELECT count(*) FROM t GROUP BY CUBE(a, b)"); + assert!( + (four_sets / one_set - 2.0).abs() < 1e-9, + "{four_sets} vs {one_set}" + ); +} + #[test] fn top_n_is_detected() { let ctx = support::ctx(); diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap new file mode 100644 index 000000000..24e2e4dc7 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap @@ -0,0 +1,7 @@ +--- +source: crates/krishiv-optimizer/tests/it/rules.rs +expression: "SELECT count(*) FROM t GROUP BY CUBE(a, b)" +--- +Aggregate [a#5, b#6, spark_grouping_id#3], [count(1) AS count(1)#4] ++- Expand [[a#0, b#1, 0L], [a#0, CAST(NULL AS STRING), 1L], [CAST(NULL AS INT), b#1, 2L], [CAST(NULL AS INT), CAST(NULL AS STRING), 3L]], [a#5, b#6, spark_grouping_id#3] + +- Scan spark_catalog.default.t [a#0, b#1] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap new file mode 100644 index 000000000..d5016c878 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap @@ -0,0 +1,7 @@ +--- +source: crates/krishiv-optimizer/tests/it/rules.rs +expression: "SELECT a, sum(c) FROM t GROUP BY GROUPING SETS ((a), (a))" +--- +Aggregate [a#5, spark_grouping_id#3, _gen_grouping_pos#6], [a#5, sum(c#2) AS sum(c)#4] ++- Expand [[c#2, a#0, 0L, 0], [c#2, a#0, 0L, 1]], [c#2, a#5, spark_grouping_id#3, _gen_grouping_pos#6] + +- Scan spark_catalog.default.t [a#0, c#2] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap new file mode 100644 index 000000000..b3333bb86 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap @@ -0,0 +1,8 @@ +--- +source: crates/krishiv-optimizer/tests/it/rules.rs +expression: "SELECT a, grouping(a), grouping_id(a, b), sum(c) FROM t WHERE c > 1 GROUP BY GROUPING SETS ((a), (b), ())" +--- +Aggregate [a#7, b#8, spark_grouping_id#3], [a#7, CAST(((grouping__id#3 div 2L) % 2L) AS TINYINT) AS grouping(a)#4, grouping__id#3 AS grouping_id(a, b)#5, sum(c#2) AS sum(c)#6] ++- Expand [[c#2, a#0, CAST(NULL AS STRING), 1L], [c#2, CAST(NULL AS INT), b#1, 2L], [c#2, CAST(NULL AS INT), CAST(NULL AS STRING), 3L]], [c#2, a#7, b#8, spark_grouping_id#3] + +- Filter (c#2 > 1L) + +- Scan spark_catalog.default.t [a#0, b#1, c#2] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap new file mode 100644 index 000000000..9c75b03d3 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap @@ -0,0 +1,7 @@ +--- +source: crates/krishiv-optimizer/tests/it/rules.rs +expression: "SELECT a, sum(c) FROM t GROUP BY ROLLUP(a, b)" +--- +Aggregate [a#5, b#6, spark_grouping_id#3], [a#5, sum(c#2) AS sum(c)#4] ++- Expand [[c#2, a#0, b#1, 0L], [c#2, a#0, CAST(NULL AS STRING), 1L], [c#2, CAST(NULL AS INT), CAST(NULL AS STRING), 3L]], [c#2, a#5, b#6, spark_grouping_id#3] + +- Scan spark_catalog.default.t [a#0, b#1, c#2] diff --git a/crates/krishiv-optimizer/tests/it/support/eval.rs b/crates/krishiv-optimizer/tests/it/support/eval.rs index bc73a19b3..e0d75ce2b 100644 --- a/crates/krishiv-optimizer/tests/it/support/eval.rs +++ b/crates/krishiv-optimizer/tests/it/support/eval.rs @@ -177,6 +177,28 @@ impl Interp<'_> { rows, } } + PlanNode::Expand { + input, + projections, + output, + } => { + let rel = self.plan(input, outer); + let ix = index(&rel.cols); + let rows = rel + .rows + .iter() + .flat_map(|r| { + projections + .iter() + .map(|row| row.iter().map(|e| self.eval(e, &ix, r, outer)).collect()) + .collect::>() + }) + .collect(); + Rel { + cols: output.iter().map(|c| c.id).collect(), + rows, + } + } PlanNode::Filter { input, predicate } => { let mut preds = Vec::new(); conjuncts(predicate, &mut preds); @@ -746,6 +768,8 @@ fn apply(name: &str, a: &[V]) -> V { "*" => num(|x, y| Some(x * y / SCALE)), "/" => num(|x, y| (y != 0).then(|| x * SCALE / y)), "%" => num(|x, y| (y != 0).then(|| x % y)), + // Integral division: the quotient truncated to a whole number. + "div" => num(|x, y| (y != 0).then(|| x / y * SCALE)), "negative" => match &a[0] { V::Num(x) => V::Num(-x), _ => V::Null, diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs index c6d34c79e..8173c6063 100644 --- a/crates/krishiv-plan/src/analyzer/analytics.rs +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -47,6 +47,10 @@ use crate::{ /// `BIGINT` with one bit each. const MAX_GROUPING: usize = 63; +/// The most expressions of a `CUBE`: 2^16 grouping sets (the SQL front end +/// bounds `CUBE (..)` items the same way). +const MAX_CUBE_EXPRS: usize = 16; + /// `grouping` or `grouping_id`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum GroupingFn { @@ -104,6 +108,8 @@ pub(super) struct Analytics { gid: Option, /// The expansion; `None` for a grand total. pub(super) expand: Option, + /// Whether the expansion numbers its sets (some two are equal). + numbered: bool, } impl Analytics { @@ -117,6 +123,17 @@ impl Analytics { .collect() } + /// `HAVING` and `ORDER BY` over numbered sets: Spark reads the last + /// grouping key as the grouping id there, and that key is the set's + /// number, so a `grouping` or `grouping_id` call is + /// `UNSUPPORTED_GROUPING_EXPRESSION`. + pub(super) fn check_late_call(&self, e: &Arc) -> Result<(), AnalysisError> { + if !self.numbered { + return Ok(()); + } + reject_grouping_calls(std::slice::from_ref(e)) + } + /// The position of `e` among the grouping expressions: a copy of one, or /// an expression equal to one. fn position(&self, e: &Expr, memo: &mut HashMap) -> Option { @@ -130,6 +147,38 @@ impl Analytics { } } +/// `UNSUPPORTED_GROUPING_EXPRESSION` if any of `exprs` calls `grouping` or +/// `grouping_id`. Spark has no source position for it. +pub(super) fn reject_grouping_calls(exprs: &[Arc]) -> Result<(), AnalysisError> { + let mut seen = std::collections::HashSet::new(); + let mut stack: Vec> = exprs.to_vec(); + while let Some(x) = stack.pop() { + if !seen.insert(Arc::as_ptr(&x) as usize) { + continue; + } + if matches!(&x.node, ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some()) { + return Err(AnalysisError::new( + UnsupportedGroupingExpression {}, + Origin::UNKNOWN, + )); + } + x.for_each_child(&mut |c| stack.push(Arc::clone(c))); + } + Ok(()) +} + +/// A resolved call `e`, stated not nullable. +fn not_null(e: &Arc) -> Arc { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) => { + let mut f = f.clone(); + f.nullable = false; + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function(f))).with_origin(e.origin)) + } + _ => Arc::clone(e), + } +} + fn literal(l: Literal) -> Arc { Arc::new(Expr::new(ExprNode::Literal(l))) } @@ -313,17 +362,25 @@ impl Resolver<'_> { aggregates: &[Arc], env: &Arc, ) -> Result { + // A cube has a set per subset of its expressions: bound it before + // building them (`WITH CUBE` and `DataFrame.cube` arrive here). + if let Grouping::Cube(v) = grouping + && v.len() > MAX_CUBE_EXPRS + { + return Err(AnalysisError::new( + GroupingSizeLimitExceeded { + max_size: (1_u64 << MAX_CUBE_EXPRS).to_string(), + }, + p.origin, + )); + } let c = self.child(input, env)?; let (flat, sets) = shape(grouping); self.check_star_ordinals(&flat, aggregates)?; // A single empty set over nothing is a grand total. let grand = flat.is_empty() && sets.len() == 1; - let gid_col = (!grand).then(|| { - self.attr( - "spark_grouping_id", - &arg_type(DataType::Long, false), - ) - }); + let gid_col = + (!grand).then(|| self.attr("spark_grouping_id", &arg_type(DataType::Long, false))); let gid = gid_col.as_ref().map(|col| Attr { col: BoundColumn { name: "grouping__id".to_owned(), @@ -375,9 +432,17 @@ impl Resolver<'_> { exprs: Vec::new(), gid: None, expand: None, + numbered: false, }; let items = self.rewrite_items(items, &an, env)?; - return self.build_aggregate(p.origin, &c, Vec::new(), &items, true, Some(Arc::new(an))); + return self.build_aggregate( + p.origin, + &c, + Vec::new(), + &items, + true, + Some(Arc::new(an)), + ); } let Some(gid_col) = gid_col else { return Err(unsupported("grouping sets without an id", p.origin)); @@ -385,7 +450,11 @@ impl Resolver<'_> { let numbered = (0..sets.len()).any(|i| sets[..i].contains(&sets[i])); // The expansion's input: the child's columns, and the grouping // expressions that are not one of them. - let mut passed: Vec> = c.out.iter().map(|a| col_expr(&a.col, Origin::UNKNOWN)).collect(); + let mut passed: Vec> = c + .out + .iter() + .map(|a| col_expr(&a.col, Origin::UNKNOWN)) + .collect(); let mut passed_out: Vec = c.out.to_vec(); let mut value: Vec = Vec::with_capacity(n); let mut copies: Vec = Vec::with_capacity(n); @@ -406,10 +475,12 @@ impl Resolver<'_> { Some(a) => a.col.clone(), None => { let v = self.attr(&name, &b.ty); - passed.push(Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { - expr: Arc::clone(&b.expr), - output: v.clone(), - })))); + passed.push(Arc::new(Expr::new(ExprNode::Resolved( + ResolvedExpr::Alias { + expr: Arc::clone(&b.expr), + output: v.clone(), + }, + )))); passed_out.push(Attr { col: v.clone(), qual: Arc::from(Vec::new()), @@ -425,7 +496,9 @@ impl Resolver<'_> { key: keys_of[j], attr: Attr { col: copy, - qual: column.as_ref().map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), + qual: column + .as_ref() + .map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), source: column.and_then(|a| a.source), }, }); @@ -437,7 +510,8 @@ impl Resolver<'_> { } else { c.clone() }; - let pos_col = numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); + let pos_col = + numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); let mut output: Vec = passed_out.iter().map(|a| a.col.clone()).collect(); output.extend(copies.iter().map(|g| g.attr.col.clone())); output.push(gid_col.clone()); @@ -459,9 +533,9 @@ impl Resolver<'_> { } row.push(literal(Literal::Long(mask))); if pos_col.is_some() { - row.push(literal(Literal::Integer( - i32::try_from(si).map_err(|_| unsupported("that many grouping sets", p.origin))?, - ))); + row.push(literal(Literal::Integer(i32::try_from(si).map_err( + |_| unsupported("that many grouping sets", p.origin), + )?))); } projections.push(row); } @@ -492,6 +566,7 @@ impl Resolver<'_> { let an = Analytics { exprs: copies, gid, + numbered, expand: Some(ExpandPlan { plan: Arc::clone(&expand.plan), ids: expand.ids, @@ -619,8 +694,7 @@ impl Resolver<'_> { Some(g) => col_expr(&g.col, Origin::UNKNOWN), None => literal(Literal::Long(0)), }; - let positions: Vec> = - f.args.iter().map(|a| an.position(a, memo)).collect(); + let positions: Vec> = f.args.iter().map(|a| an.position(a, memo)).collect(); match GroupingFn::from_name(&f.name) { Some(GroupingFn::GroupingId) => { let exact = positions.len() == n @@ -647,7 +721,7 @@ impl Resolver<'_> { Ok(id) } _ => { - let (Some(arg), Some(Some(j))) = (f.args.first(), positions.first()) else { + let (Some(_), Some(Some(j))) = (f.args.first(), positions.first()) else { return Err(AnalysisError::new( GroupingColumnMismatch { grouping: f.args.first().map(|a| with_id(a)).unwrap_or_default(), @@ -661,18 +735,20 @@ impl Resolver<'_> { Origin::UNKNOWN, )); }; - let _ = arg; // Bit `n - 1 - j` of the id, as a `TINYINT`. let shift = n - 1 - j; + // The divisors are literals that are not 0, so the id's + // bit is never null: `grouping` is not nullable (Spark). let high = if shift == 0 { id } else { - self.call( + let high = self.call( "div", vec![id, literal(Literal::Long(1_i64 << shift))], env, Origin::UNKNOWN, - )? + )?; + not_null(&high) }; let bit = self.call( "%", @@ -680,12 +756,15 @@ impl Resolver<'_> { env, Origin::UNKNOWN, )?; + let bit = not_null(&bit); let cast = Arc::new(Expr::new(ExprNode::Cast { expr: bit, to: DataType::Byte, mode: CastMode::Default, })); - Ok(self.bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + Ok(self + .bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))? + .expr) } } } diff --git a/crates/krishiv-plan/src/analyzer/exprs.rs b/crates/krishiv-plan/src/analyzer/exprs.rs index b16b69299..3bf207a15 100644 --- a/crates/krishiv-plan/src/analyzer/exprs.rs +++ b/crates/krishiv-plan/src/analyzer/exprs.rs @@ -272,6 +272,8 @@ impl Resolver<'_> { stack.push(Step::Enter(rewritten, in_agg)); continue; } + // In `ORDER BY` Spark rejects the call before it resolves + // its argument (`HAVING` reports the argument first). let inner = in_agg || self.is_aggregate(&x); stack.push(Step::Exit(Arc::clone(&x), in_agg)); let mut kids = Vec::new(); @@ -1525,6 +1527,29 @@ fn pretty_literal(v: &Literal) -> String { } } +/// A binary operator whose operands commute (Spark's `Add`, `Multiply`, +/// `And`, `Or`, `EqualTo` and `EqualNullSafe`). +fn is_commutative(node: &ExprNode) -> bool { + let ExprNode::Resolved(ResolvedExpr::Function(f)) = node else { + return false; + }; + let op = match &f.id { + FunctionId::Builtin(b) => Some(*b), + FunctionId::Registry { name, .. } => Builtin::from_name(name), + }; + matches!( + op, + Some( + Builtin::Add + | Builtin::Multiply + | Builtin::And + | Builtin::Or + | Builtin::Eq + | Builtin::EqNullSafe + ) + ) +} + /// A semantic key: equal for expressions that compute the same value, /// whatever their origins, column spellings or (top-level) alias. pub(super) fn semantic_key(e: &Expr, memo: &mut HashMap) -> u128 { @@ -1557,8 +1582,17 @@ pub(super) fn semantic_key(e: &Expr, memo: &mut HashMap) -> u128 { other => format!("u{other:?}"), } .into_bytes(); - for c in kids { - b.extend_from_slice(&memo.get(&addr(c)).copied().unwrap_or(0).to_le_bytes()); + let mut hashes: Vec = kids + .iter() + .map(|c| memo.get(&addr(*c)).copied().unwrap_or(0)) + .collect(); + // Spark canonicalizes the operands of a commutative operator into + // one order, so `a + b` and `b + a` are the same expression. + if hashes.len() == 2 && is_commutative(&x.node) { + hashes.sort_unstable(); + } + for h in hashes { + b.extend_from_slice(&h.to_le_bytes()); } memo.insert(addr(x), XxHash3_128::oneshot(&b)); } diff --git a/crates/krishiv-plan/src/analyzer/ops.rs b/crates/krishiv-plan/src/analyzer/ops.rs index 8238aec0e..aff4ece89 100644 --- a/crates/krishiv-plan/src/analyzer/ops.rs +++ b/crates/krishiv-plan/src/analyzer/ops.rs @@ -684,6 +684,7 @@ impl Resolver<'_> { } if let Some(an) = &info.analytics { for b in &mut bounds { + an.check_late_call(&b.expr)?; b.expr = self.resolve_grouping(&b.expr, an, env)?; } } @@ -960,7 +961,11 @@ impl Resolver<'_> { let (keys, scope): (Vec, Vec) = match &analytics { Some(a) => ( a.keys(), - c.out.iter().chain(a.exprs.iter().map(|g| &g.attr)).cloned().collect(), + c.out + .iter() + .chain(a.exprs.iter().map(|g| &g.attr)) + .cloned() + .collect(), ), None => ( group @@ -1400,6 +1405,11 @@ impl Resolver<'_> { .map(|(o, expr)| SortOrder { expr, ..o.clone() }) .collect() }; + // Over no grouping sets, `ORDER BY grouping(..)` is rejected before + // its argument resolves (`HAVING` reports the argument first). + if c.agg.as_ref().is_none_or(|a| a.analytics.is_none()) { + super::analytics::reject_grouping_calls(&exprs)?; + } let (bounds, below, extended) = if c.agg.is_some() { self.over_agg(&c, &exprs, env, "Sort")? } else { diff --git a/crates/krishiv-plan/src/tests/analyzer.rs b/crates/krishiv-plan/src/tests/analyzer.rs index c05b3ba6b..647bec823 100644 --- a/crates/krishiv-plan/src/tests/analyzer.rs +++ b/crates/krishiv-plan/src/tests/analyzer.rs @@ -26,8 +26,8 @@ use crate::analyzer::{ }; use crate::walk::{Node, visit_distinct}; use crate::{ - ColumnRef, Cte, Expr, ExprNode, Grouping, JoinCondition, JoinKind, Literal, LogicalPlan, - MAX_DEPTH, Origin, Plan, PlanId, PlanNode, ReadSource, SourceId, SqlArgs, Star, TableVersion, + ColumnRef, Cte, Expr, ExprNode, JoinCondition, JoinKind, Literal, LogicalPlan, MAX_DEPTH, + Origin, Plan, PlanId, PlanNode, ReadSource, SourceId, SqlArgs, Star, TableVersion, }; fn goldens() -> Value { diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index fcd401999..862da651f 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -12,8 +12,8 @@ use krishiv_types::DataType; use super::dsl::{c, int}; use super::tpch::QUERIES; use crate::{ - Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, - ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, + Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, + PlanNode, ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, }; /// The node kinds a plan uses, named as in the inventory. No wildcard arm: @@ -170,7 +170,10 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { ("filter", df(1).filter(c("a"))), ("where", df(1).filter(c("a"))), ("groupBy.agg", df(1).aggregate(vec![c("a")], vec![c("a")])), - ("rollup.agg", grouped(Grouping::Rollup(vec![Arc::new(c("a"))]))), + ( + "rollup.agg", + grouped(Grouping::Rollup(vec![Arc::new(c("a"))])), + ), ("cube.agg", grouped(Grouping::Cube(vec![Arc::new(c("a"))]))), ( "join", diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index ec013bb5d..24bfa55ad 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -426,7 +426,7 @@ impl Gen<'_> { }, 21 => PlanNode::Expand { input: self.input(d), - projections: self.many(1, 2, |g| g.exprs()), + projections: self.many(1, 2, Gen::exprs), output: self.many(0, 2, Gen::bound), }, _ => PlanNode::WithCte { @@ -824,12 +824,15 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { output, } => { site!(w, projections.push(Vec::new())); - site!(w, output.push(output.first().cloned().unwrap_or(BoundColumn { - id: AttrId(0), - name: "x".to_owned(), - data_type: DataType::Integer, - nullable: false, - }))); + site!( + w, + output.push(output.first().cloned().unwrap_or(BoundColumn { + id: AttrId(0), + name: "x".to_owned(), + data_type: DataType::Integer, + nullable: false, + })) + ); projections.iter_mut().any(|p| m_exprs(p, w)) || output.iter_mut().any(|c| m_bound(c, w)) || m_plan(input, w) diff --git a/crates/krishiv-session/tests/grouping_analytics.rs b/crates/krishiv-session/tests/grouping_analytics.rs new file mode 100644 index 000000000..a544dce4a --- /dev/null +++ b/crates/krishiv-session/tests/grouping_analytics.rs @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Spark 4.2 `SQLQueryTestSuite` goldens for grouping analytics (`ROLLUP`, `CUBE`, +//! `GROUPING SETS`, `grouping`, `grouping_id`): each statement of the vendored +//! `group-analytics` and `grouping_set` result files runs through a session and +//! its schema, rows (compared as sets of rows) or error class must match. +//! +//! The session has no `CREATE VIEW` or inline `VALUES` yet, so the driver turns +//! each literal `VALUES` table into a registered table of the same name. +#![forbid(unsafe_code)] + +use std::path::Path; +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Float64Type, Int8Type, Int16Type, Int32Type, Int64Type}; +use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use futures::executor::block_on; +use krishiv_common::error::KrishivError; +use krishiv_session::Session; + +type R = Result>; + +/// One golden statement: the query and its expected schema and output. +struct Case { + query: String, + schema: String, + output: String, +} + +fn load(file: &str) -> R> { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../tests/spark-golden/v4.2.0/results") + .join(file); + let text = std::fs::read_to_string(&path)?; + let cases: Vec = text + .split("-- !query\n") + .skip(1) + .filter_map(|block| { + let (query, rest) = block.split_once("\n-- !query schema\n")?; + let (schema, output) = rest.split_once("\n-- !query output\n")?; + Some(Case { + query: query.trim().to_owned(), + schema: schema.trim().to_owned(), + output: output.trim_end().to_owned(), + }) + }) + .collect(); + if cases.is_empty() { + return Err(format!("no statements in {file}").into()); + } + Ok(cases) +} + +fn spark_type(t: &DataType) -> String { + match t { + DataType::Int8 => "tinyint".into(), + DataType::Int16 => "smallint".into(), + DataType::Int32 => "int".into(), + DataType::Int64 => "bigint".into(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => "string".into(), + DataType::Boolean => "boolean".into(), + DataType::Float64 => "double".into(), + other => format!("{other}"), + } +} + +fn cell(c: &dyn Array, row: usize) -> String { + if c.is_null(row) { + return "NULL".into(); + } + match c.data_type() { + DataType::Int8 => c.as_primitive::().value(row).to_string(), + DataType::Int16 => c.as_primitive::().value(row).to_string(), + DataType::Int32 => c.as_primitive::().value(row).to_string(), + DataType::Int64 => c.as_primitive::().value(row).to_string(), + DataType::Utf8 => c.as_string::().value(row).to_owned(), + DataType::Boolean => c.as_boolean().value(row).to_string(), + DataType::Float64 => format!("{:?}", c.as_primitive::().value(row)), + other => format!("<{other}>"), + } +} + +fn render(schema: &Schema, batches: &[RecordBatch]) -> (String, Vec) { + let schema = schema + .fields() + .iter() + .map(|f| format!("{}:{}", f.name(), spark_type(f.data_type()))) + .collect::>() + .join(","); + let rows = batches + .iter() + .flat_map(|b| { + (0..b.num_rows()).map(|r| { + b.columns() + .iter() + .map(|c| cell(c.as_ref(), r)) + .collect::>() + .join("\t") + }) + }) + .collect(); + (format!("struct<{schema}>"), rows) +} + +/// A parsed `VALUES` literal table. +struct Table { + name: String, + columns: Vec, + rows: Vec>>, +} + +fn is_word(b: u8) -> bool { + b.is_ascii_alphanumeric() || b == b'_' +} + +fn skip_ws(bytes: &[u8], i: &mut usize) { + while bytes.get(*i).is_some_and(u8::is_ascii_whitespace) { + *i += 1; + } +} + +/// Find an inline `VALUES (..), (..) [AS] alias[(cols)]` (optionally wrapped in +/// parentheses) in `sql`; return the text with it replaced by the alias, and the table. +fn extract_values(sql: &str) -> Option<(String, Table)> { + let upper = sql.to_ascii_uppercase(); + let bytes = sql.as_bytes(); + let at = upper.match_indices("VALUES").map(|(i, _)| i).find(|&i| { + (i == 0 || bytes.get(i - 1).is_none_or(|&b| !is_word(b))) + && bytes.get(i + 6).is_none_or(|&b| !is_word(b)) + })?; + let before = sql[..at].trim_end(); + let wrapped = before.ends_with('('); + let start = if wrapped { before.len() - 1 } else { at }; + let mut i = at + 6; + let mut rows = Vec::new(); + loop { + skip_ws(bytes, &mut i); + if bytes.get(i) != Some(&b'(') { + return None; + } + let close = i + sql[i..].find(')')?; + rows.push( + sql[i + 1..close] + .split(',') + .map(|v| { + let v = v.trim(); + (!v.eq_ignore_ascii_case("NULL")).then(|| v.to_owned()) + }) + .collect::>(), + ); + i = close + 1; + skip_ws(bytes, &mut i); + if bytes.get(i) == Some(&b',') { + i += 1; + } else { + break; + } + } + if wrapped { + if bytes.get(i) != Some(&b')') { + return None; + } + i += 1; + skip_ws(bytes, &mut i); + } + if upper[i..].starts_with("AS ") { + i += 3; + skip_ws(bytes, &mut i); + } + let name_start = i; + while bytes.get(i).is_some_and(|&b| is_word(b)) { + i += 1; + } + let name = sql[name_start..i].to_owned(); + skip_ws(bytes, &mut i); + let mut columns = Vec::new(); + if bytes.get(i) == Some(&b'(') { + let close = i + sql[i..].find(')')?; + columns = sql[i + 1..close] + .split(',') + .map(|c| c.trim().to_owned()) + .collect(); + i = close + 1; + } + let text = format!("{}{}{}", &sql[..start], name, &sql[i..]); + Some(( + text, + Table { + name, + columns, + rows, + }, + )) +} + +fn register(s: &Session, t: &Table, name: &str) -> R { + let mut fields = Vec::new(); + let mut arrays: Vec = Vec::new(); + for (c, column) in t.columns.iter().enumerate() { + let col: Vec> = t + .rows + .iter() + .map(|r| r.get(c).and_then(Option::as_ref)) + .collect(); + let numeric = col.iter().flatten().all(|v| v.parse::().is_ok()); + let nullable = col.iter().any(Option::is_none); + if numeric { + fields.push(Field::new(column, DataType::Int32, nullable)); + arrays.push(Arc::new(Int32Array::from( + col.iter() + .map(|v| v.and_then(|v| v.parse().ok())) + .collect::>(), + ))); + } else { + fields.push(Field::new(column, DataType::Utf8, nullable)); + arrays.push(Arc::new(StringArray::from( + col.iter() + .map(|v| v.map(|v| v.trim_matches(['"', '\'']))) + .collect::>(), + ))); + } + } + let schema = Arc::new(Schema::new(fields)); + let batch = RecordBatch::try_new(Arc::clone(&schema), arrays)?; + s.register_batches(name, schema, vec![batch])?; + Ok(()) +} + +/// The error class a golden expects, from its JSON error body. +fn expected_class(output: &str) -> Option { + output + .lines() + .find_map(|l| l.trim().strip_prefix("\"errorClass\" : \"")) + .map(|l| l.trim_end_matches(['"', ',']).to_owned()) +} + +/// Run every statement of `file`; return the statements that disagree with Spark. +fn run(file: &str) -> R> { + let mut views: Vec<(String, Table)> = Vec::new(); + let mut bad = Vec::new(); + for case in load(file)? { + let upper = case.query.to_ascii_uppercase(); + if upper.starts_with("CREATE") { + let name = case + .query + .split_whitespace() + .skip_while(|w| !w.eq_ignore_ascii_case("VIEW")) + .nth(1); + let (Some(name), Some((_, table))) = (name, extract_values(&case.query)) else { + bad.push(format!("{}\n unsupported setup", case.query)); + continue; + }; + views.retain(|(n, _)| n != name); + views.push((name.to_owned(), table)); + continue; + } + if let Some(name) = upper.strip_prefix("DROP VIEW ") { + views.retain(|(n, _)| !n.eq_ignore_ascii_case(name.trim())); + continue; + } + let s = Session::new()?; + for (name, t) in &views { + register(&s, t, name)?; + } + let mut query = case.query.clone(); + if let Some((text, table)) = extract_values(&query) { + register(&s, &table, &table.name)?; + query = text; + } + let result = block_on(async { + let df = s.sql(&query).await?; + let schema = df.schema().await?; + Ok::<_, KrishivError>((schema, df.collect().await?)) + }); + if case.schema == "struct<>" { + let class = expected_class(&case.output); + match (result, class) { + (Err(e), Some(c)) if e.error_class().name() == c => {} + (Err(e), c) => bad.push(format!( + "{}\n error class {} expected {c:?}", + case.query, + e.error_class().name() + )), + (Ok(_), c) => bad.push(format!("{}\n expected error {c:?}", case.query)), + } + continue; + } + match result { + Err(e) => bad.push(format!("{}\n error {e}", case.query)), + Ok((schema, batches)) => { + let (schema, mut rows) = render(&schema, &batches); + let mut want: Vec = case.output.lines().map(str::to_owned).collect(); + // Rows tied under `ORDER BY` come in no fixed order. + rows.sort(); + want.sort(); + if schema != case.schema || rows != want { + bad.push(format!( + "{}\n schema {schema} want {}\n rows {rows:?}\n want {want:?}", + case.query, case.schema + )); + } + } + } + } + Ok(bad) +} + +fn check(file: &str) -> R { + let bad = run(file)?; + if bad.is_empty() { + Ok(()) + } else { + Err(format!("{} mismatches:\n{}", bad.len(), bad.join("\n")).into()) + } +} + +#[test] +fn group_analytics_golden() -> R { + check("group-analytics.sql.out") +} + +#[test] +fn grouping_set_golden() -> R { + check("grouping_set.sql.out") +} diff --git a/crates/krishiv-session/tests/grouping_analytics_api.rs b/crates/krishiv-session/tests/grouping_analytics_api.rs new file mode 100644 index 000000000..9c6381437 --- /dev/null +++ b/crates/krishiv-session/tests/grouping_analytics_api.rs @@ -0,0 +1,381 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Grouping analytics through the session: the DataFrame forms (`rollup`, +//! `cube`, `grouping_sets`) against the SQL forms, rolled-up columns' types, +//! `NULL` data against rolled-up `NULL`, and the errors of invalid uses. +//! The Spark goldens are in `grouping_analytics.rs`. +#![forbid(unsafe_code)] + +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Int8Type, Int32Type, Int64Type}; +use arrow_array::{Array, Int32Array, RecordBatch}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use futures::executor::block_on; +use krishiv_common::error::{ErrorClass, KrishivError}; +use krishiv_session::functions::{col, lit, sum}; +use krishiv_session::{DataFrame, Session}; + +type R = Result>; + +/// Tables `t(a, b, v)` (no nulls) and `n(a, v)` (with a real NULL key). +fn session() -> R { + let s = Session::new()?; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2])), + Arc::new(Int32Array::from(vec![1, 2, 1])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + )?; + s.register_batches("t", schema, vec![batch])?; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("v", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(1), None, None])), + Arc::new(Int32Array::from(vec![5, 7, 11])), + ], + )?; + s.register_batches("n", schema, vec![batch])?; + Ok(s) +} + +fn cell(c: &dyn Array, row: usize) -> String { + if c.is_null(row) { + return "NULL".into(); + } + match c.data_type() { + DataType::Int8 => c.as_primitive::().value(row).to_string(), + DataType::Int32 => c.as_primitive::().value(row).to_string(), + DataType::Int64 => c.as_primitive::().value(row).to_string(), + other => format!("<{other}>"), + } +} + +fn row_strings(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|b| { + (0..b.num_rows()).map(|r| { + b.columns() + .iter() + .map(|c| cell(c.as_ref(), r)) + .collect::>() + .join("\t") + }) + }) + .collect() +} + +/// The schema and the rows (tab-separated, sorted) of a DataFrame. +fn run(df: &DataFrame) -> Result<(SchemaRef, Vec), KrishivError> { + block_on(async { + let schema = df.schema().await?; + let mut rows = row_strings(&df.collect().await?); + rows.sort(); + Ok((schema, rows)) + }) +} + +fn sql(s: &Session, query: &str) -> Result<(SchemaRef, Vec), KrishivError> { + run(&block_on(s.sql(query))?) +} + +fn rows(s: &Session, query: &str) -> R> { + Ok(sql(s, query)?.1) +} + +fn sorted(rows: &[&str]) -> Vec { + let mut rows: Vec = rows.iter().map(|r| (*r).to_owned()).collect(); + rows.sort(); + rows +} + +fn error_class(s: &Session, query: &str) -> R { + match sql(s, query) { + Err(e) => Ok(e.error_class()), + Ok((_, rows)) => { + Err(format!("{query} should fail, it returned {} rows", rows.len()).into()) + } + } +} + +#[test] +fn dataframe_rollup_cube_and_grouping_sets_match_their_sql() -> R { + let s = session()?; + let sum_v = || [sum(col("v"))]; + let (schema, rollup) = run(&s.table("t").rollup([col("a"), col("b")]).agg(sum_v()))?; + assert_eq!( + rollup, + sorted(&[ + "1\t1\t10", + "1\t2\t20", + "2\t1\t30", + "1\tNULL\t30", + "2\tNULL\t30", + "NULL\tNULL\t60", + ]) + ); + let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, ["a", "b", "sum(v)"]); + assert_eq!( + rollup, + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY ROLLUP(a, b)")? + ); + + let (_, cube) = run(&s.table("t").cube([col("a"), col("b")]).agg(sum_v()))?; + assert_eq!( + cube, + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY CUBE(a, b)")? + ); + assert_eq!(cube.len(), 8); + + let sets = s + .table("t") + .grouping_sets( + [vec![col("a")], vec![col("b")], vec![]], + [col("a"), col("b")], + ) + .agg(sum_v()); + let (_, sets) = run(&sets)?; + assert_eq!( + sets, + sorted(&[ + "1\tNULL\t30", + "2\tNULL\t30", + "NULL\t1\t40", + "NULL\t2\t20", + "NULL\tNULL\t60", + ]) + ); + assert_eq!( + sets, + rows( + &s, + "SELECT a, b, sum(v) FROM t GROUP BY GROUPING SETS ((a), (b), ())" + )? + ); + Ok(()) +} + +#[test] +fn rolled_up_columns_become_nullable_and_grouping_is_a_tinyint() -> R { + let s = session()?; + let (schema, _) = sql( + &s, + "SELECT a, b, sum(v), grouping(a), grouping_id(a, b) FROM t GROUP BY ROLLUP(a, b)", + )?; + let fields: Vec<(String, DataType, bool)> = schema + .fields() + .iter() + .map(|f| (f.name().clone(), f.data_type().clone(), f.is_nullable())) + .collect(); + assert_eq!( + fields, + [ + ("a".to_owned(), DataType::Int32, true), + ("b".to_owned(), DataType::Int32, true), + ("sum(v)".to_owned(), DataType::Int64, true), + ("grouping(a)".to_owned(), DataType::Int8, false), + ("grouping_id(a, b)".to_owned(), DataType::Int64, false), + ] + ); + Ok(()) +} + +#[test] +fn grouping_tells_a_null_key_from_a_rolled_up_one() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, grouping(a), sum(v) FROM n GROUP BY ROLLUP(a)" + )?, + sorted(&["1\t0\t5", "NULL\t0\t18", "NULL\t1\t23"]) + ); + Ok(()) +} + +#[test] +fn grouping_id_sets_bit_n_minus_1_minus_position() -> R { + let s = session()?; + // `grouping_id(a, b)`: `a` is the high bit. + assert_eq!( + rows( + &s, + "SELECT a, b, grouping_id(a, b) FROM t GROUP BY GROUPING SETS ((a), (b), (a, b))" + )?, + sorted(&[ + "1\tNULL\t1", + "2\tNULL\t1", + "NULL\t1\t2", + "NULL\t2\t2", + "1\t1\t0", + "1\t2\t0", + "2\t1\t0", + ]) + ); + Ok(()) +} + +#[test] +fn mixed_plain_keys_and_rollup() -> R { + let s = session()?; + assert_eq!( + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY a, ROLLUP(b)")?, + sorted(&[ + "1\t1\t10", + "1\t2\t20", + "2\t1\t30", + "1\tNULL\t30", + "2\tNULL\t30" + ]) + ); + Ok(()) +} + +#[test] +fn operands_of_a_commutative_operator_are_one_grouping_expression() -> R { + let s = session()?; + // `b + a` is `a + b`: the select item reads the grouping copy, and the + // two sets are equal, so each group comes twice. + assert_eq!( + rows( + &s, + "SELECT b + a, sum(v) FROM t GROUP BY GROUPING SETS ((a + b), (b + a))" + )?, + sorted(&["2\t10", "2\t10", "3\t50", "3\t50"]) + ); + Ok(()) +} + +#[test] +fn empty_set_is_a_grand_total_even_over_no_rows() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT count(*), grouping_id() FROM t WHERE v > 100 GROUP BY GROUPING SETS (())" + )?, + ["0\t0"] + ); + assert_eq!( + rows( + &s, + "SELECT a, count(*) FROM t WHERE v > 100 GROUP BY ROLLUP(a)" + )?, + Vec::::new(), + "a rollup over no rows has no grand-total row, as in Spark" + ); + Ok(()) +} + +#[test] +fn having_and_order_by_read_grouping() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, sum(v) FROM t GROUP BY ROLLUP(a) HAVING grouping(a) = 1" + )?, + ["NULL\t60"] + ); + // `ORDER BY` follows the id: the grand total (id 1) comes first, descending. + let df = block_on( + s.sql("SELECT a, sum(v) FROM t GROUP BY ROLLUP(a) ORDER BY grouping_id(a) DESC, a"), + )?; + assert_eq!( + row_strings(&block_on(df.collect())?), + ["NULL\t60", "1\t30", "2\t30"] + ); + Ok(()) +} + +#[test] +fn invalid_uses_have_spark_error_classes() -> R { + let s = session()?; + let unsupported = ErrorClass::UnsupportedGroupingExpression; + // Outside a grouping aggregate, in any clause. + assert_eq!(error_class(&s, "SELECT grouping(a) FROM t")?, unsupported); + assert_eq!( + error_class(&s, "SELECT a FROM t GROUP BY a HAVING grouping(a) = 0")?, + unsupported + ); + assert_eq!( + error_class(&s, "SELECT count(*) FROM t ORDER BY grouping(a)")?, + unsupported + ); + // An argument that is no grouping expression, and a mismatched id list. + assert_eq!( + error_class(&s, "SELECT grouping(v) FROM t GROUP BY ROLLUP(a)")?, + ErrorClass::GroupingColumnMismatch + ); + assert_eq!( + error_class(&s, "SELECT grouping_id(b, a) FROM t GROUP BY ROLLUP(a, b)")?, + ErrorClass::GroupingIdColumnMismatch + ); + assert_eq!( + error_class(&s, "SELECT grouping(a, b) FROM t GROUP BY ROLLUP(a, b)")?, + ErrorClass::WrongNumArgsWithoutSuggestion + ); + // 64 grouping expressions do not fit the 64-bit id. + let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); + let query = format!( + "SELECT count(*) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ); + assert_eq!( + error_class(&s, &query)?, + ErrorClass::GroupingSizeLimitExceeded + ); + // A cube has 2^n sets: 17 expressions are too many, 16 still are not. + let cube = |n: usize| { + let keys: Vec = (0..n).map(|i| format!("a + {i}")).collect(); + format!( + "SELECT count(*) FROM t GROUP BY {} WITH CUBE", + keys.join(", ") + ) + }; + assert_eq!( + error_class(&s, &cube(17))?, + ErrorClass::GroupingSizeLimitExceeded + ); + let keys: Vec<_> = (0..17).map(|i| col("a").gt(lit(i))).collect(); + let df = s.table("t").cube(keys).agg([sum(col("v"))]); + assert_eq!( + run(&df).err().map(|e| e.error_class()), + Some(ErrorClass::GroupingSizeLimitExceeded) + ); + Ok(()) +} + +#[test] +fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, sum(v) FROM t GROUP BY GROUPING SETS ((a), (a))" + )?, + sorted(&["1\t30", "1\t30", "2\t30", "2\t30"]) + ); + // Spark reads the set's number where `HAVING` expects the id. + assert_eq!( + error_class( + &s, + "SELECT a FROM t GROUP BY GROUPING SETS ((a), (a)) HAVING grouping(a) = 0" + )?, + ErrorClass::UnsupportedGroupingExpression + ); + Ok(()) +} diff --git a/crates/krishiv-session/tests/scratch_rollup.rs b/crates/krishiv-session/tests/scratch_rollup.rs deleted file mode 100644 index f12c60530..000000000 --- a/crates/krishiv-session/tests/scratch_rollup.rs +++ /dev/null @@ -1,84 +0,0 @@ -//! scratch -#![cfg(test)] -use std::sync::Arc; - -use arrow_array::{Int32Array, RecordBatch, StringArray}; -use arrow_schema::{DataType as T, Field, Schema}; -use futures::executor::block_on; -use krishiv_session::Session; - -#[test] -fn scratch() { - let s = Session::new().unwrap(); - let schema = Arc::new(Schema::new(vec![ - Field::new("a", T::Int32, false), - Field::new("b", T::Int32, false), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(vec![1, 1, 2, 2, 3, 3])), - Arc::new(Int32Array::from(vec![1, 2, 1, 2, 1, 2])), - ], - ) - .unwrap(); - s.register_batches("td", schema, vec![batch]).unwrap(); - let schema = Arc::new(Schema::new(vec![ - Field::new("course", T::Utf8, false), - Field::new("year", T::Int32, false), - Field::new("earnings", T::Int32, false), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(StringArray::from(vec!["dotNET", "Java", "dotNET", "dotNET", "Java"])), - Arc::new(Int32Array::from(vec![2012, 2012, 2012, 2013, 2013])), - Arc::new(Int32Array::from(vec![10000, 20000, 5000, 48000, 30000])), - ], - ) - .unwrap(); - s.register_batches("cs", schema, vec![batch]).unwrap(); - for q in std::fs::read_to_string("/home/krishiv/krishiv-data/tmp/s-rollup/q.sql") - .unwrap() - .split(";\n") - .filter(|q| !q.trim().is_empty()) - { - match block_on(s.sql(q)) { - Err(e) => println!("ERR {q}\n {e}"), - Ok(df) => match block_on(df.collect()) { - Err(e) => println!("RUNERR {q}\n {e}"), - Ok(b) => { - println!("OK {q}"); - for b in &b { - println!("{}", arrow_cast_fmt(b)); - } - } - }, - } - } -} - -fn arrow_cast_fmt(b: &RecordBatch) -> String { - let mut out = format!("{:?}\n", b.schema().fields().iter().map(|f| format!("{}:{}{}", f.name(), f.data_type(), if f.is_nullable() {"?"} else {""})).collect::>()); - for r in 0..b.num_rows() { - let cells: Vec = b - .columns() - .iter() - .map(|c| { - if c.is_null(r) { - "NULL".to_owned() - } else { - let d = arrow_array::cast::AsArray::as_primitive_opt::(c) - .map(|a| a.value(r).to_string()) - .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) - .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) - .or_else(|| arrow_array::cast::AsArray::as_string_opt::(c).map(|a| a.value(r).to_owned())); - d.unwrap_or_else(|| "?".to_owned()) - } - }) - .collect(); - out.push_str(&cells.join("\t")); - out.push('\n'); - } - out -} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index 8daa1db39..f14c2e530 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -28,9 +28,9 @@ use krishiv_types::DecimalType; use crate::ast::{ Alias, BinaryOp, Expr, ExprKind, FromClause, GroupBy, GroupItem, GroupItemKind, GroupWith, - GroupingSet, JoinConstraint, JoinKind, Literal, - NullTreatment, NumKind, Query, Relation, RelationExt, Select, SetExpr, SetExprKind, - SetOperator, SetQuantifier, SortItem, Span, TableFactor, TableFactorKind, Truth, UnaryOp, + GroupingSet, JoinConstraint, JoinKind, Literal, NullTreatment, NumKind, Query, Relation, + RelationExt, Select, SetExpr, SetExprKind, SetOperator, SetQuantifier, SortItem, Span, + TableFactor, TableFactorKind, Truth, UnaryOp, }; use crate::validate::literal::{self, Number, TypedKind}; use crate::{ParseError, ParserOptions, parse_query}; @@ -452,10 +452,7 @@ impl Lowerer { (Some(GroupWith::Rollup), true) => return Ok(Grouping::Rollup(plain)), (Some(GroupWith::Cube), true) => return Ok(Grouping::Cube(plain)), (Some(GroupWith::GroupingSets(sets)), true) => { - let sets = sets - .iter() - .map(|s| self.set(s)) - .collect::>>()?; + let sets = sets.iter().map(|s| self.set(s)).collect::>>()?; return Ok(Grouping::GroupingSets { exprs: plain, sets }); } (Some(_), false) => { @@ -479,9 +476,8 @@ impl Lowerer { product = product .iter() .flat_map(|head| { - sets.iter().map(move |tail| { - head.iter().chain(tail).cloned().collect::>() - }) + sets.iter() + .map(move |tail| head.iter().chain(tail).cloned().collect::>()) }) .collect(); } @@ -545,7 +541,11 @@ impl Lowerer { }; chosen .into_iter() - .map(|c| c.into_iter().flat_map(|i| parts[i].iter().cloned()).collect()) + .map(|c| { + c.into_iter() + .flat_map(|i| parts[i].iter().cloned()) + .collect() + }) .collect() } }) diff --git a/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap b/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap index 2e9c27b99..ba9dd242f 100644 --- a/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap +++ b/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap @@ -21,9 +21,9 @@ SELECT a FROM t LATERAL VIEW explode(b) x AS y SELECT row_number() OVER (ORDER BY a) FROM t => [KRISHIV_UNSUPPORTED_FEATURE] The feature window functions is not supported by Krishiv yet. SQLSTATE: 0A000 @ 7..37 SELECT a FROM t GROUP BY ROLLUP (a) - => [KRISHIV_UNSUPPORTED_FEATURE] The feature ROLLUP, CUBE and GROUPING SETS is not supported by Krishiv yet. SQLSTATE: 0A000 @ 25..35 + => Aggregate [a] GROUPING SETS ((a), ()), [a] | +- Read t | SELECT a FROM t GROUP BY a WITH CUBE - => [KRISHIV_UNSUPPORTED_FEATURE] The feature ROLLUP, CUBE and GROUPING SETS is not supported by Krishiv yet. SQLSTATE: 0A000 @ 16..36 + => Aggregate CUBE [a], [a] | +- Read t | SELECT * FROM t NATURAL JOIN u => [KRISHIV_UNSUPPORTED_FEATURE] The feature NATURAL JOIN is not supported by Krishiv yet. SQLSTATE: 0A000 @ 16..30 SELECT * FROM t, LATERAL (SELECT a) From a41e41eb2161828041829a9e11a3b15f7a68d771 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:57:22 +0200 Subject: [PATCH 3/8] Address review threads for ROLLUP, CUBE and GROUPING SETS Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-common/src/stack.rs | 2 +- crates/krishiv-optimizer/src/stats.rs | 100 +++++++++++++++++- crates/krishiv-optimizer/tests/it/rules.rs | 48 ++++++++- crates/krishiv-plan/src/analyzer/analytics.rs | 100 +++++++++++++++--- .../tests/grouping_analytics.rs | 34 +++++- .../tests/grouping_analytics_api.rs | 48 +++++++-- crates/krishiv-sql/src/lower.rs | 2 +- crates/krishiv-sql/tests/lower.rs | 22 ++++ 8 files changed, 321 insertions(+), 35 deletions(-) diff --git a/crates/krishiv-common/src/stack.rs b/crates/krishiv-common/src/stack.rs index e44a30970..2a0488ec4 100644 --- a/crates/krishiv-common/src/stack.rs +++ b/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 diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index d96d5abf1..8e5c52108 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -241,11 +241,107 @@ impl<'a> Estimator<'a> { | PlanNode::ReplicateRows { input, .. } => Est::clone(&self.plan(input)), // Every input row once per projection. PlanNode::Expand { - input, projections, .. + input, + projections, + output, } => { let est = Est::clone(&self.plan(input)); let rows = est.rows * to_f64(u64::try_from(projections.len()).unwrap_or(u64::MAX)); - est.with_rows(rows) + let rows = rows.max(1.0); + if projections.is_empty() + || output.is_empty() + || projections.iter().any(|p| p.len() != output.len()) + { + return est.with_rows(rows); + } + let per = est.rows; + let mut cols = HashMap::with_capacity(output.len()); + for (j, out) in output.iter().enumerate() { + let mut null_rows = 0.0; + let mut col_id: Option = None; + let mut mixed = false; + let mut lits: Vec<&Literal> = Vec::new(); + let mut other = false; + for proj in projections.iter() { + let s = strip_alias(&proj[j]); + match &s.node { + ExprNode::Literal(Literal::Null(_)) => null_rows += per, + ExprNode::Literal(l) => { + if col_id.is_some() { + mixed = true; + } + if !lits.contains(&l) { + lits.push(l); + } + } + ExprNode::Resolved( + ResolvedExpr::Column(c) | ResolvedExpr::OuterColumn(c), + ) => match est.cols.get(&c.id) { + Some(col) => { + null_rows += col.nulls * per; + match col_id { + None => col_id = Some(c.id), + Some(id) if id == c.id => {} + Some(_) => mixed = true, + } + if !lits.is_empty() { + mixed = true; + } + } + None => other = true, + }, + _ => other = true, + } + } + let nulls = (null_rows / rows).clamp(0.0, 1.0); + let distinct_lits = to_f64(u64::try_from(lits.len()).unwrap_or(u64::MAX)); + let ndv = if other { + rows + } else if let Some(id) = col_id { + let base = if mixed { + // Different input columns (or columns plus + // literals): at most their sum. + let mut seen = Vec::new(); + let mut sum = 0.0; + for proj in projections.iter() { + let s = strip_alias(&proj[j]); + if let ExprNode::Resolved( + ResolvedExpr::Column(c) | ResolvedExpr::OuterColumn(c), + ) = &s.node + && !seen.contains(&c.id) + { + seen.push(c.id); + sum += est.cols.get(&c.id).map(|c| c.ndv).unwrap_or(per); + } + } + sum + distinct_lits + } else if lits.is_empty() { + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + } else { + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits + }; + base.clamp(1.0, rows) + } else if lits.is_empty() { + // Every projection is NULL. + 1.0 + } else { + distinct_lits.clamp(1.0, rows) + }; + let range = if !mixed && !other && lits.is_empty() { + col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) + } else { + None + }; + cols.insert( + out.id, + Col { + ndv: ndv.clamp(1.0, rows), + nulls, + range, + }, + ); + } + Est { rows, cols } } PlanNode::Union { inputs, .. } => { let rows = inputs.iter().map(|i| self.plan(i).rows).sum(); diff --git a/crates/krishiv-optimizer/tests/it/rules.rs b/crates/krishiv-optimizer/tests/it/rules.rs index f797e7a46..27189251c 100644 --- a/crates/krishiv-optimizer/tests/it/rules.rs +++ b/crates/krishiv-optimizer/tests/it/rules.rs @@ -3,13 +3,14 @@ //! (when the result is deterministic) the same rows as the unoptimised plan //! on generated data. +use std::collections::HashMap; use std::sync::Arc; use krishiv_common::error::ErrorClass; -use krishiv_optimizer::{TopN, optimize}; +use krishiv_optimizer::{ColumnStats, Statistics, TableStats, TopN, optimize}; use krishiv_plan::{ Builtin, Expr, ExprNode, FunctionId, FunctionKind, Literal, LogicalPlan, Plan, PlanNode, - ResolvedExpr, ResolvedFunction, + ResolvedExpr, ResolvedFunction, Scan, }; use krishiv_types::DataType; @@ -476,9 +477,46 @@ fn an_expansion_has_a_copy_of_its_rows_per_projection() { }; let one_set = rows("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))"); let four_sets = rows("SELECT count(*) FROM t GROUP BY CUBE(a, b)"); - assert!( - (four_sets / one_set - 2.0).abs() < 1e-9, - "{four_sets} vs {one_set}" + assert_eq!(one_set, 2_000_000.0, "duplicate sets double the rows"); + assert_eq!(four_sets, 4_000_000.0, "a 2-cube quadruples the rows"); +} + +#[test] +fn expansion_copies_keep_their_input_ndv() { + #[derive(Debug)] + struct Stats(HashMap); + impl Statistics for Stats { + fn table(&self, scan: &Scan) -> Option { + let name = match &scan.source { + krishiv_plan::ReadSource::Table { name, .. } => name.last()?.clone(), + _ => return None, + }; + self.0.get(&name).cloned() + } + } + let stats = Stats(HashMap::from([( + "t".to_owned(), + TableStats { + rows: 1_000, + columns: HashMap::from([( + "a".to_owned(), + ColumnStats { + distinct: Some(10), + nulls: Some(0), + min: None, + max: None, + }, + )]), + }, + )])); + let ctx = support::ctx(); + let plan = support::resolve("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))", &ctx); + let optimized = optimize(&plan).unwrap(); + // Two projections of 1,000 rows; the grouping copy has 10 distinct + // values, the id one and the position two: 20 groups, not 2,000 rows. + assert_eq!( + krishiv_optimizer::estimated_rows(optimized.root(), &stats), + 20.0 ); } diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs index 8173c6063..9e17591c9 100644 --- a/crates/krishiv-plan/src/analyzer/analytics.rs +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -45,7 +45,7 @@ use crate::{ /// The most grouping expressions one `GROUP BY` takes: the grouping id is a /// `BIGINT` with one bit each. -const MAX_GROUPING: usize = 63; +const MAX_GROUPING: usize = 64; /// The most expressions of a `CUBE`: 2^16 grouping sets (the SQL front end /// bounds `CUBE (..)` items the same way). @@ -623,7 +623,14 @@ impl Resolver<'_> { ) -> Result, AnalysisError> { match &e.node { ExprNode::Resolved(ResolvedExpr::Function(f)) if f.kind == FunctionKind::Aggregate => { - return Ok(Arc::clone(e)); + // `grouping` calls inside aggregate arguments still resolve, + // but ordinary columns stay on the input (no copy). + let rebuilt = + Arc::clone(e).map_children(|c| self.rewrite(&c, an, false, env, memo))?; + if Arc::ptr_eq(&rebuilt, e) { + return Ok(Arc::clone(e)); + } + return self.retype(rebuilt, env); } ExprNode::Subquery { .. } => return Ok(Arc::clone(e)), ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some() => { @@ -736,26 +743,91 @@ impl Resolver<'_> { )); }; // Bit `n - 1 - j` of the id, as a `TINYINT`. + // With 64 expressions the id can be negative (bit 63 set), + // where signed `div`/`%` on `i64::MIN` reads -1 instead of 1: + // there the quotient is floored before `pmod(_, 2)` (bit 63 + // is the sign). Smaller ids stay non-negative and keep the + // plain `div`/`%` extraction. let shift = n - 1 - j; - // The divisors are literals that are not 0, so the id's - // bit is never null: `grouping` is not nullable (Spark). - let high = if shift == 0 { - id + let bit = if n < 64 { + // The divisors are literals that are not 0, so the id's + // bit is never null: `grouping` is not nullable (Spark). + let high = if shift == 0 { + id + } else { + let high = self.call( + "div", + vec![id, literal(Literal::Long(1_i64 << shift))], + env, + Origin::UNKNOWN, + )?; + not_null(&high) + }; + self.call( + "%", + vec![high, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )? + } else if shift >= 63 { + let cond = self.call( + "<", + vec![Arc::clone(&id), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + self.call( + "when", + vec![cond, literal(Literal::Long(1)), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )? } else { + let pow = 1_i64 << shift; let high = self.call( "div", - vec![id, literal(Literal::Long(1_i64 << shift))], + vec![Arc::clone(&id), literal(Literal::Long(pow))], + env, + Origin::UNKNOWN, + )?; + let high = not_null(&high); + let rem = self.call( + "%", + vec![Arc::clone(&id), literal(Literal::Long(pow))], env, Origin::UNKNOWN, )?; - not_null(&high) + let neg = self.call( + "<", + vec![Arc::clone(&id), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + let zero = self.call( + "==", + vec![rem, literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + let nonzero = self.call("not", vec![zero], env, Origin::UNKNOWN)?; + let adjust = self.call("and", vec![neg, nonzero], env, Origin::UNKNOWN)?; + let one = literal(Literal::Long(1)); + let decremented = + self.call("-", vec![Arc::clone(&high), one], env, Origin::UNKNOWN)?; + let floor = self.call( + "when", + vec![adjust, decremented, Arc::clone(&high)], + env, + Origin::UNKNOWN, + )?; + // `pmod` (not `%`) so a negative quotient reads 0/1. + self.call( + "pmod", + vec![floor, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )? }; - let bit = self.call( - "%", - vec![high, literal(Literal::Long(2))], - env, - Origin::UNKNOWN, - )?; let bit = not_null(&bit); let cast = Arc::new(Expr::new(ExprNode::Cast { expr: bit, diff --git a/crates/krishiv-session/tests/grouping_analytics.rs b/crates/krishiv-session/tests/grouping_analytics.rs index a544dce4a..0920d7102 100644 --- a/crates/krishiv-session/tests/grouping_analytics.rs +++ b/crates/krishiv-session/tests/grouping_analytics.rs @@ -28,11 +28,8 @@ struct Case { output: String, } -fn load(file: &str) -> R> { - let path = Path::new(env!("CARGO_MANIFEST_DIR")) - .join("../../tests/spark-golden/v4.2.0/results") - .join(file); - let text = std::fs::read_to_string(&path)?; +fn parse_cases(text: &str, file: &str) -> R> { + let total = text.split("-- !query\n").skip(1).count(); let cases: Vec = text .split("-- !query\n") .skip(1) @@ -49,9 +46,36 @@ fn load(file: &str) -> R> { if cases.is_empty() { return Err(format!("no statements in {file}").into()); } + if cases.len() != total { + return Err(format!( + "dropped {} of {total} statements in {file}", + total - cases.len() + ) + .into()); + } Ok(cases) } +fn load(file: &str) -> R> { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../tests/spark-golden/v4.2.0/results") + .join(file); + let text = std::fs::read_to_string(&path)?; + parse_cases(&text, file) +} + +#[test] +fn malformed_blocks_are_not_dropped_silently() { + let good = "-- !query\nSELECT 1\n-- !query schema\nstruct\n-- !query output\n1\n"; + assert_eq!(parse_cases(good, "good").unwrap().len(), 1); + let bad = + format!("{good}-- !query\nSELECT 2\n-- !query schema\nstruct\n-- !output\n2\n"); + assert!( + parse_cases(&bad, "bad").is_err(), + "a block without the schema/output separators must fail, not vanish" + ); +} + fn spark_type(t: &DataType) -> String { match t { DataType::Int8 => "tinyint".into(), diff --git a/crates/krishiv-session/tests/grouping_analytics_api.rs b/crates/krishiv-session/tests/grouping_analytics_api.rs index 9c6381437..0259df485 100644 --- a/crates/krishiv-session/tests/grouping_analytics_api.rs +++ b/crates/krishiv-session/tests/grouping_analytics_api.rs @@ -328,14 +328,17 @@ fn invalid_uses_have_spark_error_classes() -> R { error_class(&s, "SELECT grouping(a, b) FROM t GROUP BY ROLLUP(a, b)")?, ErrorClass::WrongNumArgsWithoutSuggestion ); - // 64 grouping expressions do not fit the 64-bit id. - let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); - let query = format!( - "SELECT count(*) FROM t GROUP BY ROLLUP({})", - keys.join(", ") - ); + // 65 grouping expressions do not fit the 64-bit id; 64 still do. + let rollup = |n: usize| { + let keys: Vec = (0..n).map(|i| format!("a + {i}")).collect(); + format!( + "SELECT count(*) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ) + }; + block_on(s.sql(&rollup(64)))?; // Analyzer-only: 65 sets would still run. assert_eq!( - error_class(&s, &query)?, + error_class(&s, &rollup(65))?, ErrorClass::GroupingSizeLimitExceeded ); // A cube has 2^n sets: 17 expressions are too many, 16 still are not. @@ -346,6 +349,7 @@ fn invalid_uses_have_spark_error_classes() -> R { keys.join(", ") ) }; + block_on(s.sql(&cube(16)))?; // Analyzer-only: 65,536 sets would still run. assert_eq!( error_class(&s, &cube(17))?, ErrorClass::GroupingSizeLimitExceeded @@ -359,6 +363,26 @@ fn invalid_uses_have_spark_error_classes() -> R { Ok(()) } +#[test] +fn grouping_reads_the_top_bit_of_a_64_expression_id() -> R { + let s = session()?; + let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); + let query = format!( + "SELECT grouping(a + 0), grouping(a + 63) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ); + let got = rows(&s, &query)?; + assert!( + got.iter().any(|r| r == "1\t1"), + "the grand total (id -1) reads both bits as 1, got {got:?}" + ); + assert!( + !got.iter().any(|r| r.contains('-')), + "no negative bits, got {got:?}" + ); + Ok(()) +} + #[test] fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { let s = session()?; @@ -379,3 +403,13 @@ fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { ); Ok(()) } + +#[test] +fn grouping_calls_inside_aggregate_arguments_resolve() -> R { + let s = session()?; + assert_eq!( + rows(&s, "SELECT sum(grouping(a)) FROM t GROUP BY ROLLUP(a)")?, + sorted(&["0", "0", "3"]) + ); + Ok(()) +} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index f14c2e530..8f1fe3488 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -508,7 +508,7 @@ impl Lowerer { GroupItemKind::GroupingSets(items) => { let mut sets = Vec::new(); for item in items { - sets.extend(self.item_sets(item, written)?); + sets.extend(self.deeper(|l| l.item_sets(item, &mut *written))?); } sets } diff --git a/crates/krishiv-sql/tests/lower.rs b/crates/krishiv-sql/tests/lower.rs index 33b469bba..013f87160 100644 --- a/crates/krishiv-sql/tests/lower.rs +++ b/crates/krishiv-sql/tests/lower.rs @@ -191,3 +191,25 @@ fn deep_queries_lower_on_a_small_stack_and_long_joins_hit_the_cap() { assert!(lower(&dotted).is_ok()); }); } + +#[test] +fn nested_grouping_sets_lower_on_a_small_stack() { + on_small_stack(|| { + // 900 nested `GROUPING SETS (...)`: the parser accepts it (below + // `MAX_DEPTH`), and lowering must hop stacks through the recursion. + let sql = format!( + "SELECT a FROM t GROUP BY {}a{}", + "GROUPING SETS (".repeat(900), + ")".repeat(900) + ); + assert!(lower(&sql).is_ok(), "nested GROUPING SETS lowers"); + // A single level nests the same way and keeps its sets. + let plan = lower("SELECT a FROM t GROUP BY GROUPING SETS (GROUPING SETS ((a)))") + .expect("nested sets lower"); + assert!( + plan.root().to_string().contains("GROUPING SETS"), + "keeps grouping sets: {}", + plan.root() + ); + }); +} From c997748ea516a76f82f25eeb5f54fc302ad75853 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:59:15 +0200 Subject: [PATCH 4/8] optimizer: name the IN-list literals in full (typos lint) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-optimizer/src/stats.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index 8e5c52108..cbbd57f2c 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -260,7 +260,7 @@ impl<'a> Estimator<'a> { let mut null_rows = 0.0; let mut col_id: Option = None; let mut mixed = false; - let mut lits: Vec<&Literal> = Vec::new(); + let mut literals: Vec<&Literal> = Vec::new(); let mut other = false; for proj in projections.iter() { let s = strip_alias(&proj[j]); @@ -270,8 +270,8 @@ impl<'a> Estimator<'a> { if col_id.is_some() { mixed = true; } - if !lits.contains(&l) { - lits.push(l); + if !literals.contains(&l) { + literals.push(l); } } ExprNode::Resolved( @@ -284,7 +284,7 @@ impl<'a> Estimator<'a> { Some(id) if id == c.id => {} Some(_) => mixed = true, } - if !lits.is_empty() { + if !literals.is_empty() { mixed = true; } } @@ -294,7 +294,7 @@ impl<'a> Estimator<'a> { } } let nulls = (null_rows / rows).clamp(0.0, 1.0); - let distinct_lits = to_f64(u64::try_from(lits.len()).unwrap_or(u64::MAX)); + let distinct_lits = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); let ndv = if other { rows } else if let Some(id) = col_id { @@ -315,19 +315,19 @@ impl<'a> Estimator<'a> { } } sum + distinct_lits - } else if lits.is_empty() { + } else if literals.is_empty() { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) } else { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits }; base.clamp(1.0, rows) - } else if lits.is_empty() { + } else if literals.is_empty() { // Every projection is NULL. 1.0 } else { distinct_lits.clamp(1.0, rows) }; - let range = if !mixed && !other && lits.is_empty() { + let range = if !mixed && !other && literals.is_empty() { col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) } else { None From b3a96acfda20943ad6ee3f181d19d28e4bf9e2f2 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:00:24 +0200 Subject: [PATCH 5/8] optimizer: finish the literals rename (typos lint) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-optimizer/src/stats.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index cbbd57f2c..255ab67b3 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -294,7 +294,7 @@ impl<'a> Estimator<'a> { } } let nulls = (null_rows / rows).clamp(0.0, 1.0); - let distinct_lits = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); + let distinct_literals = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); let ndv = if other { rows } else if let Some(id) = col_id { @@ -314,18 +314,18 @@ impl<'a> Estimator<'a> { sum += est.cols.get(&c.id).map(|c| c.ndv).unwrap_or(per); } } - sum + distinct_lits + sum + distinct_literals } else if literals.is_empty() { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) } else { - est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_literals }; base.clamp(1.0, rows) } else if literals.is_empty() { // Every projection is NULL. 1.0 } else { - distinct_lits.clamp(1.0, rows) + distinct_literals.clamp(1.0, rows) }; let range = if !mixed && !other && literals.is_empty() { col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) From 80469d0ebe62b123ce31af2774ec7bb9c15987a5 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 23 Oct 2026 13:12:40 +0200 Subject: [PATCH 6/8] feat(plan-codec): carry the Expand node (format v5) PhysNode::Expand (ROLLUP, CUBE, GROUPING SETS) and the NullOf project item it uses made the codec's matches non-exhaustive. The proto gains ExpandNode (layout, projections, nullability, subqueries), ExpandProjection and ProjectNullOf. The writer format version becomes 5. A plan raises its min_reader_version to 5 only when it holds an Expand or a NullOf item, so other plans stay readable by older readers and an older reader never meets a node it does not know. The decoder refuses an expand without a projection and a projection whose item count differs from the nullability list. The v1 to v4 goldens still decode; the v5 goldens add three Expand shapes (a rollup with a grouping id, a single projection, a computed column beside a nulled one with a subquery). The TPC-DS ratchet shrinks by the four grouping-set queries (q05, q14, q18, q22) that now plan and round-trip. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-plan-codec/src/decode.rs | 35 +++++++++ crates/krishiv-plan-codec/src/encode.rs | 38 +++++++-- crates/krishiv-plan-codec/src/lib.rs | 2 +- .../src/proto/krishiv.plan.v1.pb.rs | 42 +++++++++- .../tests/common/inventory.rs | 65 ++++++++++++++++ .../tests/golden/v5/aggregate.bin | Bin 0 -> 1730 bytes .../tests/golden/v5/aggregate_final.bin | Bin 0 -> 848 bytes .../tests/golden/v5/exchange_broadcast.bin | Bin 0 -> 754 bytes .../tests/golden/v5/exchange_hash.bin | Bin 0 -> 758 bytes .../tests/golden/v5/exchange_range.bin | Bin 0 -> 764 bytes .../tests/golden/v5/exchange_round_robin.bin | Bin 0 -> 754 bytes .../tests/golden/v5/exchange_single.bin | Bin 0 -> 754 bytes .../tests/golden/v5/expand_rollup.bin | Bin 0 -> 887 bytes .../tests/golden/v5/expand_single.bin | Bin 0 -> 762 bytes .../tests/golden/v5/expand_subquery.bin | Bin 0 -> 952 bytes .../tests/golden/v5/expressions.bin | Bin 0 -> 1470 bytes .../tests/golden/v5/file_scan.bin | Bin 0 -> 495 bytes .../tests/golden/v5/filter.bin | Bin 0 -> 967 bytes .../tests/golden/v5/hash_join.bin | Bin 0 -> 1030 bytes .../tests/golden/v5/hash_join_anti.bin | Bin 0 -> 879 bytes .../tests/golden/v5/hash_join_partitioned.bin | Bin 0 -> 1030 bytes .../tests/golden/v5/limit.bin | Bin 0 -> 759 bytes .../tests/golden/v5/limit_offset.bin | Bin 0 -> 752 bytes .../tests/golden/v5/nested_loop_join.bin | Bin 0 -> 866 bytes .../tests/golden/v5/one_row.bin | Bin 0 -> 386 bytes .../tests/golden/v5/project.bin | Bin 0 -> 953 bytes .../golden/v5/project_without_nullable.bin | Bin 0 -> 392 bytes .../tests/golden/v5/replicate_rows.bin | Bin 0 -> 879 bytes .../tests/golden/v5/scan.bin | Bin 0 -> 742 bytes .../tests/golden/v5/sort.bin | Bin 0 -> 758 bytes .../tests/golden/v5/top_n.bin | Bin 0 -> 760 bytes .../tests/golden/v5/union.bin | Bin 0 -> 870 bytes .../tests/golden/v5/window_aggregates.bin | Bin 0 -> 1790 bytes .../tests/golden/v5/window_ranking.bin | Bin 0 -> 860 bytes .../tests/golden/v5/write_files.bin | Bin 0 -> 748 bytes crates/krishiv-plan-codec/tests/ratchet.rs | 6 +- crates/krishiv-plan-codec/tests/roundtrip.rs | 72 +++++++++++++++++- proto/krishiv/plan/v1/plan.proto | 28 ++++++- 38 files changed, 270 insertions(+), 18 deletions(-) create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/aggregate_final.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_range.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_rollup.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_single.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expressions.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/filter.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join_anti.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/limit.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/one_row.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/project.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/replicate_rows.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/sort.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/top_n.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/union.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/write_files.bin diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 552f70c87..80b6d2d0d 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -299,6 +299,20 @@ impl Decoder { count: self.expr_ref(r.count)?, subqueries: self.subqueries(&r.subqueries)?, }, + Kind::Expand(e) => { + expand_shape(e)?; + PhysNode::Expand { + input: self.boxed(e.input)?, + layout: layout(&e.layout), + projections: e + .projections + .iter() + .map(|p| p.items.iter().map(|i| self.item(i)).collect()) + .collect::>()?, + nullable: e.nullable.clone(), + subqueries: self.subqueries(&e.subqueries)?, + } + } Kind::Sort(s) => PhysNode::Sort { input: self.boxed(s.input)?, keys: sort_keys(&s.keys)?, @@ -364,6 +378,10 @@ impl Decoder { name: e.name.clone(), expr: self.expr_ref(e.expr)?, }, + Kind::NullOf(n) => ProjectItemPlan::NullOf { + index: usize_of(n.index)?, + name: n.name.clone(), + }, }) } @@ -545,6 +563,23 @@ fn aggregate(a: &pb::AggregateNode) -> Result { }) } +/// An expand has a projection, and every projection has one item per output +/// column. +fn expand_shape(e: &pb::ExpandNode) -> Result<()> { + if e.projections.is_empty() { + return Err(corrupt("an expand without a projection")); + } + if e.projections + .iter() + .any(|p| p.items.len() != e.nullable.len()) + { + return Err(corrupt( + "an expand projection whose items do not match the nullability", + )); + } + Ok(()) +} + fn window_plan(w: &pb::WindowNode) -> Result { Ok(WindowPlan { partition_by: usizes(&w.partition_by)?, diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index b1fa442f6..9c90308fd 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -243,10 +243,29 @@ impl Encoder { input: self.node_index(input)?, }), PhysNode::Exchange(exchange) => Kind::Exchange(self.exchange(exchange)?), - PhysNode::Expand { .. } => { - return Err(CodecError::Unsupported( - "an Expand node (not in this format yet)".to_owned(), - )); + PhysNode::Expand { + input, + layout: l, + projections, + nullable, + subqueries, + } => { + // An older reader has no `Expand` node to read. + self.min_reader = self.min_reader.max(5); + Kind::Expand(pb::ExpandNode { + input: self.node_index(input)?, + layout: layout(l), + projections: projections + .iter() + .map(|items| { + Ok(pb::ExpandProjection { + items: items.iter().map(|i| self.item(i)).collect::>()?, + }) + }) + .collect::>()?, + nullable: nullable.clone(), + subqueries: self.subqueries(subqueries)?, + }) } }; Ok(pb::Node { kind: Some(kind) }) @@ -281,10 +300,13 @@ impl Encoder { name: name.clone(), expr: self.expr(expr)?, }), - ProjectItemPlan::NullOf { .. } => { - return Err(CodecError::Unsupported( - "a NullOf project item (not in this format yet)".to_owned(), - )); + ProjectItemPlan::NullOf { index, name } => { + // An older reader has no such item to read. + self.min_reader = self.min_reader.max(5); + Kind::NullOf(pb::ProjectNullOf { + index: *index as u64, + name: name.clone(), + }) } }), }) diff --git a/crates/krishiv-plan-codec/src/lib.rs b/crates/krishiv-plan-codec/src/lib.rs index a451f1fff..0bf831492 100644 --- a/crates/krishiv-plan-codec/src/lib.rs +++ b/crates/krishiv-plan-codec/src/lib.rs @@ -54,7 +54,7 @@ pub use encode::{encode, to_message}; pub use error::{CodecError, Result}; /// The format version this crate writes. -pub const FORMAT_VERSION: u32 = 4; +pub const FORMAT_VERSION: u32 = 5; /// The `min_reader_version` this crate writes: readers older than this /// cannot understand the plan. diff --git a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs index 286966df8..7d11c857d 100644 --- a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs +++ b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs @@ -3,7 +3,7 @@ // This file is @generated by prost-build. #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalPlan { - /// The writer's format version (4). + /// The writer's format version (5). #[prost(uint32, tag = "1")] pub format_version: u32, /// The oldest reader version that understands every field in use. @@ -55,7 +55,7 @@ pub struct ExplainNode { pub struct Node { #[prost( oneof = "node::Kind", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16" )] pub kind: ::core::option::Option, } @@ -94,6 +94,10 @@ pub mod node { /// v4: a plan with a window needs a v4 reader. #[prost(message, tag = "15")] Window(super::WindowNode), + /// v5: a plan with an expand (ROLLUP, CUBE, GROUPING SETS) needs a v5 + /// reader. + #[prost(message, tag = "16")] + Expand(super::ExpandNode), } } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] @@ -130,7 +134,7 @@ pub struct BoolList { } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ProjectItem { - #[prost(oneof = "project_item::Kind", tags = "1, 2")] + #[prost(oneof = "project_item::Kind", tags = "1, 2, 3")] pub kind: ::core::option::Option, } /// Nested message and enum types in `ProjectItem`. @@ -141,6 +145,9 @@ pub mod project_item { Column(super::ProjectColumn), #[prost(message, tag = "2")] Expr(super::ProjectExpr), + /// v5: only inside an expand; a plan holding one needs a v5 reader. + #[prost(message, tag = "3")] + NullOf(super::ProjectNullOf), } } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] @@ -157,6 +164,14 @@ pub struct ProjectExpr { #[prost(uint32, tag = "2")] pub expr: u32, } +/// v5: a column of nulls in the Arrow type of an input column. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ProjectNullOf { + #[prost(uint64, tag = "1")] + pub index: u64, + #[prost(string, tag = "2")] + pub name: ::prost::alloc::string::String, +} #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct SubqueryPlan { #[prost(uint64, tag = "1")] @@ -190,6 +205,27 @@ pub struct ReplicateRowsNode { #[prost(message, repeated, tag = "4")] pub subqueries: ::prost::alloc::vec::Vec, } +/// v5: each input row once per projection (ROLLUP, CUBE, GROUPING SETS). +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ExpandNode { + #[prost(uint32, tag = "1")] + pub input: u32, + #[prost(uint64, repeated, tag = "2")] + pub layout: ::prost::alloc::vec::Vec, + /// The projections: the same output columns, each computed one way. + #[prost(message, repeated, tag = "3")] + pub projections: ::prost::alloc::vec::Vec, + /// Whether each output column is nullable. + #[prost(bool, repeated, tag = "4")] + pub nullable: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "5")] + pub subqueries: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ExpandProjection { + #[prost(message, repeated, tag = "1")] + pub items: ::prost::alloc::vec::Vec, +} #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct SortKey { #[prost(uint64, tag = "1")] diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index 123fa6257..2f673d7af 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -307,6 +307,34 @@ fn join(join_type: JoinType) -> HashJoinPlan { } } +/// `ROLLUP (a, b)` over a two-column input: the key columns are kept or +/// nulled and the last column is the grouping id. +fn expand_rollup() -> PhysNode { + let kept = |index, name: &str| ProjectItemPlan::Column { + index, + name: Some(name.to_owned()), + }; + let nulled = |index, name: &str| ProjectItemPlan::NullOf { + index, + name: name.to_owned(), + }; + let id = |n: i64| ProjectItemPlan::Expr { + name: "grouping_id".to_owned(), + expr: lit(Literal::Long(n)), + }; + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![ + vec![kept(0, "a"), kept(1, "b"), id(0)], + vec![kept(0, "a"), nulled(1, "b"), id(1)], + vec![nulled(0, "a"), nulled(1, "b"), id(3)], + ], + nullable: vec![true, true, false], + subqueries: Vec::new(), + } +} + /// One node of every variant, each as the root of its own plan. pub(crate) fn nodes() -> Vec<(&'static str, PhysNode)> { let keys = vec![SortKey::asc(0), SortKey::desc(1)]; @@ -383,6 +411,43 @@ pub(crate) fn nodes() -> Vec<(&'static str, PhysNode)> { subqueries: vec![subquery()], }, ), + ("expand_rollup", expand_rollup()), + ( + "expand_single", + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![vec![ProjectItemPlan::Column { + index: 0, + name: None, + }]], + nullable: vec![false], + subqueries: Vec::new(), + }, + ), + ( + "expand_subquery", + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![ + vec![ProjectItemPlan::Expr { + name: "shifted".to_owned(), + expr: call( + Builtin::Add, + vec![col(1, "a", DataType::Integer), lit(Literal::Integer(1))], + DataType::Integer, + ), + }], + vec![ProjectItemPlan::NullOf { + index: 0, + name: "shifted".to_owned(), + }], + ], + nullable: vec![true], + subqueries: vec![subquery()], + }, + ), ( "sort", PhysNode::Sort { diff --git a/crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin b/crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin new file mode 100644 index 0000000000000000000000000000000000000000..35a39060157d4948f6e7d071a38ec8e4588b88ea GIT binary patch literal 1730 zcmcgt&2G~`5Z+n;ZMG?OHwBEv!OE$M#3d)x15`bsiVLB~BD$)b#Hwpo_C^giE=c9Z zfd}9PcmxheJOUE0&;#!fX5zL=)bJNXXEiJD%#J_L&bMRX_&hx0@(mZBdM=;I)#^l- zM=uYx(OEH&UYzvxC^Hg_@}+%T#EBBptxTC78r7DLDUOm{-BWj^eNq;&T07NcGAu?_ z@?h!;R|F`~As0>rVMA~Q6)qalUdNK4J5oX=oixvrvMVG9?%M%urNUC=4iB1)F@b|I z`(kDcaJ7Nz93+G_k8~jH8tGr^6#TqS zl#BC6xm;)B%*ly57yA0yeCBh>h3nreC%h#&@wa!37+4|$xn)QE1_fC+YLJK zscdq*Y47+VI#TQSf_jF!jY{kN$z(E9J{h5{mf7;vbR_#F`QOjS{r1Ad7+W_ncCCrk zmyRtYXLS{~eDl-zj&4KRPm{h;kkya6uVCDj(&?pxp$P%HEot|P!YKR^7@cO=yBU9I z?I;}$)7(U-x{A_4UX)3^8wl_xpRM5a+sXG&A3wNlPv<&&Rhi!E++Uy@#08aLZh)116NHrAB& z61^;6n009ptoG6USZBFbIBsoI99uodk*!biNMHOmCf;CSTR-Hhdmw)?oadJ!%Tv+$jQ2;!sr zFmnwN0DlRK+n<4w*%GMc0K1I;Ro+Fwjr=CPB;6-v|83WG+{!gLMsy!!Ol^(nu*e}q zmdC?n-)`4=jKRVqQH0DxR~)X$5a`I9rzdi3xCmLX8-^I literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin new file mode 100644 index 0000000000000000000000000000000000000000..44186479a51888279684b067d8549bebdeb9267d GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M?28vloQ1D6FfHr8HNe6cD9Chz3K61gqG0XT{i&y{AA&gOqeB zcmQ62N1#CBVN&o8X2uRk5p?Y8`|SDsJnL0nDB}x>pQH-(OMHQ|?##5OZ%&LgRkOw* zFIHw-Sp=&?bl)|3p%wa_ZOap@7wFmMw5auceGlEUw#oJGxoL|{v+at5Ar$W@&|qSz zyaeKjk{UWPCdsSRK`<#Ibm*0JU9?Mu5@hH`a9Aowb2xrG6+$T5Lbx|91dz2)2#G^+ zN96;98cACus8U{^^;h&zM91{773(R@?YXC z^@DIhI3lq9y6<}y<9zCm$?#Dra&Ee$T-QxoO*I<*{xfX#c70- uva1_gZ>uVT1&BkbPYel}0{5FS&6)c%C-rgL;mowM74-|k;p=wGY= literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin new file mode 100644 index 0000000000000000000000000000000000000000..87676f17bdb8c6e75c1bcc692a3a0b666b25651d GIT binary patch literal 758 zcmbtSy>8S%5Z+n;?af_?vloQ1D6FfHr8HNe6cD9Chz3K61gqG0XT{i&y{AA&gOqeB zcmQ62N1#CBVN&o8X2uRk5p?Y6`|SDO*w@1c9vHo4wCH*K+Lwq0>BxZQL5dv3__)B}^@%j8pYpM;kjn2Cqr4Sc?x7jO+;ym56iG-q*L&dCYm z;Nq{r$KE31gW%;3GUj!p|L=hrs&>5_6?H*JwWixh8S%5Z+na=e@~=bM}HT7KL>cvXtg3l!8R55Te0QBEc&5-B~eqWbY}^(I6#V ziabDGAdf(S#KWZE9n6d!kRs^V)%V#mv-7*Fyg)|h62D62>o@osx7$-TRlh5)+c$TAL~cxo;G!+r)Q?g*Y&2&_j*&br9gv; zr1D~j#!6~vNf^g(69>U0h|r=}R8`(A6iSeR8^V5}9L?tF!$b(7Xba)qu@FGgE+HTe z!~-!U8=ne@LqbfTQ}PM15)-m+^}aEy2lqaPCE;kIZvCBJFK!Pa-2`eHO{gpx0=fzU zh#|R@|0;u%{veza4hby3>$;B9@_EXHasOo+Q3)g3ykyTz+zOX*b>GkD7`(WAb=Nw3 zaa{Jv2t)hgul9%BBII%Ka{E2_wIu)V*bK#3jM06SFRg}h_%%xnpB9K-TCCPK0$9w@ zP3zifdP3F|CEdD^&zvJ|;&oBk_{_9%v8w7O&t8T!=kB``e!1v=e!ICI&3#ju@^fp` z?R)Phjgk<@MO)Rj+LUDoa}fJd$wPPnsjKDF-E@}hAUr1vW&2}#n9R^IRmN%ueye{0 D;~ro$ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin new file mode 100644 index 0000000000000000000000000000000000000000..c4f0eabdd9a59858b29e2e6c03d64928e4e5353c GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M*I;vloQ1D6FfHC7P>H3KFG4hz3K61gqG0XT{i&y{AA&gA}?H zJOD4iBTz)j!=&IH%#0n7BIwxF_u2FNdDg4EP{tP$KT8$pSNIaQyED_CzCAJ4RLu&5 zyjYq|Wf80n(S6_Kg;wZywk=Pro}*`*)1uZ7^nG;C+9ub#=cX;z&890J521KVfd&&x z!MvKlpsSlf`d{yn#1wCsSrZZ7Q(${A%LuXLP#8n z9kEL>&I*YmLPDS{`H)zNDMh!2+JxO>`Ow3XaXeKw!QS8(H;1umeYF$IxkwZWE zscXUo;fTQYo4)T^jPt2KCc~4-Ct;s9FF7z155VGI@74v}gBO3ax*eFSxG&e_gmGx` z*U+Q3i1-kAxx*goIt)x;WO9 z75!@=pE$?bB{PF#pKMzc8s*jz`w(s1UEKVbw vlwIA}dQ(*q%t0JTC6C}4WUiKv_On^Khwzjlr1mFtFP))d>YUvU{C58WYo1@V literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin new file mode 100644 index 0000000000000000000000000000000000000000..b9c7afa741c7cac571b6fd349b8e0010bb9f5a35 GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M*Jk*$cv06xLP9Qktt!iV&qjhz3K61gqG0XT{i&y{AA&6Df2Q zJV0I`k3fOM!=&IH%#0n7BIwxF_u2FNdDg4EP{x-MFQf|eD}0ID-I-}m-<}w2s%C{j zUM$U~vItg(=)Q0ALM!w;+mLnz)-puxma zc?rZ5B{g(pOp@2BgJ4ob=+G)CCJc?;Gk5F=5YLODuht9g>Y|K2q0^Z5E6&t zp4g!nXNANOAt6wfd_=6ol%iWhZNl!ceCT1xIG(DTV0Z9~o2RjAeKn1x773(ZIdPHa71AHP2cw{#`)AAli|taldw;lmmHXh2Vn89ck2S~!HYjy-44uE+?Q)|!Z@_} zYv@s1M0^Om++mM(9oheTScZBmCg{E>mR3VGzM7?m4+}&uD_3h911x6fW=&%?T^wu5 zivG2bPn=_Il66_zSJfK?K`(7i_-`v uWmh-0-c(fta}Wno$s>3MnXBdFy=<25A{k~D>H(^r5aL4Mh*Xx^Nvs+>vNviG7ledV zZX9?3UVum7fW*V}z&nH)yDF-nCw4WTJwM;9zws(B^x`v#AEgTPb9{!|&C1lrheyVm zqFP{(W^=PHEP~Y`x^Jp9(+d5@*7=duWAtowoR#{vzJ>0pu2MZXF?F`A)=hT5O~qRZ zG?-W_FM+tHq=tr^Uh-n%AeadvH0b4JnblK;5@hH`u$L=Gb2@(27eXl7Lbw+w1dz8w z77~Wyh8R$et3twvEFq&Rd52JmK4rIBZo=V?-uAHM9Q4(BFl^)E{89Yw1DtC5iuf*i z=+K8_fF3J&jYB=atwmj}S9!TO%u^cNpQ8A>BA<)muG6(e!Po&12~0lAf9Yj)KgrI> z_Q}|Px!G)3AJ>y4Oxh=dPg;H2ykv_EJTun&SD&s;=asD$+|gxq06TSbm`StzVSAQIs(iCqSJKMf;QD^Za!d~8# zl`Ypr5y2S5o>X!l9>dJl^4`vDG#Mg1qzs*QnYuX{p<~LF!wy`AiLon~w<}xpe*8S%5Z+na=e@}tarS~R7KL>cvP5$gN-Sv^P3*2j6JAL<9_o-}o)w@*!zuj@^lKj}@;mI4hX zlFExA8Y`)xC1D)DNgM=|AVP~?QB`@fP$)qLZU}pYax|Nx_Y)z6qAi4b%R&H2yM%x^ z5ckA3+1M)}4hb=VyyPQdB_?Fu>SJS858-_dOTytq-S|7bUfk?Q!8uGdbrHEJQOrTa zkbIHv-;fLBHM9c*DN)BTp)UBv0B>*U@=2C zt!u040a;U&bmu}obq=+O*F|OHQ`5%9s;Zkjdl4!Kx@#vqyX<~meg8I^`=&DGr`D#^ zJFh2=k`Tv5Th+GOlw}BXJ$q8g19%RptL5X}be8NOJR=KL`(wJF%+N7a#%c$CtA7C; CJzw$w literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin b/crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin new file mode 100644 index 0000000000000000000000000000000000000000..1f9e9d7f857fbc0e6b84da5b9dd8dc000ed138bf GIT binary patch literal 952 zcmb_a&5G1O5bo+Elb)I#C)>MXY&?WHNyN*XWI-MFB;v)LR^UrvJN_I!s<1=i{VJv-B2gtW~Ub?w~f@o^ip3z1m+{W zg7=`QCNsjJ9$?rmi@D3QC^SPHn7;m{7z3nF)z{ZI3H4tqo~{eyA3m!%k%^bSlYeUr zx8Xa{G0~KWb~mOCSbg*2mv2`;zI^`F8;7>C<%`zMhUYeV zX3`k@MO)RbIxNc=#vpd2l2f5*)(q{9cB?u@zp^6IwQYC~$t*n!^>e|uT zt5PmqAr-YZ4jg)e3kMEdfD;l@5&bt17Y;peh=jzk=8boDm9)Jv*1UY?y?Jk(@AE9P zfd=;}zekxXKI9+pR&gASZ+?0Zl~I}>ad#LWMw7JUq!hrdTlp{+j5|d+P7X@Z<#w6h zj5Bdwta0mjoDW6&b~KJh`J{+nuR}qD5kaDWGCL%J$EY9$F1+xgp2Z> z4;ZI}g5{I_gqcExgBwl5Fc?0DVQFg^1Xn9K0WBCyMjK;l7ohcVhB$E--9XEXCPo`| zZm8U|^twigi>poc#O>66?wCA`5FvkIO=C}=FjyteC$6Z*N;~uQe@e-Vyoe=lI z*R%!aH4Pxpv<;!A9k{4z7rL7E;FXZ{zM?OZUy1z~=Y5=4aH{nuRaGgu)qYf%hxK1& z9F*R{M-5tAWmF55{9m8Vb5#vBRNI(txyn*?Rn8FStnJ2gZJ(WO`06uI!>XTIUV-cX z`?*9|W8iV?{rIpHB%NV&D9C36&TW|-jmm(K!7{gGo|gizP8lUBe(eUmY+e;nI7+fI zyd4!`a+KxcczDfc1ghUmfBK;M`_~`7b-HepMd_WQl&v%OMg~2fdr6VyWj0AupLB__ zPZ_;J-XqcybaO>6_d1-sjS)O;5nAmnb2G|@%59Q`8#A~en}Tz^%B%se`CL^2B5h&% z6_Jjx!l#uf@eXWagQ~V{iAMjaLN$nVh5Z()G1c(su51YJUG(5MxhpS-jcWl1)v!l@ zO#Z$1=AG}Z|Md80F+Z#=TCy?v{rfb{!;^0w)@436jmlLg-}vXR>#E#nIz6dXQSWw4 zBp2*+hphwgsRXg~1!=J9!yo_r#YT_ytYCb=*qWfs@iB41^c0)VraZ6l`f0_^XuST? zB6M2g3mXgft0=CE|F8+R@Ab~9i(X7_LC2?EvbIV3FXK)}w!1Cn2|J&Z@i?HYDcAN_ j@x42Y3pYxWsECK?ZiPgy^aOVf@(h2NeG_w+RPX--6Rp<0 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin b/crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..50724761fef2212a8852bd0822410443c91392fa GIT binary patch literal 495 zcmZWmy-ve05I(zU!k}@Ff^nG_H`@2uFMYYNOMBWIYj@6NLCuVv(i>%s~Tj~B`5M0i#t|J?m zpw<38+EnY(nM{*VHPU!6j>5f(u|vNhm}17~qN zt!KUj00I-hSuk|~pNL1qDT3yUJkNKq(tQ9Xby10GsHRj4EA#jse zc#?|tk|snnVFLlMS~!cM_i!$~!J{(@r^r&5;#}%h{^>Mb2FA&od>v?D1<(i#vMtD9 z?bM}n9`&taqp~%9#qK*&R(tVaXnn>y`0BkRvC0i2YvK{UF7Sg&R~bEw$5!8&RO5Ff c(?6FU-2Mz^w|weO_m|lHj+tDGsfgf@kZayi!-y6>yK6I#W9*rpZ_Jy3KF&klvD9 zaifLwVy;J0DBhwmiXSBoaFYP=7QCXW@@6JMa1pu@-!G)22-Z)>j4_FuG51_E#?jhC z2~mga65BzHqC(UWN{oV0;u>nn#)xk9xiPti^q#|_aWIx=!EWzoXLqzvO#6u|dC0ZZ zjlu&(qv9Tdm77;#1D{g%opsm#3~0 r*Z0y%vJ3n+B3Nw$T}~$8m?|T;!#D1~l6#xbU&*~-1N}@c{U7e{Zj^Eb literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin b/crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..ddfde900465a97dd93e490ca1dcf1663622258d3 GIT binary patch literal 1030 zcmbtSzi-n(6ux)1Fix424Oz#t!VovQ7>< zUe0$I6PLtAZX*x(N3jyCe5E+oZZ_|YHn(Hdl+EDGkJhqSI=c2;k!)pgw7x6V;BQwh zFWjAHq;ui*D{XH(J@9Gpoc@U|&CMtDee?}!SLoc#GXe%Sjv2R-NnX@HV}!q0kw(`K8^h?|HWpr(M*|~sb=z_BN5{8s6IV|_i&?L2UWNYuPD)~punFdq zWU2(swnsS?^lb=WMd@s=JW4~HMpP6^;ICCO&G4fUcF!0K>Ce+#`RB6oaZE*-#E)H0 zoy|AHU4PhIef{#eF|cJWvv-w>IydGQd4UVfw8{&Wud>Xg14>4Wvr`;XWD0htAN7JR W(1(cNYrmoWpa+J`V;wf=b@&$mJbirt literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin b/crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin new file mode 100644 index 0000000000000000000000000000000000000000..048d0341b790566fb4f1ddffe0fdfd5ede05a9fa GIT binary patch literal 1030 zcmbtSzi-n(6ux)1Fix424Oz#t!VovQ7>< zUe0$I6PLtAZX*x(N3jyCe5E+oZZ_|YHn(Hdl+EDGkJhqSI=c2;k!)pgw7x6V;BQwh zFWjAHq;ui*D{XH(J?hioIsFq`nwwAP`{*0!I=?-cOf)A#7qRnumcYdTk8{=e!<#AsEy zc|Gnc0qOKp_60m>ac+gl)Nv8ZgS@9r|`2YFw^k#js%m*IXN7nhHdSvHAw*9IwMpXI+! z!2Ny_E(k{i_Fs{VJMw;Nf@%L`8d3KVEjHOR6Hmb>_;OPhat$^>Ag`nk1E)8FGv zRnrF-@}YOEO}eQXm!6w0t=5fg%i?LQAlm)#;t!v8zrKI_I$DILG4=b-<+B^VPM%~j wj;pS*uG!Xg3=0s4QpqED0=ci{{eyg-?IS#*40ZdHx}D9@Gflx^53a+%0R$>&D*ylh literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin b/crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin new file mode 100644 index 0000000000000000000000000000000000000000..d55e972c842f15a43345d50a9d77725d3eeb3e44 GIT binary patch literal 752 zcmbtSyKdA#6rDSM&E9N?GaH1lD6FlJr8HZi6d_855DkVB2}ZGZ*NU+tdscytCQ|4q z_<;OCK7j&>k4eEln0xJj6hX&~&avp?FJy1`|u= zB@j=P)Xl4hYF zgiFE^f$cYa-?J3wQ+G^;FOyHQK208SU?yHZOMksv7jO?AykvE|U9RH3T$2;VEsMXl zJZg)GkAR0e%(1Q``+twhP>;m~-515uYN*Cvv()fmf#_xBYHeeH#SGo7X{@GaV@+An zO$+(NIo2jwm$gmKO_!9bx@n92d88ohzdO;V&;8GD*H`{LFtw>Zb~f9-^JcO*jc`(S tbz|#IRYfodaUhjEf@hGqT0Y*(X6Y`%Aw@{-Pw0L+L&wxPyB+xL{smTIUrGP~ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin b/crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..0f49b3f684db38f274f53e354754e1257d3ef418 GIT binary patch literal 866 zcmbtSy>8S%5Z+na=k?^m`SyY^7KL>cvXtg3lmenu2+?4u5E86n-<=g>NA{jVx->|k zqu>E}0Um(@iHAwSI~+5<=tKaCf?a)|J%8Wqj;*{<#+MS`OXch5_zZWuQ`4TjJ~q}= z%>wNqj3~xgA#p@V2$Uu75i2pI=uTgouzOS9YgjUlhU&^6_kMBpD3)5pU#-_`?#%hr43qxJG@w%qX!DXiGx2TsB%f~9`K-Z< zZ=kMk&sD6;H92830rAHKHoZl}6XE6ddsEku{lAklq%jN7eOk<|hHCR_<{IA35WTEi zENu)ho1mLDjn(uQSyNW@;e~wQ9BGp*%i1Pqrc262-LysiBvKHrzd6yxhxL!IUp^0} zzNt-h-q~#T#xIn`X@o)9)s3xJRTaS$#DP@u5FSJ3YI%Pzo1|lehZLc1+tBTFf{v+k cb~~_j|5M!Crv5JO`Pii!N#-+}ZkuRfS z$fPWS7>I9weL3G?Xu70#ne5`b>}iMvPKdx!O4eyEX`i!E@*6{?ntc=Pv+wYq&+%sn z$=RH%uumOuL;}Yq@zm=*qi))Du%e$wzqaC`HfFi%Hrr4MQ6C$%8pmMsrVYO9r%FOz qEUv6~+inQ%-cPRE4C7tfT$frD@oP~&y~Ll-kN4~>_k-^rr|=8UmMiH1 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/project.bin b/crates/krishiv-plan-codec/tests/golden/v5/project.bin new file mode 100644 index 0000000000000000000000000000000000000000..cb7da8071df616756d098deac40e248a474db82e GIT binary patch literal 953 zcmb_ay>8S%5Z+na=e<0FQtK5)YGtcQ|J3=%fH0>}tk4Ki|xFW>gSKb0+bvRH1&1k8q|d1mN|Ub~n_;V3{41f(z2w>NhP_Fqfh?JNXtN;K2 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin b/crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin new file mode 100644 index 0000000000000000000000000000000000000000..91115ce8affabf703ee8dab06fc19e701b060b4c GIT binary patch literal 392 zcmZWiF;2uV5FGn_i7^Nw-xX0rVO>d?iV{&Gs;jTDl-ymUD3ORIpreG2f(P&coAdHKh zB_Urz#gI{11Thfb{`ykB!O(O`A2Qj+ciB@BGc1X~SxnY(&S{^rk@Fiv#hQEr?UV2H zpU?4U2+7%;t8h#ma6|&fM)B0kBcpEXZLp%9M!&Vy=wD2oh#f@ zpuvPvc@cyoB{j69jH2hUgJ5DrXwfUGDsK{n5~S}2uwN)gvpak>7D6c6Lb#VG1dz5z z;uHGfn%E&9pYjOPtBM9ZSG(TQoJVo}vi zo;?oeao0bb;N8dduWw(!45mXGqXjO~W|AQPEbltyX0jz!b#3RB{H7Aa%98 ey_ZhnU4#eZp==x1^>~7gsWLV@urU+kuKOE#TV@6T literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/scan.bin b/crates/krishiv-plan-codec/tests/golden/v5/scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..8d48e76e5683a6f6a4c8cde8f70b6fc8c4904fb2 GIT binary patch literal 742 zcmbtSy>8S%5Z+na=e@~=ID0`Di^94JSxR#iN&!(SglI67NU(~1cUH_gviB6|Xpllj z!2|FDJOTw050ip-Ff(>QilAdg-)GPNj=d^~WO6C-gH)k@j?ZwnJ2&mwt5f4l-K;Sz z%9YvH4#DXN{WncfYK4=|wbiN93k+OyR$6^m-$DPpZ3?}2VcK%jY`gMs2qilTG?+xH zAcbVAq=t@+Y5F4b5KM*$9R`)PWxG@;K}LQIN0st4hm+ScA%r3o!oOxAfUE;TL>!4* zVvk~+6%ogTlt5YX9mWTvjd{lPD;9%Y~8zXo#qAB0Q7F@f!u zec$r{&ZnN34o{{D4V#d%$bp&o&^++-&AO0#u<$|E^^#o0eYvI}jEjoD79D4cm=|a9 z2OaA=vj6u38S1f|qW`p9ISuvrYE~NFE)j#gT5nteu$-fxH;vPDS)8eA`oBUx@J_Tz zH8S%5Z+na=e@~=ID0`Di^94JSxR#ix&oq95u%Ahi3F?IcW1@ek-eutM}w4f zDR=;0fJdM}%EP4K9n6d!kRs^V(f8Riv-7vR$_r$4De<#ZzCOWY+-}cJbN2SsSX0(3 z^s{_vHl;V(^mXyPZy7hN^ySRN3O%|wS2-;N;Kn%&X z{8s~9?uKwlI3Te8rt3N`$or`V#{HLRM8!uWykyTzJOnrJ<$j#cIe2m7>TYO`;=CM_ z5r)CVUxN>|MaT!i%k5=|Yf1m#12dFkK1TOdzO)+3;nyrRe3&D8X|Y<{2w*-%H?3={ z>G4=olyu)hK64JWiPuGC;|tTq#j2{CJbM|^Id?yt@Y7ZI^ZU22qnU3iQ+{l1x_$5E zq)`&$xM-`|R-3X6VFqF$l{|nWNL?+T?xoXY7vVX1sM;Ub!(@t%sWLV@@Z08S%5Z+na=e@}#arS~R7KL>cvP5%Llp;h|Aw+|rM1oc9yR%~K$lg;(mnKru zQSbn~0FOX{#KWZE9n6d!kRs^V(f8Riv-7vR$_r$4E%8b!U%$pzxZR$a=Jef(v8Jq7 z=x6!TY)Xq@b%5@NI?J`fQEQvx#OfJ(wm!`({YXDV_pGTiy?bt&d|hwa{AurswiIYE zkyKs`(O5|hEh*#pP2wP!1QA;FimJ++xk3pta6>pOl%v@ly`Kmn6p0Y-9SZ@Z?Gpmx zKKI3&ab+9e+oD={JORv#O)c?j=wSW*@fb?fi-c5(YG+F4Ar3_-gJ0thjN z>rA8Lz`kAj!m%MjO+{=Y|ND93z^?#p~>HI&1zS!(z=NA%KSwYCw!e2Q*b z*H+W>v8E{L#)Uj^7TUz?qO$S1Y2#v5)lHtg2{X4$MqnYqGPIz%?|uF{{mwz BU%UVS literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/union.bin b/crates/krishiv-plan-codec/tests/golden/v5/union.bin new file mode 100644 index 0000000000000000000000000000000000000000..32f0afadd3ea26360418e4dff225fd2f263f5917 GIT binary patch literal 870 zcmbtS!EVz)5Zztd@n(`xXVa>##lgx6St^mj1qDH?oDkxI;fMs=*lDb|cJOXgxaERW zZXEaket=KlfW#60qlf-Om~~R53OF&VdFz?knf04l?szP?VDO%CS3HMj&`}Fn&0kHV zmU(#!Zjw&rGS@)0@W4JRlT>hMC|zX}E&AZ-a-J6AK#m>@w3XPk?N z;sLnWi#N+uiv!qK_eO)k$#ODJ^*|1yyVs#T(c4mQsqC~!lf#A7vwzX9Nsl{xwX(ZE ze;Bzs3fvP1Ed1ZliaxTx5jOO`qh6pMqZ<3=dc8IT&3WVmVZ9q@AlU{kBdkUPrfG)o zr<-}MsbR!4if@~2vP@l*lb~)IP5x-~=5FGf#v8HgwaHUx|L=e#))+Ryewt3Tp!w!g zP6d5C1aRW)bfyDJhd6+^EVaPLs%4hr`y%YVH5M|QWrYsUr3!H}WtApRd``XfH_Kmr zSYLhp^10D>Wg+u-O2?fWBa4H`hi0aVQWwiS_i3M!5##KbKBloP*qz RJG2+|z>-B`+!nod{{p@Ne6#=n literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin b/crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin new file mode 100644 index 0000000000000000000000000000000000000000..c9a55711561e6751d820d9fdadc9e50e00a4abb3 GIT binary patch literal 1790 zcmcgt&2G~`5Z+lk$!42UcU!<%94wzuBrZ82AE4@~LR>IC7PYI|Nvs;XvNuXW98jTB zxpCkDcmW=P0}>C@1Md)Ky=j!v@{=O5t66z>Jf3g9of)%+V|%BRpHS8k&-hcmTph~t zV0T|CnH2-xij%$^Wr~9m4!2(vaUvLRR;o<*l?Zr46$eQ!?ugsmJ}iq^tc+!u42w~f z+^<#M5(5F!qpabB*Je~eg`2j&-L*K#F6U73Mw;hI*<+jnb!->HlvzTny}b@0gyAH_ zp00!d?$%KpTpe2a~_$c@my^AX&9n>z>w!S_#=JhueZZcqi=DF=G%^;d2H)bhodAM%$oM(CGBA>uF$U1U<+%o7$agMSMIgvSha` z##QhZ!i-T52s&Pv1|9m;aZ5kzp{Zc#T$4#rv|dwbKpY_mT7y`cf{^(j30{XoW*HDW zzmNXbm3LvBgQf30qclfZiInqoPm(e-yW=*KQL zz{yw3ee-_u{mbW1&7dW7nZ2x3w0z+Y5_w&hx6>*wR6fcw7Xlz5Wpo3!AhHF$w;ru^ YS2;XH4SJiYZg$tWCG%JxTX25-11?3X4gdfE literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin b/crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin new file mode 100644 index 0000000000000000000000000000000000000000..f8a20245a49e3d9fb349b5b10d66cb20427aadae GIT binary patch literal 860 zcmbtSv5wO~5ZzgClFeK$@of%aRuq;iWC;nSLRY{k6+$!^N+ejtPGZH_k-gCY6%EqS zQSbr$0G~jC#K&C0KZIH5Aa{a}UA^a>H#@VM@yL!Z^f|{5T=vv6e2Rz7+SIF;Ghl0>*HxM+iQUH5`OK;aJGNS7rMjhVqFC2eszzs~&X(1t$?mo;eIP-B z(Oh;y(0$1jG-UL{=dp)iVnk@Llb2;yPb6~Sz6jtXm!5KN{c6A%le8HVm&_QDbxiPy zeKuqx^0`|+aX<(Obc^36mTW-YgElsF@-Dp1Vab>d1;0T9Xi65!NIOALOW`fvB^2O3>fXNs|QdL&bh+2~u zG*yh>_om8(%bdD;W|}ZxlvSOjj{^z*_PZCn|G546?dz9r)H9_i-ZnNlyf&YSjsxuH uOwMJRlFf_Q!Q29;0W<)X5(FmVW{2=xufY literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/write_files.bin b/crates/krishiv-plan-codec/tests/golden/v5/write_files.bin new file mode 100644 index 0000000000000000000000000000000000000000..0e6927a4e89736bd79b3c5c33d9214f733087356 GIT binary patch literal 748 zcmbtSy>8S%5Z+na=e@~=ID0`Di^94JS)#cLrGO|ELNpjkBv{41I}_uL>^%iK8l=!s z@Bq93k3fOM!=&IH%#0n7BIwxF_u2FN?bxe=NG2B&KS&koXZRF%yED_CzC1C`)OL+w zQLfCkb_h;K=)bW=sTEE-*H$M^FEDWSv~2WkeGC1w))sp2+_dG!ZoBf{5K49wXfTOX zK?=!KNevwt)AV`fA(#viIt;3&Dchw&2{Q6yIINVXIh?$j2_Y11A^a;A0?0ZbM8uJ} zA@(T7SrKteNC}iB?+`07qv+01o3eW>A9^@4j%VsJ+#men@?j!hixjd?@?Wwn@`G?e zI3lq9qVIc_;e6_h>F{KlkmQ6mFF7z1ADv}CU#| Vec { diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index 62958449e..b9f9fd449 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -282,11 +282,10 @@ fn an_offset_limit_needs_a_v3_reader() { fn a_window_needs_a_v4_reader() { // A `Window` node sets `min_reader_version` to 4: an older reader has // no such node. Plans without one stay readable by older readers. - assert_eq!(FORMAT_VERSION, 4); assert_eq!(message("window_ranking").min_reader_version, 4); assert_eq!(message("window_aggregates").min_reader_version, 4); for (name, plan) in cases() { - if !name.starts_with("window") { + if !name.starts_with("window") && !name.starts_with("expand") { let m = to_message(&plan).unwrap(); assert!(m.min_reader_version <= 3, "{name} needs a newer reader"); } @@ -337,6 +336,75 @@ fn a_corrupt_window_is_refused() { ); } +#[test] +fn an_expand_needs_a_v5_reader() { + // An `Expand` node (and the `NullOf` item inside it) sets + // `min_reader_version` to 5: an older reader has neither. Plans + // without one keep the older minimum, and a v4 reader still reads them. + assert_eq!(FORMAT_VERSION, 5); + for name in ["expand_rollup", "expand_single", "expand_subquery"] { + assert_eq!(message(name).min_reader_version, 5, "{name}"); + } + for (name, plan) in cases() { + if !name.starts_with("expand") { + let m = to_message(&plan).unwrap(); + assert!(m.min_reader_version <= 4, "{name} needs a newer reader"); + } + } +} + +#[test] +fn an_expand_plan_survives_the_wire_with_every_field() { + for name in ["expand_rollup", "expand_single", "expand_subquery"] { + let (_, plan) = cases().into_iter().find(|(n, _)| n == name).unwrap(); + assert_eq!(round_trip(&plan), plan, "{name}"); + } +} + +#[test] +fn a_loose_null_of_item_still_needs_a_v5_reader() { + let plan = plan_of(PhysNode::Project { + input: Box::new(PhysNode::OneRow), + layout: Vec::new(), + items: vec![krishiv_exec::physical::ProjectItemPlan::NullOf { + index: 0, + name: "x".to_owned(), + }], + nullable: None, + subqueries: Vec::new(), + }); + assert_eq!(to_message(&plan).unwrap().min_reader_version, 5); + assert_eq!(round_trip(&plan), plan); +} + +#[test] +fn a_corrupt_expand_is_refused() { + use pb::node::Kind; + let edit = |f: &dyn Fn(&mut pb::ExpandNode)| { + let mut m = message("expand_rollup"); + for node in &mut m.nodes { + if let Some(Kind::Expand(e)) = &mut node.kind { + f(e); + } + } + m + }; + corrupt( + &edit(&|e| e.projections.clear()), + "an expand without a projection", + ); + corrupt( + &edit(&|e| { + e.nullable.pop(); + }), + "do not match the nullability", + ); + corrupt( + &edit(&|e| e.projections[0].items[0].kind = None), + "a project item kind", + ); +} + #[test] fn a_newer_writer_that_stays_compatible_is_read() { // N reads N+1 when the writer only added fields: unknown fields (here diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 5945829d1..dfa056764 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -19,7 +19,7 @@ syntax = "proto3"; package krishiv.plan.v1; message PhysicalPlan { - // The writer's format version (4). + // The writer's format version (5). uint32 format_version = 1; // The oldest reader version that understands every field in use. uint32 min_reader_version = 2; @@ -69,6 +69,9 @@ message Node { ExchangeNode exchange = 14; // v4: a plan with a window needs a v4 reader. WindowNode window = 15; + // v5: a plan with an expand (ROLLUP, CUBE, GROUPING SETS) needs a v5 + // reader. + ExpandNode expand = 16; } } @@ -98,6 +101,8 @@ message ProjectItem { oneof kind { ProjectColumn column = 1; ProjectExpr expr = 2; + // v5: only inside an expand; a plan holding one needs a v5 reader. + ProjectNullOf null_of = 3; } } @@ -111,6 +116,12 @@ message ProjectExpr { uint32 expr = 2; } +// v5: a column of nulls in the Arrow type of an input column. +message ProjectNullOf { + uint64 index = 1; + string name = 2; +} + message SubqueryPlan { uint64 id = 1; SubqueryValueKind kind = 2; @@ -139,6 +150,21 @@ message ReplicateRowsNode { repeated SubqueryPlan subqueries = 4; } +// v5: each input row once per projection (ROLLUP, CUBE, GROUPING SETS). +message ExpandNode { + uint32 input = 1; + repeated uint64 layout = 2; + // The projections: the same output columns, each computed one way. + repeated ExpandProjection projections = 3; + // Whether each output column is nullable. + repeated bool nullable = 4; + repeated SubqueryPlan subqueries = 5; +} + +message ExpandProjection { + repeated ProjectItem items = 1; +} + message SortKey { uint64 column = 1; bool descending = 2; From d812aa866a543a4ebf749eb125c5a3e700ddcccb Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:19:36 +0200 Subject: [PATCH 7/8] Fix review threads for grouping analytics and lowering caps Preserve DISTINCT and FILTER when retyping rewritten aggregates. Assert the full 129-row distribution for the 64-expression rollup. Count nested GROUPING SETS against the cap while accumulating. Assert exact lowered plans for nested GROUPING SETS. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-plan/src/analyzer/analytics.rs | 11 ++++- .../tests/grouping_analytics_api.rs | 39 ++++++++++++--- crates/krishiv-sql/src/lower.rs | 11 ++++- crates/krishiv-sql/tests/lower.rs | 49 ++++++++++++++++--- 4 files changed, 92 insertions(+), 18 deletions(-) diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs index b2af3a16f..f23975334 100644 --- a/crates/krishiv-plan/src/analyzer/analytics.rs +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -673,7 +673,16 @@ impl Resolver<'_> { let origin = e.origin; match &e.node { ExprNode::Resolved(ResolvedExpr::Function(f)) if !has_subquery(&e) => { - self.call(f.id.name(), f.args.clone(), env, origin) + let call = Arc::new( + Expr::new(ExprNode::Function(FunctionCall { + name: f.id.name().to_owned(), + args: f.args.clone(), + distinct: f.distinct, + filter: f.filter.clone(), + })) + .with_origin(origin), + ); + Ok(self.bind(&call, &Self::cx(&[], &[], env, "Join"))?.expr) } ExprNode::Resolved(ResolvedExpr::Function(f)) => { // A call over a subquery keeps its type, nullable. diff --git a/crates/krishiv-session/tests/grouping_analytics_api.rs b/crates/krishiv-session/tests/grouping_analytics_api.rs index 0259df485..e8f69e87c 100644 --- a/crates/krishiv-session/tests/grouping_analytics_api.rs +++ b/crates/krishiv-session/tests/grouping_analytics_api.rs @@ -372,14 +372,13 @@ fn grouping_reads_the_top_bit_of_a_64_expression_id() -> R { keys.join(", ") ); let got = rows(&s, &query)?; - assert!( - got.iter().any(|r| r == "1\t1"), - "the grand total (id -1) reads both bits as 1, got {got:?}" - ); - assert!( - !got.iter().any(|r| r.contains('-')), - "no negative bits, got {got:?}" - ); + // 65 rollup sets over `a`: the full set holds both expressions + // (`0\t0`, one row per distinct `a`), sets 63..1 hold the first but + // not the last (`0\t1`), and the empty set holds neither (`1\t1`). + let mut expected = vec!["0\t0".to_owned(), "0\t0".to_owned()]; + expected.extend(std::iter::repeat_n("0\t1".to_owned(), 126)); + expected.push("1\t1".to_owned()); + assert_eq!(got, expected); Ok(()) } @@ -413,3 +412,27 @@ fn grouping_calls_inside_aggregate_arguments_resolve() -> R { ); Ok(()) } + +#[test] +fn grouping_retype_keeps_distinct_and_filter() -> R { + let s = session()?; + // `grouping(a)` resolves inside the aggregate and retypes it: `DISTINCT` + // stays, so the grand total sums one `1`, not three. + assert_eq!( + rows( + &s, + "SELECT a, sum(DISTINCT grouping(a)) FROM t GROUP BY ROLLUP(a)" + )?, + sorted(&["1\t0", "2\t0", "NULL\t1"]) + ); + // `FILTER` stays too: the engine rejects filtered aggregates, so keeping + // the filter fails instead of silently summing to 60. + assert_eq!( + error_class( + &s, + "SELECT a, sum(v) FILTER (WHERE grouping(a) = 0) FROM t GROUP BY ROLLUP(a)" + )?, + ErrorClass::KrishivUnsupportedFeature + ); + Ok(()) +} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index 9416cc980..3dfe86607 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -649,8 +649,15 @@ impl Lowerer { } GroupItemKind::GroupingSets(items) => { let mut sets = Vec::new(); - for item in items { - sets.extend(self.deeper(|l| l.item_sets(item, &mut *written))?); + let span = item.span; + for inner in items { + sets.extend(self.deeper(|l| l.item_sets(inner, &mut *written))?); + if sets.len() > MAX_GROUPING_SETS { + return Err(unsupported( + &format!("more than {MAX_GROUPING_SETS} grouping sets"), + span, + )); + } } sets } diff --git a/crates/krishiv-sql/tests/lower.rs b/crates/krishiv-sql/tests/lower.rs index 013f87160..5d8aa1a42 100644 --- a/crates/krishiv-sql/tests/lower.rs +++ b/crates/krishiv-sql/tests/lower.rs @@ -202,14 +202,49 @@ fn nested_grouping_sets_lower_on_a_small_stack() { "GROUPING SETS (".repeat(900), ")".repeat(900) ); - assert!(lower(&sql).is_ok(), "nested GROUPING SETS lowers"); + // One set `((a))`: every level unwraps to the single inner set. + let expected = "Aggregate [a] GROUPING SETS ((a)), [a]\n+- Read t\n"; + assert_eq!( + lower(&sql) + .expect("nested GROUPING SETS lowers") + .root() + .to_string(), + expected + ); // A single level nests the same way and keeps its sets. - let plan = lower("SELECT a FROM t GROUP BY GROUPING SETS (GROUPING SETS ((a)))") - .expect("nested sets lower"); - assert!( - plan.root().to_string().contains("GROUPING SETS"), - "keeps grouping sets: {}", - plan.root() + assert_eq!( + lower("SELECT a FROM t GROUP BY GROUPING SETS (GROUPING SETS ((a)))") + .expect("nested sets lower") + .root() + .to_string(), + expected ); }); } + +#[test] +fn nested_grouping_sets_cap_counts_inside_the_loop() { + // Two 16-element cubes hold 65,536 sets each: the union passes the + // 65,536 cap on the second item, at the outer `GROUPING SETS`, before + // the cross-product check runs. + let a: Vec = (0..16).map(|i| format!("a{i}")).collect(); + let b: Vec = (0..16).map(|i| format!("b{i}")).collect(); + let sql = format!( + "SELECT a FROM t GROUP BY GROUPING SETS (CUBE ({}), CUBE ({}))", + a.join(", "), + b.join(", ") + ); + let e = lower(&sql).expect_err("two 16-cubes exceed the grouping-sets cap"); + assert!( + e.error + .to_string() + .contains("more than 65536 grouping sets"), + "cap message, got {}", + e.error + ); + assert_eq!( + e.span.start, + sql.find("GROUPING SETS").expect("the item"), + "the cap cites the GROUPING SETS item, not GROUP BY" + ); +} From e05fcdac5a0ecc199f242a894df9298b0f0c5605 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 24 Oct 2026 06:36:33 +0200 Subject: [PATCH 8/8] test(xtask): make the engine-raised sanity check structural The engine-raised error-class set grows with every implemented feature; the fixed one-tenth bound tripped once grouping sets added their classes (195 of 1,929). Bound the set at half the catalog, which still catches wholesale counting, and check on this repo that the comment and test filter drops real references. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- xtask/tests/registry_coverage.rs | 37 +++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/xtask/tests/registry_coverage.rs b/xtask/tests/registry_coverage.rs index 7de133d1d..36d40afe3 100644 --- a/xtask/tests/registry_coverage.rs +++ b/xtask/tests/registry_coverage.rs @@ -141,5 +141,40 @@ fn engine_raised_set_of_this_repo_is_a_small_subset_of_the_catalog() { let catalog = parse_error_classes(&json).unwrap(); assert!(catalog.len() > 1_900, "full Spark 4.2.0 catalog is loaded"); let in_use = raised(&catalog, &collect_engine(root)); - assert!(in_use.len() < catalog.len() / 10, "{}", in_use.len()); + // The engine-raised set grows with every implemented feature, so the + // bound is structural rather than a budget: a scan that counted classes + // wholesale (a generated registry or the catalog itself slipping into + // engine text) would reach the whole catalog. + assert!( + !in_use.is_empty() && in_use.len() < catalog.len() / 2, + "{} of {}", + in_use.len(), + catalog.len() + ); + // The comment and test filter removes real references in this repo: + // the same sources read whole name strictly more classes. + let mut raw = Vec::new(); + let mut dirs = vec![root.join("crates")]; + while let Some(dir) = dirs.pop() { + for entry in std::fs::read_dir(&dir).unwrap().flatten() { + let path = entry.path(); + if path.is_dir() { + dirs.push(path); + } else if path.extension().is_some_and(|e| e == "rs") + && path.components().any(|c| c.as_os_str() == "src") + { + raw.push(std::fs::read_to_string(&path).unwrap()); + } + } + } + let in_raw = |c: &ClassEntry| raised(std::slice::from_ref(c), &raw).len() == 1; + assert!( + in_use.iter().all(in_raw), + "engine text is a subset of the raw sources" + ); + let filtered_out = catalog + .iter() + .find(|c| !in_use.contains(c) && in_raw(c)) + .map(|c| c.name.clone()); + assert!(filtered_out.is_some(), "the filter drops no reference"); }