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-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index f38179085..157443739 100644 --- a/crates/krishiv-exec/src/lib.rs +++ b/crates/krishiv-exec/src/lib.rs @@ -115,10 +115,11 @@ pub use operator::{ FlushOutput, GlobalState, LocalState, Operator, OperatorOutput, Sink, SinkStatus, Source, }; pub use operators::{ - DEFAULT_MAX_MERGE_FAN_IN, FileScan, FileScanMetrics, FileScanSource, FileSelections, FileWrite, - FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, MemoryScan, - Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RowStageSettings, RuntimeFilters, - ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, SortMetrics, TopN, UnionAlign, + DEFAULT_MAX_MERGE_FAN_IN, Expand, FileScan, FileScanMetrics, FileScanSource, FileSelections, + FileWrite, FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, + MemoryScan, NullOf, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, + RowStageSettings, RuntimeFilters, ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, + SortMetrics, TopN, UnionAlign, }; pub use pipeline::{ExecNode, PipelineBuilder}; pub use runner::{QueryRunner, Scheduler}; diff --git a/crates/krishiv-exec/src/operators/expand.rs b/crates/krishiv-exec/src/operators/expand.rs new file mode 100644 index 000000000..aedb7a718 --- /dev/null +++ b/crates/krishiv-exec/src/operators/expand.rs @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Each row once per projection: the first step of `ROLLUP`, `CUBE` and +//! `GROUPING SETS` (Spark's `Expand`). + +use std::sync::Arc; + +use arrow_array::{ArrayRef, RecordBatch, new_null_array}; +use arrow_schema::{DataType, SchemaRef}; +use krishiv_types::batch::Batch; + +use super::{Project, ScalarExpr, local_mut, no_state}; +use crate::context::TaskContext; +use crate::error::{Result, internal}; +use crate::operator::{LocalState, Operator, OperatorOutput}; + +/// Emits, for every input batch, one batch per projection: the batch +/// projected by it. The projections share one output schema. Rows keep their +/// weights in every copy. Output is produced one projection per +/// [`Operator::execute`] call (`HaveMoreOutput` until the last), so only +/// one copy of an input batch exists at a time, and the plain columns of +/// each copy share the input's buffers. +#[derive(Debug)] +pub struct Expand { + projects: Vec>, + 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 a5aa39e67..4ebc36e92 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/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 3d7efcf80..b84ac9f2c 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -48,8 +48,8 @@ use crate::join::{ }; use crate::operator::{Operator, Sink, Source}; use crate::operators::{ - FileWrite, FileWriteSink, Filter, FilterPushdown, Predicate, Project, ProjectItem, - PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, + Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project, + ProjectItem, PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, }; use crate::pipeline::{ExecNode, PipelineBuilder}; use crate::planner::TableProviderResolver; @@ -296,6 +296,15 @@ impl Instantiation<'_> { } => self.with_subqueries(input, subqueries, |input, v| { replicate_op(input, layout, count, v) }), + PhysNode::Expand { + input, + layout, + projections, + nullable, + subqueries, + } => self.with_subqueries(input, subqueries, |input, v| { + expand_op(input, layout, projections, nullable, v) + }), PhysNode::Limit { input, offset, @@ -586,6 +595,9 @@ pub(crate) fn project_op( name.clone(), Arc::new(Compiled(compile_one(expr, layout, schema, v)?)), ), + ProjectItemPlan::NullOf { index, name } => { + ProjectItem::expr(name.clone(), Arc::new(NullOf::new(schema, *index)?)) + } }) }) .collect::>>()?; @@ -595,6 +607,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 a717f823d..f82c33d70 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -56,8 +56,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, - 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}; @@ -213,6 +213,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. @@ -333,6 +347,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`. @@ -444,6 +466,7 @@ plan_data!(enum PhysNode { Project { input, layout, items, nullable, subqueries }, Limit { input, offset, limit }, ReplicateRows { input, layout, count, subqueries }, + Expand { input, layout, projections, nullable, subqueries }, Sort { input, keys }, TopN { input, keys, limit }, Window { input, window }, @@ -458,7 +481,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, @@ -487,6 +510,9 @@ impl PhysNode { | Self::Project { input, subqueries, .. } + | Self::Expand { + input, subqueries, .. + } | Self::ReplicateRows { input, subqueries, .. } => std::iter::once(input.as_ref()) @@ -515,6 +541,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 6d3f36d49..fa3259dec 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -41,8 +41,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: @@ -363,6 +363,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::Window { input, exprs, @@ -784,6 +792,95 @@ 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 b4503247c..5f265997a 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())); @@ -199,6 +213,29 @@ 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, @@ -332,6 +369,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 b1cddaf41..2484702a4 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -241,6 +241,110 @@ 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, + output, + } => { + let est = Est::clone(&self.plan(input)); + let rows = est.rows * to_f64(u64::try_from(projections.len()).unwrap_or(u64::MAX)); + 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 literals: 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 !literals.contains(&l) { + literals.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 !literals.is_empty() { + mixed = true; + } + } + None => other = true, + }, + _ => other = true, + } + } + let nulls = (null_rows / rows).clamp(0.0, 1.0); + 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 { + 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_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_literals + }; + base.clamp(1.0, rows) + } else if literals.is_empty() { + // Every projection is NULL. + 1.0 + } else { + 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)) + } 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(); Est::unknown(rows, ids()) diff --git a/crates/krishiv-optimizer/tests/it/rules.rs b/crates/krishiv-optimizer/tests/it/rules.rs index d0597e11b..1aeb3f0b7 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; @@ -398,6 +399,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 +461,65 @@ 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_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 + ); +} + #[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 ceb389a88..057651766 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); @@ -755,6 +777,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-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 43ad12dd6..9c90308fd 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -243,6 +243,30 @@ impl Encoder { input: self.node_index(input)?, }), PhysNode::Exchange(exchange) => Kind::Exchange(self.exchange(exchange)?), + 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) }) } @@ -276,6 +300,14 @@ impl Encoder { name: name.clone(), expr: self.expr(expr)?, }), + 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 000000000..35a390601 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/aggregate_final.bin b/crates/krishiv-plan-codec/tests/golden/v5/aggregate_final.bin new file mode 100644 index 000000000..2c066dd7e Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/aggregate_final.bin differ 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 000000000..44186479a Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin differ 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 000000000..87676f17b Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_range.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_range.bin new file mode 100644 index 000000000..de8148094 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/exchange_range.bin differ 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 000000000..c4f0eabdd Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin differ 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 000000000..b9c7afa74 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/expand_rollup.bin b/crates/krishiv-plan-codec/tests/golden/v5/expand_rollup.bin new file mode 100644 index 000000000..4e73e9063 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/expand_rollup.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/expand_single.bin b/crates/krishiv-plan-codec/tests/golden/v5/expand_single.bin new file mode 100644 index 000000000..33a252bd6 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/expand_single.bin differ 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 000000000..1f9e9d7f8 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/expressions.bin b/crates/krishiv-plan-codec/tests/golden/v5/expressions.bin new file mode 100644 index 000000000..76e79c742 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/expressions.bin differ 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 000000000..50724761f Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/filter.bin b/crates/krishiv-plan-codec/tests/golden/v5/filter.bin new file mode 100644 index 000000000..56b679721 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/filter.bin differ 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 000000000..ddfde9004 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/hash_join_anti.bin b/crates/krishiv-plan-codec/tests/golden/v5/hash_join_anti.bin new file mode 100644 index 000000000..4722f9e39 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/hash_join_anti.bin differ 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 000000000..048d0341b Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/limit.bin b/crates/krishiv-plan-codec/tests/golden/v5/limit.bin new file mode 100644 index 000000000..e40801d41 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/limit.bin differ 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 000000000..d55e972c8 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin differ 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 000000000..0f49b3f68 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/one_row.bin b/crates/krishiv-plan-codec/tests/golden/v5/one_row.bin new file mode 100644 index 000000000..b72ea66cf Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/one_row.bin differ 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 000000000..cb7da8071 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/project.bin differ 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 000000000..91115ce8a Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/replicate_rows.bin b/crates/krishiv-plan-codec/tests/golden/v5/replicate_rows.bin new file mode 100644 index 000000000..8aec8a9b2 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/replicate_rows.bin differ 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 000000000..8d48e76e5 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/scan.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/sort.bin b/crates/krishiv-plan-codec/tests/golden/v5/sort.bin new file mode 100644 index 000000000..04264c411 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/sort.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/v5/top_n.bin b/crates/krishiv-plan-codec/tests/golden/v5/top_n.bin new file mode 100644 index 000000000..004e79b81 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/top_n.bin differ 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 000000000..32f0afadd Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/union.bin differ 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 000000000..c9a557115 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin differ 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 000000000..f8a20245a Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin differ 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 000000000..0e6927a4e Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/v5/write_files.bin differ diff --git a/crates/krishiv-plan-codec/tests/ratchet.rs b/crates/krishiv-plan-codec/tests/ratchet.rs index a1d0306fc..a8f2c437c 100644 --- a/crates/krishiv-plan-codec/tests/ratchet.rs +++ b/crates/krishiv-plan-codec/tests/ratchet.rs @@ -21,9 +21,8 @@ use suites::{Suite, queries}; /// fails when a query on it plans, so its entry is removed in the change /// that makes it plan. const TPCDS_NOT_PLANNED: &[&str] = &[ - "q05", "q10", "q14", "q16", "q17", "q18", "q22", "q30", "q32", "q35", "q39", "q41", "q45", - "q50", "q51", "q62", "q67", "q70", "q72", "q77", "q80", "q86", "q92", "q94", "q95", "q97", - "q99", + "q10", "q16", "q17", "q30", "q32", "q35", "q39", "q41", "q45", "q50", "q51", "q62", "q70", + "q72", "q86", "q92", "q94", "q95", "q97", "q99", ]; fn assert_round_trips(name: &str, plan: &PhysicalPlan) -> Vec { diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index acf090030..b9f9fd449 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -40,6 +40,7 @@ variants!( Project, Limit, ReplicateRows, + Expand, Sort, TopN, Window, @@ -59,6 +60,7 @@ fn variant_of(node: &PhysNode) -> Variant { PhysNode::Project { .. } => Variant::Project, PhysNode::Limit { .. } => Variant::Limit, PhysNode::ReplicateRows { .. } => Variant::ReplicateRows, + PhysNode::Expand { .. } => Variant::Expand, PhysNode::Sort { .. } => Variant::Sort, PhysNode::TopN { .. } => Variant::TopN, PhysNode::Window { .. } => Variant::Window, @@ -280,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"); } @@ -335,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/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs new file mode 100644 index 000000000..e8e9d29b3 --- /dev/null +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -0,0 +1,855 @@ +// 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 = 64; + +/// 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 { + 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, + /// Whether the expansion numbers its sets (some two are equal). + numbered: bool, +} + +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() + } + + /// `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 { + 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) + } +} + +/// `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))) +} + +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, + win: 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, zone: Option) -> String { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Column(c)) => format!("{}#{}", c.name, c.id.0), + _ => pretty(e, zone), + } +} + +/// 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), + win: kids.iter().find_map(|k| k.win), + 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 { + // 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 = 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, &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, + 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)), + ); + } + 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(), self.ctx.conf.session_zone), + }; + 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, + numbered, + 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 => { + // `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() => { + 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) => { + 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. + 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 zone = self.ctx.conf.session_zone; + 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, zone)) + .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(_), Some(Some(j))) = (f.args.first(), positions.first()) else { + return Err(AnalysisError::new( + GroupingColumnMismatch { + grouping: f.args.first().map(|a| with_id(a, zone)).unwrap_or_default(), + grouping_columns: an + .exprs + .iter() + .map(|g| format!("{}#{}", g.attr.col.name, g.attr.col.id.0)) + .collect::>() + .join(","), + }, + Origin::UNKNOWN, + )); + }; + // 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; + 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![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, + )?; + 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 = 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) + } + } + } +} 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 c663d8c8a..f1abcd22c 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}; @@ -51,6 +52,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, @@ -317,7 +322,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) @@ -803,6 +808,9 @@ impl Resolver<'_> { { return Err(super::windows::window_in_filter(filter, 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_op(&f.name, n) { 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 1422f2b74..fd3051de8 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 f48cc5d70..10497d4d6 100644 --- a/crates/krishiv-plan/src/analyzer/ops.rs +++ b/crates/krishiv-plan/src/analyzer/ops.rs @@ -15,6 +15,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; @@ -260,7 +261,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); } if items.iter().any(|b| b.win.is_some()) { return self.window_project(p.origin, c, &items, &[]); @@ -277,6 +278,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(); @@ -293,6 +308,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(); cx.windows = WindowUse::Allowed; self.bind(e, &cx)? }; @@ -737,6 +754,7 @@ 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(); if op == "Filter" { cx.windows = WindowUse::Having; } @@ -784,6 +802,12 @@ impl Resolver<'_> { Err(err) => return Err(err), } } + 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)?; + } + } for b in &mut bounds { b.expr = Arc::clone(&b.expr).transform_down(&mut |x: Arc| match &x.node { ExprNode::Resolved(ResolvedExpr::Function(f)) @@ -958,8 +982,14 @@ impl Resolver<'_> { info: AggInfo, pass: &[Attr], ) -> Result { + // 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(), }; @@ -981,6 +1011,7 @@ impl Resolver<'_> { } } }; + r.ids = super::resolve::hull(r.ids, defined); let info = Arc::new(info); // A column the select list dropped comes back from the aggregate // (rebuilt with its windows and projections above it). @@ -998,27 +1029,57 @@ 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, &items, env)?; + // No grouping (`HAVING` without `GROUP BY`): Spark's + // `MISSING_GROUP_BY` for an ungrouped column. + let global = group.is_empty(); + self.build_aggregate(p.origin, &c, group, &items, global, 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], + 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)?); } } // Spark resolves the select list first, then ordinals, then the @@ -1026,7 +1087,7 @@ impl Resolver<'_> { // or implicit (`SELECT 1 GROUP BY `1`` groups by the literal): the // aliases are the items' output names over the expressions that // compute them, searched after the input. - let aliases: Vec<(String, Arc)> = if conf.group_by_aliases { + let aliases: Vec<(String, Arc)> = if self.ctx.conf.group_by_aliases { items .iter() .map(|b| { @@ -1041,7 +1102,7 @@ impl Resolver<'_> { } else { Vec::new() }; - 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"); @@ -1050,11 +1111,8 @@ impl Resolver<'_> { } } let group: Vec = group.into_iter().flatten().collect(); - self.check_grouping(&group, &items)?; - // No grouping (`HAVING` without `GROUP BY`): Spark's - // `MISSING_GROUP_BY` for an ungrouped column. - let global = group.is_empty(); - self.build_aggregate(p.origin, &c, group, &items, global) + self.check_grouping(&group, items)?; + Ok(group) } /// Spark's `GROUP_BY_AGGREGATE` and @@ -1155,6 +1213,7 @@ impl Resolver<'_> { group: Vec, items: &[Bound], global: bool, + analytics: Option>, ) -> Result { let mut memo = HashMap::new(); let gkeys: Vec<(u128, ArgType)> = group @@ -1170,22 +1229,37 @@ 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, having: None, }; if base.len() == items.len() { @@ -1198,7 +1272,7 @@ impl Resolver<'_> { let mut top = Vec::new(); for b in items { if b.lateral.is_empty() && b.win.is_none() { - 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() @@ -1715,6 +1789,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)?; + } // Over an aggregate with windows, Spark's `ORDER BY` reads the // select list only (nothing resolves through the windows). let (bounds, below, extended) = if c.agg.as_ref().is_some_and(|i| !i.has_windows()) { diff --git a/crates/krishiv-plan/src/analyzer/resolve.rs b/crates/krishiv-plan/src/analyzer/resolve.rs index 9f554889d..e3f6f1e47 100644 --- a/crates/krishiv-plan/src/analyzer/resolve.rs +++ b/crates/krishiv-plan/src/analyzer/resolve.rs @@ -52,6 +52,9 @@ pub(super) struct AggInfo { /// 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>, /// With windows: the `HAVING` filtering the aggregate below them. pub(super) having: Option>, } @@ -220,7 +223,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, @@ -564,6 +567,7 @@ impl Resolver<'_> { aliases: &[], extra: None, agg_scope: None, + grouping: false, lenient: false, outer: true, windows: super::exprs::WindowUse::Other, @@ -638,6 +642,7 @@ impl Resolver<'_> { | PlanNode::CteRef(_) | PlanNode::Shift { .. } | PlanNode::ReplicateRows { .. } + | PlanNode::Expand { .. } | PlanNode::JoinHint { .. } => self.leaf_node(p, env), } } @@ -722,6 +727,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::JoinHint { input, strategy } => { let c = self.child(input, env)?; let node = PlanNode::JoinHint { diff --git a/crates/krishiv-plan/src/analyzer/setops.rs b/crates/krishiv-plan/src/analyzer/setops.rs index 93a5d74dd..af50faaf9 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 0b9b2551b..01b5ca52f 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -302,6 +302,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 b5f032a88..95533323a 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -795,6 +795,17 @@ impl<'k, 'a> Encoder<'k, 'a> { self.exprs(partition_by); self.orders(order_by); } + PlanNode::Expand { + projections, + output, + .. + } => { + self.u8(24); + 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 dbe5233a3..74bee1a1c 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` | @@ -82,6 +82,7 @@ //! | `PlanNode::Distinct` | U | - | - | //! | `PlanNode::SetOp` | U | - | - | //! | `PlanNode::ReplicateRows` | R | - | - | +//! | `PlanNode::Expand` | R | - | - | //! | `PlanNode::Window` | R | - | - | //! | `ExprNode::Column` | U | all | `col`, `df["a"]` | //! | `ExprNode::Literal` | U, R | all | `lit` | diff --git a/crates/krishiv-plan/src/plan.rs b/crates/krishiv-plan/src/plan.rs index 3914f57e5..7c2daab2e 100644 --- a/crates/krishiv-plan/src/plan.rs +++ b/crates/krishiv-plan/src/plan.rs @@ -267,6 +267,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`]. @@ -588,6 +601,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } + | PlanNode::Expand { input, .. } | PlanNode::Window { input, .. } | PlanNode::ReplicateRows { input, .. } => f(input), PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { @@ -621,6 +635,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } + | PlanNode::Expand { input, .. } | PlanNode::Window { input, .. } | PlanNode::ReplicateRows { input, .. } => f(input), PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { @@ -682,6 +697,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), PlanNode::Window { exprs, partition_by, @@ -742,6 +758,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); + } PlanNode::Window { exprs, partition_by, @@ -806,6 +825,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 0c8e03c93..35008652b 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 { @@ -1468,13 +1468,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 e88269d67..21682aa47 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, 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: @@ -42,6 +42,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", PlanNode::Window { .. } => "PlanNode::Window", }); plan.for_each_expr(&mut |e| expr_kinds(e, out)); @@ -157,6 +158,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), @@ -166,6 +174,11 @@ 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 52ccaf8cd..1540afd6a 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -334,7 +334,7 @@ impl Gen<'_> { fn plan(&mut self, depth: u8) -> LogicalPlan { let origin = self.origin(); - let kinds = if depth == 0 { 4 } else { 23 }; + let kinds = if depth == 0 { 4 } else { 24 }; let d = depth.saturating_sub(1); let node = match self.pick(kinds) { 0 => PlanNode::Read(self.source()), @@ -468,6 +468,11 @@ impl Gen<'_> { partition_by: self.exprs(), order_by: self.orders(2), }, + 23 => PlanNode::Expand { + input: self.input(d), + projections: self.many(1, 2, Gen::exprs), + output: self.many(0, 2, Gen::bound), + }, _ => PlanNode::WithCte { input: self.input(d), ctes: self.many(1, 2, |g| Cte { @@ -914,6 +919,25 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { | PlanNode::Limit { input, limit: e } | PlanNode::Offset { input, offset: 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 a6bdaa88e..699160fd2 100644 --- a/crates/krishiv-plan/src/walk.rs +++ b/crates/krishiv-plan/src/walk.rs @@ -156,6 +156,7 @@ pub(crate) fn all_resolved(root: Node<'_>) -> bool { | PlanNode::SubqueryAlias { .. } | PlanNode::WithCte { .. } | PlanNode::JoinHint { .. } + | PlanNode::Expand { .. } | PlanNode::Window { .. } | PlanNode::ReplicateRows { .. } => true, }, @@ -267,6 +268,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 e0255c177..aeef39dd4 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -313,9 +313,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, } } @@ -852,11 +916,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 { @@ -874,9 +953,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/grouping_analytics.rs b/crates/krishiv-session/tests/grouping_analytics.rs new file mode 100644 index 000000000..0920d7102 --- /dev/null +++ b/crates/krishiv-session/tests/grouping_analytics.rs @@ -0,0 +1,350 @@ +// 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 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) + .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()); + } + 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(), + 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..e8f69e87c --- /dev/null +++ b/crates/krishiv-session/tests/grouping_analytics_api.rs @@ -0,0 +1,438 @@ +// 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 + ); + // 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, &rollup(65))?, + 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(", ") + ) + }; + block_on(s.sql(&cube(16)))?; // Analyzer-only: 65,536 sets would still run. + 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 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)?; + // 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(()) +} + +#[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(()) +} + +#[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(()) +} + +#[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 e204a9d28..c7761a6cf 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -31,17 +31,20 @@ use krishiv_plan::{ use krishiv_types::DecimalType; use crate::ast::{ - Alias, BinaryOp, Expr, ExprKind, FrameBound, FrameBoundKind, FrameUnits, FromClause, - GroupItemKind, Ident, JoinConstraint, JoinKind, Literal, NamedWindow, NullTreatment, NumKind, - Query, Relation, RelationExt, Select, SelectItem, SetExpr, SetExprKind, SetOperator, - SetQuantifier, SortItem, Span, TableFactor, TableFactorKind, Truth, UnaryOp, WindowRef, - WindowSpec, + Alias, BinaryOp, Expr, ExprKind, FrameBound, FrameBoundKind, FrameUnits, FromClause, GroupBy, + GroupItem, GroupItemKind, GroupWith, GroupingSet, Ident, JoinConstraint, JoinKind, Literal, + NamedWindow, NullTreatment, NumKind, Query, Relation, RelationExt, Select, SelectItem, SetExpr, + SetExprKind, SetOperator, SetQuantifier, SortItem, Span, TableFactor, TableFactorKind, Truth, + UnaryOp, WindowRef, WindowSpec, }; use crate::validate::literal::{self, Number, TypedKind}; use crate::{ParseError, ParserOptions, parse_named_expression, 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`. /// /// Typed datetime literals (`TIMESTAMP '...'`, `TIMESTAMP_LTZ '...'`, @@ -589,22 +592,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, @@ -637,6 +632,136 @@ 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(); + 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 + } + 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)); diff --git a/crates/krishiv-sql/tests/lower.rs b/crates/krishiv-sql/tests/lower.rs index 33b469bba..5d8aa1a42 100644 --- a/crates/krishiv-sql/tests/lower.rs +++ b/crates/krishiv-sql/tests/lower.rs @@ -191,3 +191,60 @@ 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) + ); + // 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. + 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" + ); +} 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 ba0cd529d..3641abad0 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 => Project [row_number() OVER (ORDER BY a ASC NULLS FIRST)] | +- Read t | 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) 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; diff --git a/xtask/tests/registry_coverage.rs b/xtask/tests/registry_coverage.rs index 79417e335..3d4316336 100644 --- a/xtask/tests/registry_coverage.rs +++ b/xtask/tests/registry_coverage.rs @@ -146,4 +146,30 @@ fn engine_raised_set_of_this_repo_is_a_small_subset_of_the_catalog() { // with each function family (193 of the 1,929 classes after the bitwise // and number-format functions), so the bound is a quarter, not a tenth. assert!(in_use.len() < catalog.len() / 4, "{}", in_use.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"); }