From 127fcdd1768e7d68369df70d4462f57a174fb96c Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:20:12 +0200 Subject: [PATCH 1/9] wip: rollup/cube/grouping sets Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/lib.rs | 4 +- crates/krishiv-exec/src/operators/expand.rs | 144 ++++ crates/krishiv-exec/src/operators/mod.rs | 2 + .../krishiv-exec/src/physical/instantiate.rs | 28 +- crates/krishiv-exec/src/physical/mod.rs | 33 +- crates/krishiv-exec/src/planner/build.rs | 98 ++- crates/krishiv-optimizer/src/prune.rs | 46 ++ crates/krishiv-optimizer/src/stats.rs | 8 + crates/krishiv-plan/src/analyzer/analytics.rs | 692 ++++++++++++++++++ crates/krishiv-plan/src/analyzer/bind.rs | 4 + crates/krishiv-plan/src/analyzer/exprs.rs | 10 +- crates/krishiv-plan/src/analyzer/mod.rs | 1 + crates/krishiv-plan/src/analyzer/ops.rs | 127 +++- crates/krishiv-plan/src/analyzer/resolve.rs | 11 +- crates/krishiv-plan/src/analyzer/setops.rs | 2 +- crates/krishiv-plan/src/display.rs | 20 + crates/krishiv-plan/src/encode.rs | 11 + crates/krishiv-plan/src/lib.rs | 3 +- crates/krishiv-plan/src/plan.rs | 20 + crates/krishiv-plan/src/tests/analyzer.rs | 7 +- crates/krishiv-plan/src/tests/inventory.rs | 12 +- crates/krishiv-plan/src/tests/props.rs | 23 +- crates/krishiv-plan/src/walk.rs | 7 + crates/krishiv-session/src/dataframe.rs | 95 ++- .../krishiv-session/tests/scratch_rollup.rs | 84 +++ crates/krishiv-sql/src/lower.rs | 147 +++- 26 files changed, 1575 insertions(+), 64 deletions(-) create mode 100644 crates/krishiv-exec/src/operators/expand.rs create mode 100644 crates/krishiv-plan/src/analyzer/analytics.rs create mode 100644 crates/krishiv-session/tests/scratch_rollup.rs diff --git a/crates/krishiv-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index f6efa4c0..b9081da9 100644 --- a/crates/krishiv-exec/src/lib.rs +++ b/crates/krishiv-exec/src/lib.rs @@ -113,8 +113,8 @@ pub use operator::{ FlushOutput, GlobalState, LocalState, Operator, OperatorOutput, Sink, SinkStatus, Source, }; pub use operators::{ - DEFAULT_MAX_MERGE_FAN_IN, FileScan, FileScanMetrics, FileScanSource, FileSelections, FileWrite, - FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, MemoryScan, + DEFAULT_MAX_MERGE_FAN_IN, Expand, NullOf, FileScan, FileScanMetrics, FileScanSource, FileSelections, + FileWrite, FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, MemoryScan, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RowStageSettings, RuntimeFilters, ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, SortMetrics, TopN, UnionAlign, }; diff --git a/crates/krishiv-exec/src/operators/expand.rs b/crates/krishiv-exec/src/operators/expand.rs new file mode 100644 index 00000000..5f0a5c23 --- /dev/null +++ b/crates/krishiv-exec/src/operators/expand.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Each row once per projection: the first step of `ROLLUP`, `CUBE` and +//! `GROUPING SETS` (Spark's `Expand`). + +use std::sync::Arc; + +use arrow_array::{ArrayRef, RecordBatch, new_null_array}; +use arrow_schema::{DataType, SchemaRef}; +use krishiv_types::batch::Batch; + +use super::{Project, ScalarExpr, local_mut, no_state}; +use crate::context::TaskContext; +use crate::error::{Result, internal}; +use crate::operator::{LocalState, Operator, OperatorOutput}; + +/// Emits, for every input batch, one batch per projection: the batch +/// projected by it. The projections share one output schema. Rows keep their +/// weights in every copy. Output is produced one projection per +/// [`Operator::execute`] call (`HaveMoreOutput` until the last), so only +/// one copy of an input batch exists at a time, and the plain columns of +/// each copy share the input's buffers. +#[derive(Debug)] +pub struct Expand { + projects: Vec>, + schema: SchemaRef, +} + +impl Expand { + /// Expands batches by `projects`. + /// + /// # Errors + /// `INTERNAL_ERROR` if there is no projection, or two of them output + /// different columns. + pub fn try_new(projects: Vec>) -> Result> { + let first = projects + .first() + .ok_or_else(|| internal("an expansion without projections".to_owned()))?; + let schema = first.schema(); + for p in &projects { + let (a, b) = (p.schema(), &schema); + let same = a.fields().len() == b.fields().len() + && a.fields() + .iter() + .zip(b.fields()) + .all(|(x, y)| { + x.name() == y.name() + && x.data_type() == y.data_type() + && x.is_nullable() == y.is_nullable() + }); + if !same { + return Err(internal(format!( + "the projections of an expansion differ: {:?} and {:?}", + a.fields(), + b.fields() + ))); + } + } + Ok(Arc::new(Self { projects, schema })) + } +} + +/// The next projection to emit for the current input batch. +struct Cursor(usize); + +impl Operator for Expand { + fn name(&self) -> &'static str { + "Expand" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn local_state(&self, _ctx: &TaskContext) -> Result { + Ok(Box::new(Cursor(0))) + } + + fn execute(&self, input: &Batch, local: &mut LocalState) -> Result { + let cursor = local_mut::(local, "Expand")?; + if input.num_rows() == 0 { + cursor.0 = 0; + return Ok(OperatorOutput::NeedMoreInput(None)); + } + let project = self + .projects + .get(cursor.0) + .ok_or_else(|| internal("an expansion past its last projection".to_owned()))?; + let OperatorOutput::NeedMoreInput(Some(out)) = project.execute(input, &mut no_state())? + else { + return Err(internal("a projection did not output a batch".to_owned())); + }; + cursor.0 += 1; + if cursor.0 < self.projects.len() { + return Ok(OperatorOutput::HaveMoreOutput(out)); + } + cursor.0 = 0; + Ok(OperatorOutput::NeedMoreInput(Some(out))) + } +} + +/// A column of nulls, in the Arrow type of an input column. An expansion +/// nulls a grouping column in the carrier its other copies have: strings +/// reach an operator as `Utf8` from a registered table and as `Utf8View` +/// from a file. +#[derive(Debug)] +pub struct NullOf { + index: usize, + data_type: DataType, +} + +impl NullOf { + /// Nulls of input column `index` of batches of `input`. + /// + /// # Errors + /// `INTERNAL_ERROR` if the column is not in `input`. + pub fn new(input: &SchemaRef, index: usize) -> Result { + let field = input + .fields() + .get(index) + .ok_or_else(|| internal(format!("null-of index {index} out of range")))?; + Ok(Self { + index, + data_type: field.data_type().clone(), + }) + } +} + +impl ScalarExpr for NullOf { + fn data_type(&self) -> DataType { + self.data_type.clone() + } + + fn nullable(&self) -> bool { + true + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let column = batch + .columns() + .get(self.index) + .ok_or_else(|| internal(format!("null-of index {} out of range", self.index)))?; + Ok(new_null_array(column.data_type(), batch.num_rows())) + } +} diff --git a/crates/krishiv-exec/src/operators/mod.rs b/crates/krishiv-exec/src/operators/mod.rs index f8cf3b0a..da85c64c 100644 --- a/crates/krishiv-exec/src/operators/mod.rs +++ b/crates/krishiv-exec/src/operators/mod.rs @@ -25,6 +25,7 @@ //! - **`UNION`.** See [`union`] for how `UNION ALL` fan-in and the `UNION` //! (distinct) breaker are split between tasks. +mod expand; mod expr; mod file_scan; mod file_write; @@ -37,6 +38,7 @@ mod scan; mod sort; mod union; +pub use expand::{Expand, NullOf}; pub use expr::{FnPredicate, Predicate, ScalarExpr}; pub use file_scan::{FileScan, FileScanMetrics, FileScanSource, METADATA_COLUMN}; pub use file_write::{FileWrite, FileWriteSink}; diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index 266642cb..728028a0 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -30,7 +30,7 @@ use crate::join::{ }; use crate::operator::{Operator, Sink, Source}; use crate::operators::{ - FileWrite, FileWriteSink, Filter, FilterPushdown, Predicate, Project, ProjectItem, + Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, }; use crate::pipeline::{ExecNode, PipelineBuilder}; @@ -278,6 +278,15 @@ impl Instantiation<'_> { } => self.with_subqueries(input, subqueries, |input, v| { replicate_op(input, layout, count, v) }), + PhysNode::Expand { + input, + layout, + projections, + nullable, + subqueries, + } => self.with_subqueries(input, subqueries, |input, v| { + expand_op(input, layout, projections, nullable, v) + }), PhysNode::Limit { input, limit } => { let input = self.node(input)?; let op = crate::operators::Limit::new(Arc::clone(&input.schema), *limit); @@ -547,6 +556,9 @@ pub(crate) fn project_op( name.clone(), Arc::new(Compiled(compile_one(expr, layout, schema, v)?)), ), + ProjectItemPlan::NullOf { index, name } => { + ProjectItem::expr(name.clone(), Arc::new(NullOf::new(schema, *index)?)) + } }) }) .collect::>>()?; @@ -556,6 +568,20 @@ pub(crate) fn project_op( } } +pub(crate) fn expand_op( + schema: &SchemaRef, + layout: &[AttrId], + projections: &[Vec], + nullable: &[bool], + v: &Values<'_>, +) -> Result> { + let projects = projections + .iter() + .map(|items| project_op(schema, layout, items, Some(nullable), v)) + .collect::>>()?; + Ok(Expand::try_new(projects)?) +} + pub(crate) fn replicate_op( schema: &SchemaRef, layout: &[AttrId], diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index f4f1ce66..2ca8958d 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -51,7 +51,7 @@ pub(crate) use instantiate::node_ids; pub use instantiate::{SpillScopes, WorkerContext, instantiate}; pub(crate) use instantiate::{ Unbuilt, Values, aggregate_node, aggregate_spec, filter_op, hash_join, nested_loop_join, - project_op, replicate_op, scan_filter, sort_op, top_n_op, + expand_op, project_op, replicate_op, scan_filter, sort_op, top_n_op, }; use crate::aggregate::{AggregateExpr, AggregateOutput}; @@ -144,6 +144,20 @@ pub enum PhysNode { /// The subqueries `count` reads. subqueries: Vec, }, + /// Each row of `input` once per projection (`ROLLUP`, `CUBE`, + /// `GROUPING SETS`). + Expand { + /// The input. + input: Box, + /// The attribute of each input column. + layout: Vec, + /// The projections: the same output columns, each computed one way. + projections: Vec>, + /// Whether each output column is nullable. + nullable: Vec, + /// The subqueries the projections read. + subqueries: Vec, + }, /// `input` sorted by `keys` (columns of `input`). Sort { /// The input. @@ -252,6 +266,14 @@ pub enum ProjectItemPlan { /// The value. expr: Arc, }, + /// A column of nulls in the Arrow type of an input column (the cell an + /// expansion nulls). + NullOf { + /// The input column whose type it takes. + index: usize, + /// The output name. + name: String, + }, } /// Scalar or `EXISTS`. @@ -346,6 +368,7 @@ plan_data!(enum PhysNode { Project { input, layout, items, nullable, subqueries }, Limit { input, limit }, ReplicateRows { input, layout, count, subqueries }, + Expand { input, layout, projections, nullable, subqueries }, Sort { input, keys }, TopN { input, keys, limit }, Aggregate { input, aggregate }, @@ -358,7 +381,7 @@ plan_data!(struct ScanNode { id, scan, projection, filter, runtime_filters, row_ plan_data!(struct ScanFilter { layout, predicate, exact }); plan_data!(struct ScanRuntimeFilter { id, column, key_type }); plan_data!(struct JoinFilterPlan { id, key, options }); -plan_data!(enum ProjectItemPlan { Column { index, name }, Expr { name, expr } }); +plan_data!(enum ProjectItemPlan { Column { index, name }, Expr { name, expr }, NullOf { index, name } }); plan_data!( enum SubqueryValueKind { Scalar, @@ -386,6 +409,9 @@ impl PhysNode { | Self::Project { input, subqueries, .. } + | Self::Expand { + input, subqueries, .. + } | Self::ReplicateRows { input, subqueries, .. } => std::iter::once(input.as_ref()) @@ -412,6 +438,9 @@ impl PhysNode { | Self::Project { input, subqueries, .. } + | Self::Expand { + input, subqueries, .. + } | Self::ReplicateRows { input, subqueries, .. } => std::iter::once(input.as_mut()) diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index b7aadefb..e9fcc721 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -40,8 +40,8 @@ use crate::operator::{Operator, Source}; use crate::operators::{FilterPushdown, Limit, Project, ProjectItem, ScanPushdowns, SortKey, TopN}; use crate::physical::{ PhysNode, PhysicalPlan, ProjectItemPlan, ScanFilter, ScanId, ScanNode, ScanRuntimeFilter, - SubqueryId, SubqueryPlan, SubqueryValueKind, Values, filter_op, project_op, replicate_op, - scan_filter, sort_op, top_n_op, + SubqueryId, SubqueryPlan, SubqueryValueKind, Values, expand_op, filter_op, project_op, + replicate_op, scan_filter, sort_op, top_n_op, }; /// How far above the planner's first fresh attribute subquery ids start: @@ -348,6 +348,14 @@ impl<'a> Planner<'a> { let input = self.plan(input)?; self.replicate(input, count) } + PlanNode::Expand { + input, + projections, + output, + } => { + let input = self.plan(input)?; + self.expand(input, projections, output) + } PlanNode::OneRowRelation => Ok(one_row()), _ => { let text = p.to_string(); @@ -696,6 +704,92 @@ impl<'a> Planner<'a> { Ok(planned) } + /// Each row once per projection (`ROLLUP`, `CUBE`, `GROUPING SETS`): + /// plain columns pass through without a copy. + fn expand( + &mut self, + input: Planned, + projections: &[Vec>], + output: &[BoundColumn], + ) -> Result { + if projections.iter().any(|row| row.len() != output.len()) { + return Err(internal("an expansion row of the wrong width")); + } + // A column an expansion nulls keeps the Arrow type of the column + // another projection copies. + let copied: Vec> = (0..output.len()) + .map(|j| { + projections + .iter() + .find_map(|row| row.get(j).and_then(|e| column_of(e))) + .and_then(|c| input.column(c).ok()) + }) + .collect(); + let null_of = |j: usize, e: &Expr| { + let null = matches!(&e.node, ExprNode::Literal(Literal::Null(Some(_)))); + copied.get(j).copied().flatten().filter(|_| null) + }; + let computed: Vec<&Expr> = projections + .iter() + .flat_map(|row| row.iter().enumerate()) + .filter(|(j, e)| column_of(e).is_none() && null_of(*j, e).is_none()) + .map(|(_, e)| e.as_ref()) + .collect(); + let prepared = self.prepare(&computed)?; + let mut values = prepared.exprs.iter(); + let mut rows = Vec::with_capacity(projections.len()); + for row in projections { + let mut items = Vec::with_capacity(row.len()); + for (j, (e, out)) in row.iter().zip(output).enumerate() { + let name = out.name.clone(); + if let Some(c) = column_of(e) { + items.push(ProjectItemPlan::Column { + index: input.column(c)?, + name: Some(name), + }); + } else if let Some(index) = null_of(j, e) { + items.push(ProjectItemPlan::NullOf { index, name }); + } else { + let expr = values.next().ok_or_else(|| internal("expansion arity"))?; + items.push(ProjectItemPlan::Expr { + name, + expr: Arc::clone(expr), + }); + } + } + rows.push(items); + } + let nullable: Vec = output.iter().map(|c| c.nullable).collect(); + let layout: Vec = output.iter().map(|c| c.id).collect(); + let op = expand_op( + &input.schema, + &input.layout, + &rows, + &nullable, + &self.values(&prepared), + )?; + let names: Vec<&str> = output.iter().map(|c| c.name.as_str()).collect(); + let text = names.join(", "); + let mut planned = Self::attach( + input, + prepared, + op.schema(), + layout, + ("Expand", format!("{} projections [{text}]", projections.len())), + |input, layout, subqueries| PhysNode::Expand { + input, + layout, + projections: rows, + nullable, + subqueries, + }, + ); + // The rows are repeated with columns nulled: no runtime filter may + // reach a scan through the expansion. + planned.sites.clear(); + Ok(planned) + } + /// `SELECT exprs`: plain columns pass through without a copy. pub(crate) fn project(&mut self, input: Planned, exprs: &[Arc]) -> Result { let computed: Vec<&Expr> = exprs diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index cf56e0ef..7fb5f23a 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -76,6 +76,20 @@ impl Graph { grouping_refs(p, &mut need); child(input, Some(need)); } + PlanNode::Expand { + input, + projections, + output, + } => { + let mut need = Attrs::new(); + for j in expand_keep(output, req) { + projections + .iter() + .filter_map(|row| row.get(j)) + .for_each(|e| references_into(e, &mut need)); + } + child(input, Some(need)); + } PlanNode::SubqueryAlias { input, columns, .. } => { // A renaming alias keeps its input whole. child(input, columns.is_empty().then(|| req.clone())); @@ -197,6 +211,26 @@ impl Rewrite<'_> { replace(&rebuilt, PlanNode::Project { input, exprs: kept }) } } + PlanNode::Expand { + input, + projections, + output, + } => { + let kept = expand_keep(output, req); + if kept.len() == output.len() { + rebuilt + } else { + let pick = |row: &Vec>| -> Vec> { + kept.iter().filter_map(|&j| row.get(j).cloned()).collect() + }; + let node = PlanNode::Expand { + input: Arc::clone(input), + projections: projections.iter().map(pick).collect(), + output: kept.iter().filter_map(|&j| output.get(j).cloned()).collect(), + }; + replace(&rebuilt, node) + } + } PlanNode::Aggregate { input, grouping, @@ -304,6 +338,18 @@ fn keep(exprs: &[Arc], req: &Attrs) -> Vec> { kept } +/// The output columns of an expansion that are read (by position), keeping +/// at least one so the row count survives. +fn expand_keep(output: &[BoundColumn], req: &Attrs) -> Vec { + let mut kept: Vec = (0..output.len()) + .filter(|&j| req.contains(&output[j].id)) + .collect(); + if kept.is_empty() && !output.is_empty() { + kept.push(0); + } + kept +} + fn refs(exprs: &[Arc]) -> Attrs { let mut out = Attrs::new(); exprs.iter().for_each(|e| references_into(e, &mut out)); diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index fd00375b..d96d5abf 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -239,6 +239,14 @@ impl<'a> Estimator<'a> { // Each distinct row of an `INTERSECT ALL` / `EXCEPT ALL` count: // no count is known, the input's rows bound it from below. | PlanNode::ReplicateRows { input, .. } => Est::clone(&self.plan(input)), + // Every input row once per projection. + PlanNode::Expand { + input, projections, .. + } => { + let est = Est::clone(&self.plan(input)); + let rows = est.rows * to_f64(u64::try_from(projections.len()).unwrap_or(u64::MAX)); + est.with_rows(rows) + } PlanNode::Union { inputs, .. } => { let rows = inputs.iter().map(|i| self.plan(i).rows).sum(); Est::unknown(rows, ids()) diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs new file mode 100644 index 00000000..c6d34c79 --- /dev/null +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -0,0 +1,692 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Grouping analytics: `GROUP BY ROLLUP`, `CUBE` and `GROUPING SETS`, and the +//! `grouping` and `grouping_id` functions. +//! +//! The rewrite is Spark's (the `ResolveGroupingAnalytics` rule and the +//! `Expand` operator): every distinct grouping expression gets a nullable +//! copy `g`, `Expand` repeats each input row once per grouping set (`g` is +//! the expression for a set that has it and `NULL` for a set that has not) +//! with the set's `spark_grouping_id`, and an ordinary aggregate groups by +//! the copies and the id. Bit `n - 1 - j` of the id is 1 when the `j`-th +//! distinct grouping expression (in order of first appearance) is absent +//! from the set. Select items, `HAVING` and `ORDER BY` read the copies where +//! they name a grouping expression (aggregate calls keep reading the +//! input's own columns, which stay in the expansion), `grouping(e)` is bit +//! `e` of the id (a `TINYINT`) and `grouping_id(..)` is the id (a `BIGINT`). +//! +//! Sets that are equal as sets stay separate rows (`GROUPING SETS ((a), (a))` +//! has two rows per group): the expansion then also numbers its sets +//! (`_gen_grouping_pos`) and the aggregate groups by that too. A single +//! empty set over no grouping expression is a grand total, an aggregate +//! without `GROUP BY` (one row over empty input) in which `grouping_id()` is 0. +//! The virtual column `grouping__id` names the id. +//! +//! Exact behaviour, Spark 4.2: `tests/spark-golden/v4.2.0` (`group-analytics`, +//! `grouping_set`, `postgreSQL/groupingsets`). + +use std::collections::HashMap; +use std::sync::Arc; + +use krishiv_common::error::{ + GroupingColumnMismatch, GroupingIdColumnMismatch, GroupingSizeLimitExceeded, + UnsupportedGroupingExpression, WrongNumArgsWithoutSuggestion, +}; +use krishiv_types::DataType; + +use super::bind::unsupported; +use super::exprs::{Bound, Cx, pretty, semantic_key}; +use super::ops::col_expr; +use super::resolve::{Attr, Env, Rel, Resolver}; +use super::{AnalysisError, ArgType}; +use crate::{ + BoundColumn, CastMode, Expr, ExprNode, FunctionCall, FunctionKind, Grouping, Literal, + LogicalPlan, Origin, PlanNode, ResolvedExpr, +}; + +/// The most grouping expressions one `GROUP BY` takes: the grouping id is a +/// `BIGINT` with one bit each. +const MAX_GROUPING: usize = 63; + +/// `grouping` or `grouping_id`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum GroupingFn { + Grouping, + GroupingId, +} + +impl GroupingFn { + pub(super) fn from_name(name: &str) -> Option { + if name.eq_ignore_ascii_case("grouping") { + Some(Self::Grouping) + } else if name.eq_ignore_ascii_case("grouping_id") { + Some(Self::GroupingId) + } else { + None + } + } + + const fn name(self) -> &'static str { + match self { + Self::Grouping => "grouping", + Self::GroupingId => "grouping_id", + } + } + + const fn data_type(self) -> DataType { + match self { + Self::Grouping => DataType::Byte, + Self::GroupingId => DataType::Long, + } + } +} + +/// One distinct grouping expression. +pub(super) struct GroupExpr { + /// Its semantic key. + key: u128, + /// Its attribute in the expansion: a nullable copy. + pub(super) attr: Attr, +} + +/// The expansion an aggregate reads. +pub(super) struct ExpandPlan { + pub(super) plan: Arc, + /// The span of attribute ids the expansion defines. + pub(super) ids: Option<(u64, u64)>, +} + +/// An aggregate's grouping sets. +pub(super) struct Analytics { + /// The distinct grouping expressions, in order of first appearance. + pub(super) exprs: Vec, + /// The grouping id as the query names it (`grouping__id`); `None` for a + /// grand total, where it is 0. + gid: Option, + /// The expansion; `None` for a grand total. + pub(super) expand: Option, +} + +impl Analytics { + /// The attributes `HAVING` and `ORDER BY` may name: the grouping + /// expressions' copies and `grouping__id`. + pub(super) fn keys(&self) -> Vec { + self.exprs + .iter() + .map(|g| g.attr.clone()) + .chain(self.gid.iter().cloned()) + .collect() + } + + /// The position of `e` among the grouping expressions: a copy of one, or + /// an expression equal to one. + fn position(&self, e: &Expr, memo: &mut HashMap) -> Option { + if let ExprNode::Resolved(ResolvedExpr::Column(c)) = &e.node + && let Some(j) = self.exprs.iter().position(|g| g.attr.col.id == c.id) + { + return Some(j); + } + let key = semantic_key(e, memo); + self.exprs.iter().position(|g| g.key == key) + } +} + +fn literal(l: Literal) -> Arc { + Arc::new(Expr::new(ExprNode::Literal(l))) +} + +fn arg_type(data_type: DataType, nullable: bool) -> ArgType { + ArgType { + data_type: Some(data_type), + nullable, + } +} + +fn bound_of(col: &BoundColumn) -> Bound { + Bound { + expr: col_expr(col, Origin::UNKNOWN), + ty: arg_type(col.data_type.clone(), col.nullable), + agg: None, + outer: None, + lateral: Vec::new(), + shown: None, + } +} + +/// Spark's text for an expression in a grouping error: a column carries its +/// attribute id. +fn with_id(e: &Expr) -> String { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Column(c)) => format!("{}#{}", c.name, c.id.0), + _ => pretty(e), + } +} + +/// Whether a subquery is anywhere in `e`. +fn has_subquery(e: &Arc) -> bool { + let mut stack = vec![Arc::clone(e)]; + while let Some(x) = stack.pop() { + if matches!(x.node, ExprNode::Subquery { .. }) { + return true; + } + x.for_each_child(&mut |c| stack.push(Arc::clone(c))); + } + false +} + +/// The shape of a grouping: all its expressions in order, and each grouping +/// set as positions into them. +fn shape(grouping: &Grouping) -> (Vec>, Vec>) { + match grouping { + Grouping::GroupBy(v) => (v.clone(), vec![(0..v.len()).collect()]), + Grouping::Rollup(v) => ( + v.clone(), + (0..=v.len()).rev().map(|n| (0..n).collect()).collect(), + ), + Grouping::Cube(v) => { + // The full set first, then down to the empty one. + let n = v.len(); + let sets = (0..1_usize << n) + .rev() + .map(|m| (0..n).filter(|i| m >> (n - 1 - i) & 1 == 1).collect()) + .collect(); + (v.clone(), sets) + } + Grouping::GroupingSets { exprs, sets } => { + let mut flat = exprs.clone(); + let mut out = Vec::with_capacity(sets.len()); + for set in sets { + out.push((flat.len()..flat.len() + set.len()).collect()); + flat.extend(set.iter().cloned()); + } + (flat, out) + } + } +} + +impl Resolver<'_> { + /// `grouping(e)` or `grouping_id(..)` in a select list, `HAVING` or + /// `ORDER BY`: a call the aggregate resolves once it knows its grouping + /// sets ([`Self::resolve_grouping`]). Anywhere else it is + /// `UNSUPPORTED_GROUPING_EXPRESSION`. + #[inline(never)] + pub(super) fn grouping_call( + &mut self, + which: GroupingFn, + e: &Arc, + kids: &[Bound], + cx: &Cx<'_>, + ) -> Result { + // Spark has no source position for the error in `HAVING` and + // `ORDER BY`. + let origin = if matches!(cx.op, "Filter" | "Sort") { + Origin::UNKNOWN + } else { + e.origin + }; + if !cx.grouping { + return Err(AnalysisError::new(UnsupportedGroupingExpression {}, origin)); + } + if which == GroupingFn::Grouping && kids.len() != 1 { + return Err(AnalysisError::new( + WrongNumArgsWithoutSuggestion { + function_name: "`grouping`".to_owned(), + expected_num: "1".to_owned(), + actual_num: kids.len().to_string(), + docroot: super::exprs::SPARK_DOC_ROOT.to_owned(), + }, + e.origin, + )); + } + let call = FunctionCall { + name: which.name().to_owned(), + args: kids.iter().map(|k| Arc::clone(&k.expr)).collect(), + distinct: false, + filter: None, + }; + Ok(Bound { + expr: Arc::new(Expr::new(ExprNode::Function(call)).with_origin(e.origin)), + ty: arg_type(which.data_type(), false), + agg: kids.iter().find_map(|k| k.agg), + outer: kids.iter().filter_map(|k| k.outer).min(), + lateral: kids.iter().flat_map(|k| k.lateral.clone()).collect(), + shown: None, + }) + } + + /// A resolved `Expand` read again: its expressions bound over its + /// input, its output as it states it. + #[inline(never)] + pub(super) fn expand_node( + &mut self, + origin: Origin, + input: &Arc, + projections: &[Vec>], + output: &[BoundColumn], + env: &Arc, + ) -> Result { + let c = self.child(input, env)?; + let rels = [&c]; + let cx = Self::cx(&c.out, &rels, env, "Expand"); + let mut bound = Vec::with_capacity(projections.len()); + for row in projections { + let mut exprs = Vec::with_capacity(row.len()); + for e in row { + exprs.push(self.bind(e, &cx)?.expr); + } + bound.push(exprs); + } + let out = output + .iter() + .map(|col| Attr { + col: col.clone(), + qual: Arc::from(Vec::new()), + source: None, + }) + .collect(); + let node = PlanNode::Expand { + input: Arc::clone(&c.plan), + projections: bound, + output: output.to_vec(), + }; + Ok(self.make(origin, node, out, &[&c])) + } + + /// `e` with its `grouping` and `grouping_id` calls resolved against + /// `an`. + #[inline(never)] + pub(super) fn resolve_grouping( + &mut self, + e: &Arc, + an: &Analytics, + env: &Arc, + ) -> Result, AnalysisError> { + self.rewrite(e, an, false, env, &mut HashMap::new()) + } + + /// An aggregate over `ROLLUP`, `CUBE` or `GROUPING SETS`. + #[inline(never)] + pub(super) fn grouping_sets( + &mut self, + p: &LogicalPlan, + input: &Arc, + grouping: &Grouping, + aggregates: &[Arc], + env: &Arc, + ) -> Result { + let c = self.child(input, env)?; + let (flat, sets) = shape(grouping); + self.check_star_ordinals(&flat, aggregates)?; + // A single empty set over nothing is a grand total. + let grand = flat.is_empty() && sets.len() == 1; + let gid_col = (!grand).then(|| { + self.attr( + "spark_grouping_id", + &arg_type(DataType::Long, false), + ) + }); + let gid = gid_col.as_ref().map(|col| Attr { + col: BoundColumn { + name: "grouping__id".to_owned(), + ..col.clone() + }, + qual: Arc::from(Vec::new()), + source: None, + }); + let keys: Vec = gid.iter().cloned().collect(); + let items = self.select_items_in(&c, aggregates, env, "Aggregate", Some(&keys))?; + let bound = self.group_keys(&c, &flat, aggregates, &items, env)?; + let mut memo = HashMap::new(); + // The distinct expressions, and the sets as positions among them. + let mut distinct: Vec = Vec::new(); + let mut keys_of: Vec = Vec::new(); + let mut at = Vec::with_capacity(bound.len()); + for (i, b) in bound.iter().enumerate() { + let key = semantic_key(&b.expr, &mut memo); + let j = match keys_of.iter().position(|k| *k == key) { + Some(j) => j, + None => { + keys_of.push(key); + distinct.push(i); + keys_of.len() - 1 + } + }; + at.push(j); + } + let n = distinct.len(); + if n > MAX_GROUPING { + return Err(AnalysisError::new( + GroupingSizeLimitExceeded { + max_size: MAX_GROUPING.to_string(), + }, + p.origin, + )); + } + let sets: Vec> = sets + .into_iter() + .map(|s| { + let mut s: Vec = s.into_iter().map(|i| at[i]).collect(); + s.sort_unstable(); + s.dedup(); + s + }) + .collect(); + if grand { + let an = Analytics { + exprs: Vec::new(), + gid: None, + expand: None, + }; + let items = self.rewrite_items(items, &an, env)?; + return self.build_aggregate(p.origin, &c, Vec::new(), &items, true, Some(Arc::new(an))); + } + let Some(gid_col) = gid_col else { + return Err(unsupported("grouping sets without an id", p.origin)); + }; + let numbered = (0..sets.len()).any(|i| sets[..i].contains(&sets[i])); + // The expansion's input: the child's columns, and the grouping + // expressions that are not one of them. + let mut passed: Vec> = c.out.iter().map(|a| col_expr(&a.col, Origin::UNKNOWN)).collect(); + let mut passed_out: Vec = c.out.to_vec(); + let mut value: Vec = Vec::with_capacity(n); + let mut copies: Vec = Vec::with_capacity(n); + let mut computed = false; + for (j, &i) in distinct.iter().enumerate() { + let b = &bound[i]; + let column = match &b.expr.node { + ExprNode::Resolved(ResolvedExpr::Column(col)) => { + c.out.iter().find(|a| a.col.id == col.id).cloned() + } + _ => None, + }; + let name = match &column { + Some(a) => a.col.name.clone(), + None => pretty(&b.shown()), + }; + let v = match &column { + Some(a) => a.col.clone(), + None => { + let v = self.attr(&name, &b.ty); + passed.push(Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { + expr: Arc::clone(&b.expr), + output: v.clone(), + })))); + passed_out.push(Attr { + col: v.clone(), + qual: Arc::from(Vec::new()), + source: None, + }); + computed = true; + v + } + }; + let ty = arg_type(v.data_type.clone(), true); + let copy = self.attr(&name, &ty); + copies.push(GroupExpr { + key: keys_of[j], + attr: Attr { + col: copy, + qual: column.as_ref().map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), + source: column.and_then(|a| a.source), + }, + }); + value.push(v); + } + let below = if computed { + let out = passed_out.clone(); + self.projection(Origin::UNKNOWN, &c, passed, out) + } else { + c.clone() + }; + let pos_col = numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); + let mut output: Vec = passed_out.iter().map(|a| a.col.clone()).collect(); + output.extend(copies.iter().map(|g| g.attr.col.clone())); + output.push(gid_col.clone()); + output.extend(pos_col.iter().cloned()); + let mut projections = Vec::with_capacity(sets.len()); + for (si, set) in sets.iter().enumerate() { + let mut row: Vec> = passed_out + .iter() + .map(|a| col_expr(&a.col, Origin::UNKNOWN)) + .collect(); + let mut mask: i64 = 0; + for (j, v) in value.iter().enumerate() { + if set.contains(&j) { + row.push(col_expr(v, Origin::UNKNOWN)); + } else { + row.push(literal(Literal::Null(Some(v.data_type.clone())))); + mask |= 1 << (n - 1 - j); + } + } + row.push(literal(Literal::Long(mask))); + if pos_col.is_some() { + row.push(literal(Literal::Integer( + i32::try_from(si).map_err(|_| unsupported("that many grouping sets", p.origin))?, + ))); + } + projections.push(row); + } + let out_attrs: Vec = passed_out + .iter() + .cloned() + .chain(copies.iter().map(|g| g.attr.clone())) + .chain(std::iter::once(Attr { + col: gid_col.clone(), + qual: Arc::from(Vec::new()), + source: None, + })) + .chain(pos_col.iter().map(|col| Attr { + col: col.clone(), + qual: Arc::from(Vec::new()), + source: None, + })) + .collect(); + let node = PlanNode::Expand { + input: Arc::clone(&below.plan), + projections, + output, + }; + let expand = self.make(Origin::UNKNOWN, node, out_attrs, &[&below]); + let mut group: Vec = copies.iter().map(|g| bound_of(&g.attr.col)).collect(); + group.push(bound_of(&gid_col)); + group.extend(pos_col.iter().map(bound_of)); + let an = Analytics { + exprs: copies, + gid, + expand: Some(ExpandPlan { + plan: Arc::clone(&expand.plan), + ids: expand.ids, + }), + }; + let items = self.rewrite_items(items, &an, env)?; + self.build_aggregate(p.origin, &c, group, &items, false, Some(Arc::new(an))) + } + + /// The select `items` over the grouping sets: grouping expressions read + /// their copies, `grouping` and `grouping_id` are resolved. + #[inline(never)] + fn rewrite_items( + &mut self, + items: Vec, + an: &Analytics, + env: &Arc, + ) -> Result, AnalysisError> { + let mut memo = HashMap::new(); + let mut out = Vec::with_capacity(items.len()); + for b in items { + let expr = self.rewrite(&b.expr, an, true, env, &mut memo)?; + let ty = if Arc::ptr_eq(&expr, &b.expr) { + b.ty.clone() + } else { + self.type_of(&expr, env)? + }; + out.push(Bound { expr, ty, ..b }); + } + Ok(out) + } + + /// The type of the resolved `e`. + fn type_of(&mut self, e: &Arc, env: &Arc) -> Result { + let inner = match &e.node { + ExprNode::Resolved(ResolvedExpr::Alias { expr, .. }) => expr, + _ => e, + }; + Ok(self.bind(inner, &Self::cx(&[], &[], env, "Aggregate"))?.ty) + } + + /// `e` over the grouping sets: with `replace`, a subtree that is a + /// grouping expression reads its copy (aggregate calls and subqueries + /// are left as they are); `grouping` and `grouping_id` are resolved. + /// Nodes above a change are typed again, as the copies are nullable. + #[inline(never)] + fn rewrite( + &mut self, + e: &Arc, + an: &Analytics, + replace: bool, + env: &Arc, + memo: &mut HashMap, + ) -> Result, AnalysisError> { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) if f.kind == FunctionKind::Aggregate => { + return Ok(Arc::clone(e)); + } + ExprNode::Subquery { .. } => return Ok(Arc::clone(e)), + ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some() => { + return self.grouping_value(f, an, env, memo); + } + ExprNode::Resolved(ResolvedExpr::Alias { expr, output }) => { + let inner = self.rewrite(expr, an, replace, env, memo)?; + if Arc::ptr_eq(&inner, expr) { + return Ok(Arc::clone(e)); + } + let ty = self.type_of(&inner, env)?; + let output = BoundColumn { + nullable: ty.nullable, + ..output.clone() + }; + return Ok(Arc::new( + Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { + expr: inner, + output, + })) + .with_origin(e.origin), + )); + } + _ => {} + } + if replace && let Some(j) = an.position(e, memo) { + return Ok(col_expr(&an.exprs[j].attr.col, e.origin)); + } + let rebuilt = Arc::clone(e).map_children(|c| self.rewrite(&c, an, replace, env, memo))?; + if Arc::ptr_eq(&rebuilt, e) { + return Ok(rebuilt); + } + self.retype(rebuilt, env) + } + + /// `e`, whose children changed, typed again. + fn retype(&mut self, e: Arc, env: &Arc) -> Result, AnalysisError> { + let origin = e.origin; + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) if !has_subquery(&e) => { + self.call(f.id.name(), f.args.clone(), env, origin) + } + ExprNode::Resolved(ResolvedExpr::Function(f)) => { + // A call over a subquery keeps its type, nullable. + let mut f = f.clone(); + f.nullable = true; + Ok(Arc::new( + Expr::new(ExprNode::Resolved(ResolvedExpr::Function(f))).with_origin(origin), + )) + } + ExprNode::Cast { .. } if !has_subquery(&e) => { + Ok(self.bind(&e, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + } + _ => Ok(e), + } + } + + /// The value of a `grouping` or `grouping_id` call. + fn grouping_value( + &mut self, + f: &FunctionCall, + an: &Analytics, + env: &Arc, + memo: &mut HashMap, + ) -> Result, AnalysisError> { + let n = an.exprs.len(); + let id = match &an.gid { + Some(g) => col_expr(&g.col, Origin::UNKNOWN), + None => literal(Literal::Long(0)), + }; + let positions: Vec> = + f.args.iter().map(|a| an.position(a, memo)).collect(); + match GroupingFn::from_name(&f.name) { + Some(GroupingFn::GroupingId) => { + let exact = positions.len() == n + && positions.iter().enumerate().all(|(i, p)| *p == Some(i)); + if !f.args.is_empty() && !exact { + return Err(AnalysisError::new( + GroupingIdColumnMismatch { + grouping_id_column: f + .args + .iter() + .map(|a| with_id(a)) + .collect::>() + .join(","), + group_by_columns: an + .exprs + .iter() + .map(|g| format!("{}#{}", g.attr.col.name, g.attr.col.id.0)) + .collect::>() + .join(","), + }, + Origin::UNKNOWN, + )); + } + Ok(id) + } + _ => { + let (Some(arg), Some(Some(j))) = (f.args.first(), positions.first()) else { + return Err(AnalysisError::new( + GroupingColumnMismatch { + grouping: f.args.first().map(|a| with_id(a)).unwrap_or_default(), + grouping_columns: an + .exprs + .iter() + .map(|g| format!("{}#{}", g.attr.col.name, g.attr.col.id.0)) + .collect::>() + .join(","), + }, + Origin::UNKNOWN, + )); + }; + let _ = arg; + // Bit `n - 1 - j` of the id, as a `TINYINT`. + let shift = n - 1 - j; + let high = if shift == 0 { + id + } else { + self.call( + "div", + vec![id, literal(Literal::Long(1_i64 << shift))], + env, + Origin::UNKNOWN, + )? + }; + let bit = self.call( + "%", + vec![high, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )?; + let cast = Arc::new(Expr::new(ExprNode::Cast { + expr: bit, + to: DataType::Byte, + mode: CastMode::Default, + })); + Ok(self.bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + } + } + } +} diff --git a/crates/krishiv-plan/src/analyzer/bind.rs b/crates/krishiv-plan/src/analyzer/bind.rs index 3a8dcc69..02f44f54 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 6eb5d5a2..b16b6929 100644 --- a/crates/krishiv-plan/src/analyzer/exprs.rs +++ b/crates/krishiv-plan/src/analyzer/exprs.rs @@ -21,6 +21,7 @@ use krishiv_common::error::{ use krishiv_types::DataType; use twox_hash::XxHash3_128; +use super::analytics::GroupingFn; use super::bind::{addr, same, unsupported}; use super::catalog::quoted; use super::coerce::{self, Arg, BadCast, Mismatch, ordinal, quoted_type}; @@ -50,6 +51,10 @@ pub(super) struct Cx<'s> { /// Over an aggregate: its input, where aggregate-function arguments /// resolve. pub(super) agg_scope: Option<&'s [Attr]>, + /// Whether `grouping(..)` and `grouping_id(..)` may appear: in the + /// select list, `HAVING` and `ORDER BY` of an aggregate over grouping + /// sets (see `analytics`). + pub(super) grouping: bool, /// A name `scope` holds several times is looked up further rather /// than ambiguous (Spark's `ORDER BY` over a select list). pub(super) lenient: bool, @@ -204,7 +209,7 @@ fn arg_of(b: &Bound) -> Arg<'_> { } /// Spark's documentation root, as its errors cite it. -const SPARK_DOC_ROOT: &str = "https://spark.apache.org/docs/latest"; +pub(super) const SPARK_DOC_ROOT: &str = "https://spark.apache.org/docs/latest"; fn column(c: &BoundColumn) -> ArgType { arg_type(c.data_type.clone(), c.nullable) @@ -444,6 +449,9 @@ impl Resolver<'_> { e.origin, )); } + if let Some(which) = GroupingFn::from_name(&f.name) { + return self.grouping_call(which, e, &kids[..n], cx); + } if let Some(op) = Builtin::from_name(&f.name) { return self.operator_call(op, e.origin, &kids[..n]); } diff --git a/crates/krishiv-plan/src/analyzer/mod.rs b/crates/krishiv-plan/src/analyzer/mod.rs index 44bca066..560064ff 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 1e932bd8..8238aec0 100644 --- a/crates/krishiv-plan/src/analyzer/ops.rs +++ b/crates/krishiv-plan/src/analyzer/ops.rs @@ -14,6 +14,7 @@ use krishiv_common::error::{ UnsupportedFeatureLateralColumnAliasInGroupBy, }; +use super::analytics::Analytics; use super::bind::{BoundTable, same, unsupported}; use super::catalog::quoted; use super::coerce::quoted_type; @@ -244,7 +245,7 @@ impl Resolver<'_> { ) -> Result { let items = self.select_items(c, exprs, env, "Aggregate")?; if items.iter().any(|b| b.agg.is_some()) { - return self.build_aggregate(p.origin, c, Vec::new(), &items, true); + return self.build_aggregate(p.origin, c, Vec::new(), &items, true, None); } Ok(self.lateral_project(p.origin, c, &items, &[])) } @@ -258,6 +259,20 @@ impl Resolver<'_> { exprs: &[Arc], env: &Arc, op: &'static str, + ) -> Result, AnalysisError> { + self.select_items_in(c, exprs, env, op, None) + } + + /// [`Self::select_items`] of an aggregate over grouping sets, which + /// may call `grouping` and `grouping_id` and name `grouping__id` + /// (`extra`, searched after the input). + pub(super) fn select_items_in( + &mut self, + c: &Rel, + exprs: &[Arc], + env: &Arc, + op: &'static str, + extra: Option<&[Attr]>, ) -> Result, AnalysisError> { let rels = [c]; let mut items = Vec::new(); @@ -274,6 +289,8 @@ impl Resolver<'_> { let b = { let mut cx = Self::cx(&c.out, &rels, env, op); cx.lateral = lca.then_some(&lateral[..]); + cx.extra = extra; + cx.grouping = extra.is_some(); self.bind(e, &cx)? }; let b = self.named(b)?; @@ -660,10 +677,16 @@ impl Resolver<'_> { let mut cx = Self::cx(&agg.out, &rels, env, op); cx.extra = Some(&info.keys); cx.agg_scope = Some(&info.input.out); + cx.grouping = info.analytics.is_some(); let mut bounds = Vec::new(); for e in exprs { bounds.push(self.bind(e, &cx)?); } + if let Some(an) = &info.analytics { + for b in &mut bounds { + b.expr = self.resolve_grouping(&b.expr, an, env)?; + } + } let mut memo = HashMap::new(); let mut known: Vec<(u128, Attr)> = info .aggregates @@ -754,8 +777,14 @@ impl Resolver<'_> { /// aliases if it has any; `pass` (attributes of the aggregate) are /// output after the select items. fn aggregate_rel(&self, origin: Origin, info: AggInfo, pass: &[Attr]) -> Rel { + // Over grouping sets, the aggregate reads their expansion; the + // attributes it defines are part of the aggregate's id span. + let (input, defined) = match info.analytics.as_deref().and_then(|a| a.expand.as_ref()) { + Some(e) => (Arc::clone(&e.plan), e.ids), + None => (Arc::clone(&info.input.plan), None), + }; let node = PlanNode::Aggregate { - input: Arc::clone(&info.input.plan), + input, grouping: Grouping::GroupBy(info.grouping.clone()), aggregates: info.aggregates.clone(), }; @@ -766,6 +795,7 @@ impl Resolver<'_> { self.lateral_project(origin, &a, top, pass) } }; + r.ids = super::resolve::hull(r.ids, defined); let info = Arc::new(info); r.widen = Some(Arc::new(Widen::Aggregate(Arc::clone(&info)))); r.agg = Some(info); @@ -781,38 +811,66 @@ impl Resolver<'_> { env: &Arc, ) -> Result { let Grouping::GroupBy(keys) = grouping else { - return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", p.origin)); + return self.grouping_sets(p, input, grouping, aggregates, env); }; let c = self.child(input, env)?; - let conf = self.ctx.conf; - let ordinal = |k: &Expr| match &k.node { - ExprNode::Literal(Literal::Integer(n)) if conf.group_by_ordinal => Some(*n), - _ => None, - }; + self.check_star_ordinals(keys, aggregates)?; + // Spark resolves the select list first, then ordinals, then the + // other keys. + let items = self.select_items(&c, aggregates, env, "Aggregate")?; + let group = self.group_keys(&c, keys, aggregates, &items, env)?; + self.build_aggregate(p.origin, &c, group, &items, false, None) + } + + /// `SELECT *` with `GROUP BY n` is an error (Spark's + /// `STAR_GROUP_BY_POS`). + pub(super) fn check_star_ordinals( + &self, + keys: &[Arc], + aggregates: &[Arc], + ) -> Result<(), AnalysisError> { let star = aggregates .iter() .any(|e| matches!(e.node, ExprNode::Star(_))); - if star && keys.iter().any(|k| ordinal(k).is_some()) { + if star && keys.iter().any(|k| self.ordinal(k).is_some()) { return Err(AnalysisError::new(StarGroupByPos {}, Origin::UNKNOWN)); } - // Spark resolves the select list first, then ordinals, then the - // other keys. - let items = self.select_items(&c, aggregates, env, "Aggregate")?; + Ok(()) + } + + /// `GROUP BY n` with ordinals on: `n`. + fn ordinal(&self, k: &Expr) -> Option { + match &k.node { + ExprNode::Literal(Literal::Integer(n)) if self.ctx.conf.group_by_ordinal => Some(*n), + _ => None, + } + } + + /// The grouping expressions `keys` over `c`: an ordinal names a select + /// item, other names may be select-list aliases. + pub(super) fn group_keys( + &mut self, + c: &Rel, + keys: &[Arc], + aggregates: &[Arc], + items: &[Bound], + env: &Arc, + ) -> Result, AnalysisError> { let mut group: Vec> = vec![None; keys.len()]; for (slot, k) in group.iter_mut().zip(keys) { - if let Some(n) = ordinal(k) { - *slot = Some(self.group_ordinal(n, k.origin, &items, &c)?); + if let Some(n) = self.ordinal(k) { + *slot = Some(self.group_ordinal(n, k.origin, items, c)?); } } let aliases: Vec<(String, Arc)> = aggregates .iter() - .filter(|_| conf.group_by_aliases) + .filter(|_| self.ctx.conf.group_by_aliases) .filter_map(|e| match &e.node { ExprNode::Alias { expr, name } => Some((name.clone(), Arc::clone(expr))), _ => None, }) .collect(); - let rels = [&c]; + let rels = [c]; for (slot, k) in group.iter_mut().zip(keys) { if slot.is_none() { let mut cx = Self::cx(&c.out, &rels, env, "Aggregate"); @@ -820,8 +878,7 @@ impl Resolver<'_> { *slot = Some(self.bind(k, &cx)?); } } - let group = group.into_iter().flatten().collect(); - self.build_aggregate(p.origin, &c, group, &items, false) + Ok(group.into_iter().flatten().collect()) } /// `GROUP BY n`: the `n`-th select item, unaliased. @@ -885,6 +942,7 @@ impl Resolver<'_> { group: Vec, items: &[Bound], global: bool, + analytics: Option>, ) -> Result { let mut memo = HashMap::new(); let gkeys: Vec<(u128, ArgType)> = group @@ -898,22 +956,33 @@ impl Resolver<'_> { .cloned() .collect(); self.check_aggregation(&base, &key_set, &mut memo, origin, global)?; - let keys: Vec = group - .iter() - .filter_map(|b| match &b.expr.node { - ExprNode::Resolved(ResolvedExpr::Column(col)) => { - c.out.iter().find(|a| a.col.id == col.id).cloned() - } - _ => None, - }) - .collect(); + // Over grouping sets, the grouping attributes are the expansion's. + let (keys, scope): (Vec, Vec) = match &analytics { + Some(a) => ( + a.keys(), + c.out.iter().chain(a.exprs.iter().map(|g| &g.attr)).cloned().collect(), + ), + None => ( + group + .iter() + .filter_map(|b| match &b.expr.node { + ExprNode::Resolved(ResolvedExpr::Column(col)) => { + c.out.iter().find(|a| a.col.id == col.id).cloned() + } + _ => None, + }) + .collect(), + c.out.to_vec(), + ), + }; let mut info = AggInfo { input: c.clone(), grouping: group.into_iter().map(|b| b.expr).collect(), aggregates: base.iter().map(|b| Arc::clone(&b.expr)).collect(), - out: base.iter().map(|b| out_attr(b, &c.out)).collect(), + out: base.iter().map(|b| out_attr(b, &scope)).collect(), keys, top: None, + analytics, }; if base.len() == items.len() { return Ok(self.aggregate_rel(origin, info, &[])); @@ -925,7 +994,7 @@ impl Resolver<'_> { let mut top = Vec::new(); for b in items { if b.lateral.is_empty() { - let a = out_attr(b, &c.out); + let a = out_attr(b, &scope); top.push(Bound { expr: col_expr(&a.col, b.expr.origin), ..b.clone() diff --git a/crates/krishiv-plan/src/analyzer/resolve.rs b/crates/krishiv-plan/src/analyzer/resolve.rs index da70c8c9..0c6f4f56 100644 --- a/crates/krishiv-plan/src/analyzer/resolve.rs +++ b/crates/krishiv-plan/src/analyzer/resolve.rs @@ -39,6 +39,9 @@ pub(super) struct AggInfo { /// above the aggregate: every select item, over the aggregate's output /// (`None` when no item does). pub(super) top: Option>, + /// The grouping sets, for an aggregate over `ROLLUP`, `CUBE` or + /// `GROUPING SETS`: its input is then their expansion. + pub(super) analytics: Option>, } /// How a relation gets back columns it dropped: Spark adds the attributes @@ -193,7 +196,7 @@ fn first_free_id(root: &LogicalPlan) -> Result { Ok(next) } -fn hull(a: Option<(u64, u64)>, b: Option<(u64, u64)>) -> Option<(u64, u64)> { +pub(super) fn hull(a: Option<(u64, u64)>, b: Option<(u64, u64)>) -> Option<(u64, u64)> { match (a, b) { (Some(x), Some(y)) => Some((x.0.min(y.0), x.1.max(y.1))), (x, None) | (None, x) => x, @@ -446,6 +449,7 @@ impl Resolver<'_> { aliases: &[], extra: None, agg_scope: None, + grouping: false, lenient: false, outer: true, env, @@ -624,6 +628,11 @@ impl Resolver<'_> { }; Ok(self.make(o, node, c.out.to_vec(), &[&c])) } + PlanNode::Expand { + input, + projections, + output, + } => self.expand_node(o, input, projections, output, env), PlanNode::Distinct { input } => self.distinct(p, input, env), PlanNode::Hint { input, diff --git a/crates/krishiv-plan/src/analyzer/setops.rs b/crates/krishiv-plan/src/analyzer/setops.rs index 9e25f448..0ac2ae78 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 7483abc9..092b6eee 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -288,6 +288,26 @@ impl Display for NodeLine<'_> { PlanNode::ReplicateRows { count, .. } => { write!(f, "Generate replicaterows({})", Text(count)) } + // Spark's `Expand [[a, null, 1], [a, b, 0]], [a, b, spark_grouping_id]`. + PlanNode::Expand { + projections, + output, + .. + } => { + f.write_str("Expand [")?; + for (i, p) in projections.iter().enumerate() { + write!(f, "{}[", sep(i))?; + for (j, e) in p.iter().enumerate() { + write!(f, "{}{}", sep(j), Text(e))?; + } + f.write_str("]")?; + } + f.write_str("], [")?; + for (i, c) in output.iter().enumerate() { + write!(f, "{}{}#{}", sep(i), c.name, c.id.0)?; + } + f.write_str("]") + } PlanNode::WithCte { ctes, .. } => { f.write_str("WithCte [")?; for (i, c) in ctes.iter().enumerate() { diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index 2957c5fc..976a4f49 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -791,6 +791,17 @@ impl<'k, 'a> Encoder<'k, 'a> { self.u8(21); self.expr(count); } + PlanNode::Expand { + projections, + output, + .. + } => { + self.u8(22); + self.len(projections.len()); + projections.iter().for_each(|p| self.exprs(p)); + self.len(output.len()); + output.iter().for_each(|c| self.bound(c)); + } } } diff --git a/crates/krishiv-plan/src/lib.rs b/crates/krishiv-plan/src/lib.rs index 35ddaf2e..79607784 100644 --- a/crates/krishiv-plan/src/lib.rs +++ b/crates/krishiv-plan/src/lib.rs @@ -67,7 +67,7 @@ //! | `PlanNode::Project` | U, R | Q2, Q4, Q6-Q9, Q11, Q14-Q17, Q19-Q22 | `select` | //! | `PlanNode::WithColumns` | U | - | `withColumn` | //! | `PlanNode::Filter` | U, R | Q1-Q12, Q14-Q22 | `filter`, `where` | -//! | `PlanNode::Aggregate` | U, R | Q1, Q3-Q5, Q7-Q13, Q15, Q16, Q18, Q21, Q22 | `groupBy.agg` | +//! | `PlanNode::Aggregate` | U, R | Q1, Q3-Q5, Q7-Q13, Q15, Q16, Q18, Q21, Q22 | `groupBy.agg`, `rollup.agg`, `cube.agg` | //! | `PlanNode::Join` | U, R | Q2, Q3, Q5, Q7-Q21 | `join` | //! | `PlanNode::Sort` | U, R | Q1-Q5, Q7-Q13, Q15, Q16, Q18, Q20-Q22 | `orderBy`, `sort` | //! | `PlanNode::Limit` | U, R | Q2, Q3, Q10, Q18, Q21 | `limit` | @@ -81,6 +81,7 @@ //! | `PlanNode::Distinct` | U | - | - | //! | `PlanNode::SetOp` | U | - | - | //! | `PlanNode::ReplicateRows` | R | - | - | +//! | `PlanNode::Expand` | R | - | - | //! | `ExprNode::Column` | U | all | `col`, `df["a"]` | //! | `ExprNode::Literal` | U, R | all | `lit` | //! | `ExprNode::Star` | U | Q1, Q4, Q13, Q21, Q22 | `select("*")` | diff --git a/crates/krishiv-plan/src/plan.rs b/crates/krishiv-plan/src/plan.rs index 59df152d..d0e791be 100644 --- a/crates/krishiv-plan/src/plan.rs +++ b/crates/krishiv-plan/src/plan.rs @@ -243,6 +243,19 @@ pub enum PlanNode { /// A `BIGINT` expression over the input: the copies of each row. count: Arc, }, + /// Each input row once per projection (Spark's `Expand`, the first step + /// of `ROLLUP`, `CUBE` and `GROUPING SETS`): row `i` of the output for + /// an input row is the row evaluated by `projections[i]`. Resolved + /// plans only: the analyzer builds it, with one expression per output + /// column in every projection. + Expand { + /// Input. + input: Arc, + /// The projections, in output order; at least one. + projections: Vec>>, + /// The output attributes, aligned with each projection. + output: Vec, + }, } /// The operation of a [`PlanNode::SetOp`]. @@ -554,6 +567,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } + | PlanNode::Expand { input, .. } | PlanNode::ReplicateRows { input, .. } => f(input), PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { f(left); @@ -585,6 +599,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } + | PlanNode::Expand { input, .. } | PlanNode::ReplicateRows { input, .. } => f(input), PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { f(left); @@ -644,6 +659,7 @@ impl LogicalPlan { JoinCondition::None | JoinCondition::Using(_) => {} }, PlanNode::Sort { order, .. } => order.iter().for_each(|o| f(&o.expr)), + PlanNode::Expand { projections, .. } => projections.iter().flatten().for_each(f), } } @@ -694,6 +710,9 @@ impl LogicalPlan { JoinCondition::None | JoinCondition::Using(_) => {} }, PlanNode::Sort { order, .. } => order.iter_mut().for_each(|o| f(&mut o.expr)), + PlanNode::Expand { projections, .. } => { + projections.iter_mut().flatten().for_each(f); + } } } @@ -748,6 +767,7 @@ impl LogicalPlan { ); } PlanNode::OneRowRelation => return Some(Vec::new()), + PlanNode::Expand { output, .. } => return Some(output.clone()), PlanNode::Scan(s) => return Some(s.output.clone()), PlanNode::CteRef(r) => return Some(r.output.clone()), PlanNode::Shift { input, offset } => { diff --git a/crates/krishiv-plan/src/tests/analyzer.rs b/crates/krishiv-plan/src/tests/analyzer.rs index 1e280ceb..c05b3ba6 100644 --- a/crates/krishiv-plan/src/tests/analyzer.rs +++ b/crates/krishiv-plan/src/tests/analyzer.rs @@ -1423,13 +1423,8 @@ fn unsupported_inputs_fail_explicitly() { paths: vec![krishiv_common::secret::OptionValue::from_user("path", "/d")], options: std::collections::BTreeMap::new(), })); - let rollup = LogicalPlan::new(PlanNode::Aggregate { - input: Arc::new(LogicalPlan::table(["t"])), - grouping: Grouping::Rollup(vec![Arc::new(c("a"))]), - aggregates: vec![Arc::new(c("a"))], - }); let sha = LogicalPlan::table(["t"]).project(vec![Expr::function("factorial", vec![c("b")])]); - for plan in [files, rollup, sha] { + for plan in [files, sha] { assert_eq!(class(&run(plan, &ctx)), "KRISHIV_UNSUPPORTED_FEATURE"); } // An untyped NULL is the void type; an interval is not supported yet. diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index 011e3a13..fcd40199 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -12,7 +12,7 @@ use krishiv_types::DataType; use super::dsl::{c, int}; use super::tpch::QUERIES; use crate::{ - Builtin, ColumnRef, Expr, ExprNode, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, + Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, }; @@ -41,6 +41,7 @@ pub(super) fn kinds(plan: &LogicalPlan, out: &mut BTreeSet<&'static str>) { PlanNode::Distinct { .. } => "PlanNode::Distinct", PlanNode::SetOp { .. } => "PlanNode::SetOp", PlanNode::ReplicateRows { .. } => "PlanNode::ReplicateRows", + PlanNode::Expand { .. } => "PlanNode::Expand", }); plan.for_each_expr(&mut |e| expr_kinds(e, out)); plan.for_each_child(&mut |c| kinds(c, out)); @@ -153,6 +154,13 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { columns: vec![("b".to_owned(), Arc::new(c("a")))], }); let sum = Expr::builtin(Builtin::Add, vec![c("a"), int(1)]); + let grouped = |grouping| { + LogicalPlan::new(PlanNode::Aggregate { + input: Arc::new(df(1)), + grouping, + aggregates: vec![Arc::new(c("a"))], + }) + }; vec![ ("read.table", df(1)), ("read.parquet", parquet), @@ -162,6 +170,8 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { ("filter", df(1).filter(c("a"))), ("where", df(1).filter(c("a"))), ("groupBy.agg", df(1).aggregate(vec![c("a")], vec![c("a")])), + ("rollup.agg", grouped(Grouping::Rollup(vec![Arc::new(c("a"))]))), + ("cube.agg", grouped(Grouping::Cube(vec![Arc::new(c("a"))]))), ( "join", df(1).join(df(2), JoinKind::Inner, JoinCondition::None), diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index 715382a2..ec013bb5 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -288,7 +288,7 @@ impl Gen<'_> { fn plan(&mut self, depth: u8) -> LogicalPlan { let origin = self.origin(); - let kinds = if depth == 0 { 4 } else { 21 }; + let kinds = if depth == 0 { 4 } else { 22 }; let d = depth.saturating_sub(1); let node = match self.pick(kinds) { 0 => PlanNode::Read(self.source()), @@ -424,6 +424,11 @@ impl Gen<'_> { input: self.input(d), count: Arc::new(self.expr(1)), }, + 21 => PlanNode::Expand { + input: self.input(d), + projections: self.many(1, 2, |g| g.exprs()), + output: self.many(0, 2, Gen::bound), + }, _ => PlanNode::WithCte { input: self.input(d), ctes: self.many(1, 2, |g| Cte { @@ -813,6 +818,22 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { } | PlanNode::Limit { input, limit: e } | PlanNode::ReplicateRows { input, count: e } => m_expr(e, w) || m_plan(input, w), + PlanNode::Expand { + input, + projections, + output, + } => { + site!(w, projections.push(Vec::new())); + site!(w, output.push(output.first().cloned().unwrap_or(BoundColumn { + id: AttrId(0), + name: "x".to_owned(), + data_type: DataType::Integer, + nullable: false, + }))); + projections.iter_mut().any(|p| m_exprs(p, w)) + || output.iter_mut().any(|c| m_bound(c, w)) + || m_plan(input, w) + } PlanNode::SetOp { left, right, diff --git a/crates/krishiv-plan/src/walk.rs b/crates/krishiv-plan/src/walk.rs index f79c4154..ee0d9ed3 100644 --- a/crates/krishiv-plan/src/walk.rs +++ b/crates/krishiv-plan/src/walk.rs @@ -155,6 +155,7 @@ pub(crate) fn all_resolved(root: Node<'_>) -> bool { | PlanNode::SubqueryAlias { .. } | PlanNode::WithCte { .. } | PlanNode::JoinHint { .. } + | PlanNode::Expand { .. } | PlanNode::ReplicateRows { .. } => true, }, Node::Expr(e) => match &e.node { @@ -257,6 +258,12 @@ impl Teardown { self.plans.push(input); self.exprs.push(e); } + PlanNode::Expand { + input, projections, .. + } => { + self.plans.push(input); + self.exprs.extend(projections.into_iter().flatten()); + } PlanNode::Aggregate { input, grouping, diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index de0b9a94..5c499c5b 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -255,9 +255,73 @@ impl DataFrame { /// ``` #[must_use] pub fn group_by(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::GroupBy) + } + + /// Groups rows by `keys` and by each prefix of them, down to the grand + /// total, for [`GroupedData::agg`] (PySpark's `rollup`): a key a row is + /// rolled up over reads `NULL`. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s.table("t").rollup([col("a"), col("b")]).agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn rollup(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::Rollup) + } + + /// Groups rows by every subset of `keys`, for [`GroupedData::agg`] + /// (PySpark's `cube`): a key a row is rolled up over reads `NULL`. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s.table("t").cube([col("a"), col("b")]).agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn cube(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::Cube) + } + + /// Groups rows by each of `sets` (grouping expressions among `keys`), + /// for [`GroupedData::agg`] (PySpark's `groupingSets`): the output has + /// the columns of `keys`, `NULL` where a set does not group by them. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s + /// .table("t") + /// .grouping_sets([vec![col("a")], vec![col("b")], vec![]], [col("a"), col("b")]) + /// .agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn grouping_sets( + &self, + sets: impl IntoIterator>, + keys: impl IntoIterator>, + ) -> GroupedData { + let sets = sets + .into_iter() + .map(|set| set.into_iter().map(|c| c.expr).collect()) + .collect(); + self.grouped(keys, Sets::Explicit(sets)) + } + + fn grouped( + &self, + keys: impl IntoIterator>, + sets: Sets, + ) -> GroupedData { GroupedData { df: self.clone(), keys: keys.into_iter().map(|k| k.into().expr).collect(), + sets, } } @@ -749,11 +813,26 @@ impl DataFrame { } } -/// Rows grouped by keys ([`DataFrame::group_by`]). +/// Which groups of its keys a [`GroupedData`] aggregates. +#[derive(Debug, Clone)] +enum Sets { + /// All the keys together. + GroupBy, + /// Each prefix of the keys. + Rollup, + /// Each subset of the keys. + Cube, + /// These sets. + Explicit(Vec>), +} + +/// Rows grouped by keys ([`DataFrame::group_by`], [`DataFrame::rollup`], +/// [`DataFrame::cube`], [`DataFrame::grouping_sets`]). #[derive(Debug, Clone)] pub struct GroupedData { df: DataFrame, keys: Vec, + sets: Sets, } impl GroupedData { @@ -771,9 +850,21 @@ impl GroupedData { let keys: Vec> = self.keys.iter().cloned().map(Arc::new).collect(); let mut output = keys.clone(); output.extend(aggregates.into_iter().map(|c| Arc::new(c.expr))); + let grouping = match &self.sets { + Sets::GroupBy => Grouping::GroupBy(keys), + Sets::Rollup => Grouping::Rollup(keys), + Sets::Cube => Grouping::Cube(keys), + Sets::Explicit(sets) => Grouping::GroupingSets { + exprs: keys, + sets: sets + .iter() + .map(|set| set.iter().cloned().map(Arc::new).collect()) + .collect(), + }, + }; self.df.derive(PlanNode::Aggregate { input: Arc::clone(&self.df.plan), - grouping: Grouping::GroupBy(keys), + grouping, aggregates: output, }) } diff --git a/crates/krishiv-session/tests/scratch_rollup.rs b/crates/krishiv-session/tests/scratch_rollup.rs new file mode 100644 index 00000000..f12c6053 --- /dev/null +++ b/crates/krishiv-session/tests/scratch_rollup.rs @@ -0,0 +1,84 @@ +//! scratch +#![cfg(test)] +use std::sync::Arc; + +use arrow_array::{Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType as T, Field, Schema}; +use futures::executor::block_on; +use krishiv_session::Session; + +#[test] +fn scratch() { + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", T::Int32, false), + Field::new("b", T::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2, 3, 3])), + Arc::new(Int32Array::from(vec![1, 2, 1, 2, 1, 2])), + ], + ) + .unwrap(); + s.register_batches("td", schema, vec![batch]).unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("course", T::Utf8, false), + Field::new("year", T::Int32, false), + Field::new("earnings", T::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["dotNET", "Java", "dotNET", "dotNET", "Java"])), + Arc::new(Int32Array::from(vec![2012, 2012, 2012, 2013, 2013])), + Arc::new(Int32Array::from(vec![10000, 20000, 5000, 48000, 30000])), + ], + ) + .unwrap(); + s.register_batches("cs", schema, vec![batch]).unwrap(); + for q in std::fs::read_to_string("/home/krishiv/krishiv-data/tmp/s-rollup/q.sql") + .unwrap() + .split(";\n") + .filter(|q| !q.trim().is_empty()) + { + match block_on(s.sql(q)) { + Err(e) => println!("ERR {q}\n {e}"), + Ok(df) => match block_on(df.collect()) { + Err(e) => println!("RUNERR {q}\n {e}"), + Ok(b) => { + println!("OK {q}"); + for b in &b { + println!("{}", arrow_cast_fmt(b)); + } + } + }, + } + } +} + +fn arrow_cast_fmt(b: &RecordBatch) -> String { + let mut out = format!("{:?}\n", b.schema().fields().iter().map(|f| format!("{}:{}{}", f.name(), f.data_type(), if f.is_nullable() {"?"} else {""})).collect::>()); + for r in 0..b.num_rows() { + let cells: Vec = b + .columns() + .iter() + .map(|c| { + if c.is_null(r) { + "NULL".to_owned() + } else { + let d = arrow_array::cast::AsArray::as_primitive_opt::(c) + .map(|a| a.value(r).to_string()) + .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) + .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) + .or_else(|| arrow_array::cast::AsArray::as_string_opt::(c).map(|a| a.value(r).to_owned())); + d.unwrap_or_else(|| "?".to_owned()) + } + }) + .collect(); + out.push_str(&cells.join("\t")); + out.push('\n'); + } + out +} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index 6a4e5462..8daa1db3 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -27,7 +27,8 @@ use krishiv_plan::{ use krishiv_types::DecimalType; use crate::ast::{ - Alias, BinaryOp, Expr, ExprKind, FromClause, GroupItemKind, JoinConstraint, JoinKind, Literal, + Alias, BinaryOp, Expr, ExprKind, FromClause, GroupBy, GroupItem, GroupItemKind, GroupWith, + GroupingSet, JoinConstraint, JoinKind, Literal, NullTreatment, NumKind, Query, Relation, RelationExt, Select, SetExpr, SetExprKind, SetOperator, SetQuantifier, SortItem, Span, TableFactor, TableFactorKind, Truth, UnaryOp, }; @@ -36,6 +37,9 @@ use crate::{ParseError, ParserOptions, parse_query}; type R = Result; +/// The most grouping sets one `GROUP BY` expands to. +const MAX_GROUPING_SETS: usize = 1 << 16; + /// Parses and lowers one query; spans index `source`. /// /// # Errors @@ -389,22 +393,14 @@ impl Lowerer { s.span, )); } - let mut keys = Vec::new(); - if let Some(g) = &s.group_by { - if g.with.is_some() { - return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", g.span)); - } - for item in &g.items { - match &item.kind { - GroupItemKind::Expr(e) => keys.push(arc(self.expr(e)?)), - _ => return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", item.span)), - } - } - } + let grouping = match &s.group_by { + Some(g) => self.grouping(g)?, + None => Grouping::GroupBy(Vec::new()), + }; plan = self.node( PlanNode::Aggregate { input: arc(plan), - grouping: Grouping::GroupBy(keys), + grouping, aggregates: items, }, s.span, @@ -437,6 +433,129 @@ impl Lowerer { Ok(plan) } + /// `GROUP BY`, as Spark's `AstBuilder` shapes it: plain keys, `WITH + /// ROLLUP` / `WITH CUBE` over them, `GROUPING SETS (..)` after them (the + /// keys only fix the grouping expressions), or items that are + /// `ROLLUP`, `CUBE` and `GROUPING SETS` among keys, whose grouping sets + /// are the cross product of the items' (a key is a set of itself). + #[inline(never)] + fn grouping(&mut self, g: &GroupBy) -> R { + let mut plain = Vec::with_capacity(g.items.len()); + for item in &g.items { + if let GroupItemKind::Expr(e) = &item.kind { + plain.push(arc(self.expr(e)?)); + } + } + let all_plain = plain.len() == g.items.len(); + match (&g.with, all_plain) { + (None, true) => return Ok(Grouping::GroupBy(plain)), + (Some(GroupWith::Rollup), true) => return Ok(Grouping::Rollup(plain)), + (Some(GroupWith::Cube), true) => return Ok(Grouping::Cube(plain)), + (Some(GroupWith::GroupingSets(sets)), true) => { + let sets = sets + .iter() + .map(|s| self.set(s)) + .collect::>>()?; + return Ok(Grouping::GroupingSets { exprs: plain, sets }); + } + (Some(_), false) => { + return Err(unsupported( + "ROLLUP, CUBE and GROUPING SETS items with WITH ROLLUP, WITH CUBE or GROUPING SETS", + g.span, + )); + } + (None, false) => {} + } + let mut exprs = Vec::new(); + let mut product: Vec>> = vec![Vec::new()]; + for item in &g.items { + let sets = self.item_sets(item, &mut exprs)?; + if product.len().saturating_mul(sets.len()) > MAX_GROUPING_SETS { + return Err(unsupported( + &format!("more than {MAX_GROUPING_SETS} grouping sets"), + g.span, + )); + } + product = product + .iter() + .flat_map(|head| { + sets.iter().map(move |tail| { + head.iter().chain(tail).cloned().collect::>() + }) + }) + .collect(); + } + Ok(Grouping::GroupingSets { + exprs, + sets: product, + }) + } + + /// The grouping sets one `GROUP BY` item stands for; the expressions it + /// names are appended to `written` in order. + fn item_sets( + &mut self, + item: &GroupItem, + written: &mut Vec>, + ) -> R>>> { + Ok(match &item.kind { + GroupItemKind::Expr(e) => { + let e = arc(self.expr(e)?); + written.push(Arc::clone(&e)); + vec![vec![e]] + } + GroupItemKind::Set(s) => { + let set = self.set(s)?; + written.extend(set.iter().cloned()); + vec![set] + } + GroupItemKind::GroupingSets(items) => { + let mut sets = Vec::new(); + for item in items { + sets.extend(self.item_sets(item, written)?); + } + sets + } + GroupItemKind::Rollup(elements) | GroupItemKind::Cube(elements) => { + let rollup = matches!(item.kind, GroupItemKind::Rollup(_)); + let mut parts = Vec::with_capacity(elements.len()); + for element in elements { + let part = self.set(element)?; + written.extend(part.iter().cloned()); + parts.push(part); + } + // `ROLLUP (a, b)`: (a, b), (a), (); `CUBE (a, b)`: every + // subset of the elements. + let n = parts.len(); + if !rollup && n >= usize::BITS as usize - 1 + || !rollup && (1_usize << n) > MAX_GROUPING_SETS + { + return Err(unsupported( + &format!("more than {MAX_GROUPING_SETS} grouping sets"), + item.span, + )); + } + let chosen: Vec> = if rollup { + (0..=n).rev().map(|k| (0..k).collect()).collect() + } else { + (0..1_usize << n) + .rev() + .map(|m| (0..n).filter(|i| m >> (n - 1 - i) & 1 == 1).collect()) + .collect() + }; + chosen + .into_iter() + .map(|c| c.into_iter().flat_map(|i| parts[i].iter().cloned()).collect()) + .collect() + } + }) + } + + /// The expressions of one grouping set. + fn set(&mut self, s: &GroupingSet) -> R>> { + s.exprs.iter().map(|e| Ok(arc(self.expr(e)?))).collect() + } + fn from(&mut self, f: &FromClause) -> R { if f.pivot.is_some() || f.unpivot.is_some() { return Err(unsupported("PIVOT and UNPIVOT", f.span)); From 7bb422483cacecf444a1a884e5a4e967cae8dfc6 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:59:20 +0200 Subject: [PATCH 2/9] feat(sql): ROLLUP, CUBE, GROUPING SETS, grouping() and grouping_id() via Expand Spark's ResolveGroupingAnalytics shape: an Expand plan node (plan data, streamed one copy per projection) feeds an ordinary aggregate. Adds the DataFrame rollup/cube/grouping_sets, Spark golden tests, optimizer cases and commutative semantic keys. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/lib.rs | 9 +- crates/krishiv-exec/src/operators/expand.rs | 13 +- crates/krishiv-exec/src/operators/tests.rs | 124 ++++++ .../krishiv-exec/src/physical/instantiate.rs | 4 +- crates/krishiv-exec/src/physical/mod.rs | 4 +- crates/krishiv-exec/src/planner/build.rs | 5 +- crates/krishiv-optimizer/src/prune.rs | 5 +- crates/krishiv-optimizer/tests/it/rules.rs | 39 ++ .../it__rules__expand_cube_pruned.snap | 7 + .../it__rules__expand_duplicate_sets.snap | 7 + .../it__rules__expand_grouping_sets.snap | 8 + .../snapshots/it__rules__expand_rollup.snap | 7 + .../tests/it/support/eval.rs | 24 ++ crates/krishiv-plan/src/analyzer/analytics.rs | 127 ++++-- crates/krishiv-plan/src/analyzer/exprs.rs | 38 +- crates/krishiv-plan/src/analyzer/ops.rs | 12 +- crates/krishiv-plan/src/tests/analyzer.rs | 4 +- crates/krishiv-plan/src/tests/inventory.rs | 9 +- crates/krishiv-plan/src/tests/props.rs | 17 +- .../tests/grouping_analytics.rs | 326 +++++++++++++++ .../tests/grouping_analytics_api.rs | 381 ++++++++++++++++++ .../krishiv-session/tests/scratch_rollup.rs | 84 ---- crates/krishiv-sql/src/lower.rs | 22 +- ...cts_lower_or_are_named_as_unsupported.snap | 4 +- 24 files changed, 1126 insertions(+), 154 deletions(-) create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap create mode 100644 crates/krishiv-session/tests/grouping_analytics.rs create mode 100644 crates/krishiv-session/tests/grouping_analytics_api.rs delete mode 100644 crates/krishiv-session/tests/scratch_rollup.rs diff --git a/crates/krishiv-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index b9081da9..d189d9e1 100644 --- a/crates/krishiv-exec/src/lib.rs +++ b/crates/krishiv-exec/src/lib.rs @@ -113,10 +113,11 @@ pub use operator::{ FlushOutput, GlobalState, LocalState, Operator, OperatorOutput, Sink, SinkStatus, Source, }; pub use operators::{ - DEFAULT_MAX_MERGE_FAN_IN, Expand, NullOf, FileScan, FileScanMetrics, FileScanSource, FileSelections, - FileWrite, FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, MemoryScan, - Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RowStageSettings, RuntimeFilters, - ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, SortMetrics, TopN, UnionAlign, + DEFAULT_MAX_MERGE_FAN_IN, Expand, FileScan, FileScanMetrics, FileScanSource, FileSelections, + FileWrite, FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, + MemoryScan, NullOf, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, + RowStageSettings, RuntimeFilters, ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, + SortMetrics, TopN, UnionAlign, }; pub use pipeline::{ExecNode, PipelineBuilder}; pub use runner::{QueryRunner, Scheduler}; diff --git a/crates/krishiv-exec/src/operators/expand.rs b/crates/krishiv-exec/src/operators/expand.rs index 5f0a5c23..aedb7a71 100644 --- a/crates/krishiv-exec/src/operators/expand.rs +++ b/crates/krishiv-exec/src/operators/expand.rs @@ -39,14 +39,11 @@ impl Expand { for p in &projects { let (a, b) = (p.schema(), &schema); let same = a.fields().len() == b.fields().len() - && a.fields() - .iter() - .zip(b.fields()) - .all(|(x, y)| { - x.name() == y.name() - && x.data_type() == y.data_type() - && x.is_nullable() == y.is_nullable() - }); + && a.fields().iter().zip(b.fields()).all(|(x, y)| { + x.name() == y.name() + && x.data_type() == y.data_type() + && x.is_nullable() == y.is_nullable() + }); if !same { return Err(internal(format!( "the projections of an expansion differ: {:?} and {:?}", diff --git a/crates/krishiv-exec/src/operators/tests.rs b/crates/krishiv-exec/src/operators/tests.rs index 2f344aff..d628fb75 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 728028a0..58d5971f 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -30,8 +30,8 @@ use crate::join::{ }; use crate::operator::{Operator, Sink, Source}; use crate::operators::{ - Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project, ProjectItem, - PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, + Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project, + ProjectItem, PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, }; use crate::pipeline::{ExecNode, PipelineBuilder}; use crate::planner::TableProviderResolver; diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index 2ca8958d..9d4833bf 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -50,8 +50,8 @@ pub use data::{PlanData, assert_plan_data}; pub(crate) use instantiate::node_ids; pub use instantiate::{SpillScopes, WorkerContext, instantiate}; pub(crate) use instantiate::{ - Unbuilt, Values, aggregate_node, aggregate_spec, filter_op, hash_join, nested_loop_join, - expand_op, project_op, replicate_op, scan_filter, sort_op, top_n_op, + Unbuilt, Values, aggregate_node, aggregate_spec, expand_op, filter_op, hash_join, + nested_loop_join, project_op, replicate_op, scan_filter, sort_op, top_n_op, }; use crate::aggregate::{AggregateExpr, AggregateOutput}; diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index e9fcc721..a004d431 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -775,7 +775,10 @@ impl<'a> Planner<'a> { prepared, op.schema(), layout, - ("Expand", format!("{} projections [{text}]", projections.len())), + ( + "Expand", + format!("{} projections [{text}]", projections.len()), + ), |input, layout, subqueries| PhysNode::Expand { input, layout, diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index 7fb5f23a..332c99b1 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -226,7 +226,10 @@ impl Rewrite<'_> { let node = PlanNode::Expand { input: Arc::clone(input), projections: projections.iter().map(pick).collect(), - output: kept.iter().filter_map(|&j| output.get(j).cloned()).collect(), + output: kept + .iter() + .filter_map(|&j| output.get(j).cloned()) + .collect(), }; replace(&rebuilt, node) } diff --git a/crates/krishiv-optimizer/tests/it/rules.rs b/crates/krishiv-optimizer/tests/it/rules.rs index b69d424b..f797e7a4 100644 --- a/crates/krishiv-optimizer/tests/it/rules.rs +++ b/crates/krishiv-optimizer/tests/it/rules.rs @@ -398,6 +398,23 @@ const CASES: &[(&str, &str)] = &[ "prune_exists", "SELECT a FROM t WHERE EXISTS (SELECT * FROM u WHERE u.a = t.a)", ), + // Grouping sets: the expansion keeps what the aggregate reads. + ( + "expand_rollup", + "SELECT a, sum(c) FROM t GROUP BY ROLLUP(a, b)", + ), + ( + "expand_cube_pruned", + "SELECT count(*) FROM t GROUP BY CUBE(a, b)", + ), + ( + "expand_grouping_sets", + "SELECT a, grouping(a), grouping_id(a, b), sum(c) FROM t WHERE c > 1 GROUP BY GROUPING SETS ((a), (b), ())", + ), + ( + "expand_duplicate_sets", + "SELECT a, sum(c) FROM t GROUP BY GROUPING SETS ((a), (a))", + ), ]; #[test] @@ -443,6 +460,28 @@ fn rules() { } } +#[test] +fn an_expansion_has_a_copy_of_its_rows_per_projection() { + let ctx = support::ctx(); + let rows = |sql: &str| { + let plan = support::resolve(sql, &ctx); + let optimized = optimize(&plan).unwrap(); + let mut node = Arc::clone(optimized.root()); + while !matches!(node.node, PlanNode::Expand { .. }) { + let mut next = None; + node.for_each_child(&mut |c| next = Some(Arc::clone(c))); + node = next.expect("an expansion below"); + } + krishiv_optimizer::estimated_rows(&node, &krishiv_optimizer::NoStatistics) + }; + let one_set = rows("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))"); + let four_sets = rows("SELECT count(*) FROM t GROUP BY CUBE(a, b)"); + assert!( + (four_sets / one_set - 2.0).abs() < 1e-9, + "{four_sets} vs {one_set}" + ); +} + #[test] fn top_n_is_detected() { let ctx = support::ctx(); diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap new file mode 100644 index 00000000..24e2e4dc --- /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 00000000..d5016c87 --- /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 00000000..b3333bb8 --- /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 00000000..9c75b03d --- /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 bc73a19b..e0d75ce2 100644 --- a/crates/krishiv-optimizer/tests/it/support/eval.rs +++ b/crates/krishiv-optimizer/tests/it/support/eval.rs @@ -177,6 +177,28 @@ impl Interp<'_> { rows, } } + PlanNode::Expand { + input, + projections, + output, + } => { + let rel = self.plan(input, outer); + let ix = index(&rel.cols); + let rows = rel + .rows + .iter() + .flat_map(|r| { + projections + .iter() + .map(|row| row.iter().map(|e| self.eval(e, &ix, r, outer)).collect()) + .collect::>() + }) + .collect(); + Rel { + cols: output.iter().map(|c| c.id).collect(), + rows, + } + } PlanNode::Filter { input, predicate } => { let mut preds = Vec::new(); conjuncts(predicate, &mut preds); @@ -746,6 +768,8 @@ fn apply(name: &str, a: &[V]) -> V { "*" => num(|x, y| Some(x * y / SCALE)), "/" => num(|x, y| (y != 0).then(|| x * SCALE / y)), "%" => num(|x, y| (y != 0).then(|| x % y)), + // Integral division: the quotient truncated to a whole number. + "div" => num(|x, y| (y != 0).then(|| x / y * SCALE)), "negative" => match &a[0] { V::Num(x) => V::Num(-x), _ => V::Null, diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs index c6d34c79..8173c606 100644 --- a/crates/krishiv-plan/src/analyzer/analytics.rs +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -47,6 +47,10 @@ use crate::{ /// `BIGINT` with one bit each. const MAX_GROUPING: usize = 63; +/// The most expressions of a `CUBE`: 2^16 grouping sets (the SQL front end +/// bounds `CUBE (..)` items the same way). +const MAX_CUBE_EXPRS: usize = 16; + /// `grouping` or `grouping_id`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum GroupingFn { @@ -104,6 +108,8 @@ pub(super) struct Analytics { gid: Option, /// The expansion; `None` for a grand total. pub(super) expand: Option, + /// Whether the expansion numbers its sets (some two are equal). + numbered: bool, } impl Analytics { @@ -117,6 +123,17 @@ impl Analytics { .collect() } + /// `HAVING` and `ORDER BY` over numbered sets: Spark reads the last + /// grouping key as the grouping id there, and that key is the set's + /// number, so a `grouping` or `grouping_id` call is + /// `UNSUPPORTED_GROUPING_EXPRESSION`. + pub(super) fn check_late_call(&self, e: &Arc) -> Result<(), AnalysisError> { + if !self.numbered { + return Ok(()); + } + reject_grouping_calls(std::slice::from_ref(e)) + } + /// The position of `e` among the grouping expressions: a copy of one, or /// an expression equal to one. fn position(&self, e: &Expr, memo: &mut HashMap) -> Option { @@ -130,6 +147,38 @@ impl Analytics { } } +/// `UNSUPPORTED_GROUPING_EXPRESSION` if any of `exprs` calls `grouping` or +/// `grouping_id`. Spark has no source position for it. +pub(super) fn reject_grouping_calls(exprs: &[Arc]) -> Result<(), AnalysisError> { + let mut seen = std::collections::HashSet::new(); + let mut stack: Vec> = exprs.to_vec(); + while let Some(x) = stack.pop() { + if !seen.insert(Arc::as_ptr(&x) as usize) { + continue; + } + if matches!(&x.node, ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some()) { + return Err(AnalysisError::new( + UnsupportedGroupingExpression {}, + Origin::UNKNOWN, + )); + } + x.for_each_child(&mut |c| stack.push(Arc::clone(c))); + } + Ok(()) +} + +/// A resolved call `e`, stated not nullable. +fn not_null(e: &Arc) -> Arc { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) => { + let mut f = f.clone(); + f.nullable = false; + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function(f))).with_origin(e.origin)) + } + _ => Arc::clone(e), + } +} + fn literal(l: Literal) -> Arc { Arc::new(Expr::new(ExprNode::Literal(l))) } @@ -313,17 +362,25 @@ impl Resolver<'_> { aggregates: &[Arc], env: &Arc, ) -> Result { + // A cube has a set per subset of its expressions: bound it before + // building them (`WITH CUBE` and `DataFrame.cube` arrive here). + if let Grouping::Cube(v) = grouping + && v.len() > MAX_CUBE_EXPRS + { + return Err(AnalysisError::new( + GroupingSizeLimitExceeded { + max_size: (1_u64 << MAX_CUBE_EXPRS).to_string(), + }, + p.origin, + )); + } let c = self.child(input, env)?; let (flat, sets) = shape(grouping); self.check_star_ordinals(&flat, aggregates)?; // A single empty set over nothing is a grand total. let grand = flat.is_empty() && sets.len() == 1; - let gid_col = (!grand).then(|| { - self.attr( - "spark_grouping_id", - &arg_type(DataType::Long, false), - ) - }); + let gid_col = + (!grand).then(|| self.attr("spark_grouping_id", &arg_type(DataType::Long, false))); let gid = gid_col.as_ref().map(|col| Attr { col: BoundColumn { name: "grouping__id".to_owned(), @@ -375,9 +432,17 @@ impl Resolver<'_> { exprs: Vec::new(), gid: None, expand: None, + numbered: false, }; let items = self.rewrite_items(items, &an, env)?; - return self.build_aggregate(p.origin, &c, Vec::new(), &items, true, Some(Arc::new(an))); + return self.build_aggregate( + p.origin, + &c, + Vec::new(), + &items, + true, + Some(Arc::new(an)), + ); } let Some(gid_col) = gid_col else { return Err(unsupported("grouping sets without an id", p.origin)); @@ -385,7 +450,11 @@ impl Resolver<'_> { let numbered = (0..sets.len()).any(|i| sets[..i].contains(&sets[i])); // The expansion's input: the child's columns, and the grouping // expressions that are not one of them. - let mut passed: Vec> = c.out.iter().map(|a| col_expr(&a.col, Origin::UNKNOWN)).collect(); + let mut passed: Vec> = c + .out + .iter() + .map(|a| col_expr(&a.col, Origin::UNKNOWN)) + .collect(); let mut passed_out: Vec = c.out.to_vec(); let mut value: Vec = Vec::with_capacity(n); let mut copies: Vec = Vec::with_capacity(n); @@ -406,10 +475,12 @@ impl Resolver<'_> { Some(a) => a.col.clone(), None => { let v = self.attr(&name, &b.ty); - passed.push(Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { - expr: Arc::clone(&b.expr), - output: v.clone(), - })))); + passed.push(Arc::new(Expr::new(ExprNode::Resolved( + ResolvedExpr::Alias { + expr: Arc::clone(&b.expr), + output: v.clone(), + }, + )))); passed_out.push(Attr { col: v.clone(), qual: Arc::from(Vec::new()), @@ -425,7 +496,9 @@ impl Resolver<'_> { key: keys_of[j], attr: Attr { col: copy, - qual: column.as_ref().map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), + qual: column + .as_ref() + .map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), source: column.and_then(|a| a.source), }, }); @@ -437,7 +510,8 @@ impl Resolver<'_> { } else { c.clone() }; - let pos_col = numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); + let pos_col = + numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); let mut output: Vec = passed_out.iter().map(|a| a.col.clone()).collect(); output.extend(copies.iter().map(|g| g.attr.col.clone())); output.push(gid_col.clone()); @@ -459,9 +533,9 @@ impl Resolver<'_> { } row.push(literal(Literal::Long(mask))); if pos_col.is_some() { - row.push(literal(Literal::Integer( - i32::try_from(si).map_err(|_| unsupported("that many grouping sets", p.origin))?, - ))); + row.push(literal(Literal::Integer(i32::try_from(si).map_err( + |_| unsupported("that many grouping sets", p.origin), + )?))); } projections.push(row); } @@ -492,6 +566,7 @@ impl Resolver<'_> { let an = Analytics { exprs: copies, gid, + numbered, expand: Some(ExpandPlan { plan: Arc::clone(&expand.plan), ids: expand.ids, @@ -619,8 +694,7 @@ impl Resolver<'_> { Some(g) => col_expr(&g.col, Origin::UNKNOWN), None => literal(Literal::Long(0)), }; - let positions: Vec> = - f.args.iter().map(|a| an.position(a, memo)).collect(); + let positions: Vec> = f.args.iter().map(|a| an.position(a, memo)).collect(); match GroupingFn::from_name(&f.name) { Some(GroupingFn::GroupingId) => { let exact = positions.len() == n @@ -647,7 +721,7 @@ impl Resolver<'_> { Ok(id) } _ => { - let (Some(arg), Some(Some(j))) = (f.args.first(), positions.first()) else { + let (Some(_), Some(Some(j))) = (f.args.first(), positions.first()) else { return Err(AnalysisError::new( GroupingColumnMismatch { grouping: f.args.first().map(|a| with_id(a)).unwrap_or_default(), @@ -661,18 +735,20 @@ impl Resolver<'_> { Origin::UNKNOWN, )); }; - let _ = arg; // Bit `n - 1 - j` of the id, as a `TINYINT`. let shift = n - 1 - j; + // The divisors are literals that are not 0, so the id's + // bit is never null: `grouping` is not nullable (Spark). let high = if shift == 0 { id } else { - self.call( + let high = self.call( "div", vec![id, literal(Literal::Long(1_i64 << shift))], env, Origin::UNKNOWN, - )? + )?; + not_null(&high) }; let bit = self.call( "%", @@ -680,12 +756,15 @@ impl Resolver<'_> { env, Origin::UNKNOWN, )?; + let bit = not_null(&bit); let cast = Arc::new(Expr::new(ExprNode::Cast { expr: bit, to: DataType::Byte, mode: CastMode::Default, })); - Ok(self.bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + Ok(self + .bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))? + .expr) } } } diff --git a/crates/krishiv-plan/src/analyzer/exprs.rs b/crates/krishiv-plan/src/analyzer/exprs.rs index b16b6929..3bf207a1 100644 --- a/crates/krishiv-plan/src/analyzer/exprs.rs +++ b/crates/krishiv-plan/src/analyzer/exprs.rs @@ -272,6 +272,8 @@ impl Resolver<'_> { stack.push(Step::Enter(rewritten, in_agg)); continue; } + // In `ORDER BY` Spark rejects the call before it resolves + // its argument (`HAVING` reports the argument first). let inner = in_agg || self.is_aggregate(&x); stack.push(Step::Exit(Arc::clone(&x), in_agg)); let mut kids = Vec::new(); @@ -1525,6 +1527,29 @@ fn pretty_literal(v: &Literal) -> String { } } +/// A binary operator whose operands commute (Spark's `Add`, `Multiply`, +/// `And`, `Or`, `EqualTo` and `EqualNullSafe`). +fn is_commutative(node: &ExprNode) -> bool { + let ExprNode::Resolved(ResolvedExpr::Function(f)) = node else { + return false; + }; + let op = match &f.id { + FunctionId::Builtin(b) => Some(*b), + FunctionId::Registry { name, .. } => Builtin::from_name(name), + }; + matches!( + op, + Some( + Builtin::Add + | Builtin::Multiply + | Builtin::And + | Builtin::Or + | Builtin::Eq + | Builtin::EqNullSafe + ) + ) +} + /// A semantic key: equal for expressions that compute the same value, /// whatever their origins, column spellings or (top-level) alias. pub(super) fn semantic_key(e: &Expr, memo: &mut HashMap) -> u128 { @@ -1557,8 +1582,17 @@ pub(super) fn semantic_key(e: &Expr, memo: &mut HashMap) -> u128 { other => format!("u{other:?}"), } .into_bytes(); - for c in kids { - b.extend_from_slice(&memo.get(&addr(c)).copied().unwrap_or(0).to_le_bytes()); + let mut hashes: Vec = kids + .iter() + .map(|c| memo.get(&addr(*c)).copied().unwrap_or(0)) + .collect(); + // Spark canonicalizes the operands of a commutative operator into + // one order, so `a + b` and `b + a` are the same expression. + if hashes.len() == 2 && is_commutative(&x.node) { + hashes.sort_unstable(); + } + for h in hashes { + b.extend_from_slice(&h.to_le_bytes()); } memo.insert(addr(x), XxHash3_128::oneshot(&b)); } diff --git a/crates/krishiv-plan/src/analyzer/ops.rs b/crates/krishiv-plan/src/analyzer/ops.rs index 8238aec0..aff4ece8 100644 --- a/crates/krishiv-plan/src/analyzer/ops.rs +++ b/crates/krishiv-plan/src/analyzer/ops.rs @@ -684,6 +684,7 @@ impl Resolver<'_> { } if let Some(an) = &info.analytics { for b in &mut bounds { + an.check_late_call(&b.expr)?; b.expr = self.resolve_grouping(&b.expr, an, env)?; } } @@ -960,7 +961,11 @@ impl Resolver<'_> { let (keys, scope): (Vec, Vec) = match &analytics { Some(a) => ( a.keys(), - c.out.iter().chain(a.exprs.iter().map(|g| &g.attr)).cloned().collect(), + c.out + .iter() + .chain(a.exprs.iter().map(|g| &g.attr)) + .cloned() + .collect(), ), None => ( group @@ -1400,6 +1405,11 @@ impl Resolver<'_> { .map(|(o, expr)| SortOrder { expr, ..o.clone() }) .collect() }; + // Over no grouping sets, `ORDER BY grouping(..)` is rejected before + // its argument resolves (`HAVING` reports the argument first). + if c.agg.as_ref().is_none_or(|a| a.analytics.is_none()) { + super::analytics::reject_grouping_calls(&exprs)?; + } let (bounds, below, extended) = if c.agg.is_some() { self.over_agg(&c, &exprs, env, "Sort")? } else { diff --git a/crates/krishiv-plan/src/tests/analyzer.rs b/crates/krishiv-plan/src/tests/analyzer.rs index c05b3ba6..647bec82 100644 --- a/crates/krishiv-plan/src/tests/analyzer.rs +++ b/crates/krishiv-plan/src/tests/analyzer.rs @@ -26,8 +26,8 @@ use crate::analyzer::{ }; use crate::walk::{Node, visit_distinct}; use crate::{ - ColumnRef, Cte, Expr, ExprNode, Grouping, JoinCondition, JoinKind, Literal, LogicalPlan, - MAX_DEPTH, Origin, Plan, PlanId, PlanNode, ReadSource, SourceId, SqlArgs, Star, TableVersion, + ColumnRef, Cte, Expr, ExprNode, JoinCondition, JoinKind, Literal, LogicalPlan, MAX_DEPTH, + Origin, Plan, PlanId, PlanNode, ReadSource, SourceId, SqlArgs, Star, TableVersion, }; fn goldens() -> Value { diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index fcd40199..862da651 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -12,8 +12,8 @@ use krishiv_types::DataType; use super::dsl::{c, int}; use super::tpch::QUERIES; use crate::{ - Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, - ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, + Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, + PlanNode, ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, }; /// The node kinds a plan uses, named as in the inventory. No wildcard arm: @@ -170,7 +170,10 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { ("filter", df(1).filter(c("a"))), ("where", df(1).filter(c("a"))), ("groupBy.agg", df(1).aggregate(vec![c("a")], vec![c("a")])), - ("rollup.agg", grouped(Grouping::Rollup(vec![Arc::new(c("a"))]))), + ( + "rollup.agg", + grouped(Grouping::Rollup(vec![Arc::new(c("a"))])), + ), ("cube.agg", grouped(Grouping::Cube(vec![Arc::new(c("a"))]))), ( "join", diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index ec013bb5..24bfa55a 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -426,7 +426,7 @@ impl Gen<'_> { }, 21 => PlanNode::Expand { input: self.input(d), - projections: self.many(1, 2, |g| g.exprs()), + projections: self.many(1, 2, Gen::exprs), output: self.many(0, 2, Gen::bound), }, _ => PlanNode::WithCte { @@ -824,12 +824,15 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { output, } => { site!(w, projections.push(Vec::new())); - site!(w, output.push(output.first().cloned().unwrap_or(BoundColumn { - id: AttrId(0), - name: "x".to_owned(), - data_type: DataType::Integer, - nullable: false, - }))); + site!( + w, + output.push(output.first().cloned().unwrap_or(BoundColumn { + id: AttrId(0), + name: "x".to_owned(), + data_type: DataType::Integer, + nullable: false, + })) + ); projections.iter_mut().any(|p| m_exprs(p, w)) || output.iter_mut().any(|c| m_bound(c, w)) || m_plan(input, w) diff --git a/crates/krishiv-session/tests/grouping_analytics.rs b/crates/krishiv-session/tests/grouping_analytics.rs new file mode 100644 index 00000000..a544dce4 --- /dev/null +++ b/crates/krishiv-session/tests/grouping_analytics.rs @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Spark 4.2 `SQLQueryTestSuite` goldens for grouping analytics (`ROLLUP`, `CUBE`, +//! `GROUPING SETS`, `grouping`, `grouping_id`): each statement of the vendored +//! `group-analytics` and `grouping_set` result files runs through a session and +//! its schema, rows (compared as sets of rows) or error class must match. +//! +//! The session has no `CREATE VIEW` or inline `VALUES` yet, so the driver turns +//! each literal `VALUES` table into a registered table of the same name. +#![forbid(unsafe_code)] + +use std::path::Path; +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Float64Type, Int8Type, Int16Type, Int32Type, Int64Type}; +use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use futures::executor::block_on; +use krishiv_common::error::KrishivError; +use krishiv_session::Session; + +type R = Result>; + +/// One golden statement: the query and its expected schema and output. +struct Case { + query: String, + schema: String, + output: String, +} + +fn load(file: &str) -> R> { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../tests/spark-golden/v4.2.0/results") + .join(file); + let text = std::fs::read_to_string(&path)?; + let cases: Vec = text + .split("-- !query\n") + .skip(1) + .filter_map(|block| { + let (query, rest) = block.split_once("\n-- !query schema\n")?; + let (schema, output) = rest.split_once("\n-- !query output\n")?; + Some(Case { + query: query.trim().to_owned(), + schema: schema.trim().to_owned(), + output: output.trim_end().to_owned(), + }) + }) + .collect(); + if cases.is_empty() { + return Err(format!("no statements in {file}").into()); + } + Ok(cases) +} + +fn spark_type(t: &DataType) -> String { + match t { + DataType::Int8 => "tinyint".into(), + DataType::Int16 => "smallint".into(), + DataType::Int32 => "int".into(), + DataType::Int64 => "bigint".into(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => "string".into(), + DataType::Boolean => "boolean".into(), + DataType::Float64 => "double".into(), + other => format!("{other}"), + } +} + +fn cell(c: &dyn Array, row: usize) -> String { + if c.is_null(row) { + return "NULL".into(); + } + match c.data_type() { + DataType::Int8 => c.as_primitive::().value(row).to_string(), + DataType::Int16 => c.as_primitive::().value(row).to_string(), + DataType::Int32 => c.as_primitive::().value(row).to_string(), + DataType::Int64 => c.as_primitive::().value(row).to_string(), + DataType::Utf8 => c.as_string::().value(row).to_owned(), + DataType::Boolean => c.as_boolean().value(row).to_string(), + DataType::Float64 => format!("{:?}", c.as_primitive::().value(row)), + other => format!("<{other}>"), + } +} + +fn render(schema: &Schema, batches: &[RecordBatch]) -> (String, Vec) { + let schema = schema + .fields() + .iter() + .map(|f| format!("{}:{}", f.name(), spark_type(f.data_type()))) + .collect::>() + .join(","); + let rows = batches + .iter() + .flat_map(|b| { + (0..b.num_rows()).map(|r| { + b.columns() + .iter() + .map(|c| cell(c.as_ref(), r)) + .collect::>() + .join("\t") + }) + }) + .collect(); + (format!("struct<{schema}>"), rows) +} + +/// A parsed `VALUES` literal table. +struct Table { + name: String, + columns: Vec, + rows: Vec>>, +} + +fn is_word(b: u8) -> bool { + b.is_ascii_alphanumeric() || b == b'_' +} + +fn skip_ws(bytes: &[u8], i: &mut usize) { + while bytes.get(*i).is_some_and(u8::is_ascii_whitespace) { + *i += 1; + } +} + +/// Find an inline `VALUES (..), (..) [AS] alias[(cols)]` (optionally wrapped in +/// parentheses) in `sql`; return the text with it replaced by the alias, and the table. +fn extract_values(sql: &str) -> Option<(String, Table)> { + let upper = sql.to_ascii_uppercase(); + let bytes = sql.as_bytes(); + let at = upper.match_indices("VALUES").map(|(i, _)| i).find(|&i| { + (i == 0 || bytes.get(i - 1).is_none_or(|&b| !is_word(b))) + && bytes.get(i + 6).is_none_or(|&b| !is_word(b)) + })?; + let before = sql[..at].trim_end(); + let wrapped = before.ends_with('('); + let start = if wrapped { before.len() - 1 } else { at }; + let mut i = at + 6; + let mut rows = Vec::new(); + loop { + skip_ws(bytes, &mut i); + if bytes.get(i) != Some(&b'(') { + return None; + } + let close = i + sql[i..].find(')')?; + rows.push( + sql[i + 1..close] + .split(',') + .map(|v| { + let v = v.trim(); + (!v.eq_ignore_ascii_case("NULL")).then(|| v.to_owned()) + }) + .collect::>(), + ); + i = close + 1; + skip_ws(bytes, &mut i); + if bytes.get(i) == Some(&b',') { + i += 1; + } else { + break; + } + } + if wrapped { + if bytes.get(i) != Some(&b')') { + return None; + } + i += 1; + skip_ws(bytes, &mut i); + } + if upper[i..].starts_with("AS ") { + i += 3; + skip_ws(bytes, &mut i); + } + let name_start = i; + while bytes.get(i).is_some_and(|&b| is_word(b)) { + i += 1; + } + let name = sql[name_start..i].to_owned(); + skip_ws(bytes, &mut i); + let mut columns = Vec::new(); + if bytes.get(i) == Some(&b'(') { + let close = i + sql[i..].find(')')?; + columns = sql[i + 1..close] + .split(',') + .map(|c| c.trim().to_owned()) + .collect(); + i = close + 1; + } + let text = format!("{}{}{}", &sql[..start], name, &sql[i..]); + Some(( + text, + Table { + name, + columns, + rows, + }, + )) +} + +fn register(s: &Session, t: &Table, name: &str) -> R { + let mut fields = Vec::new(); + let mut arrays: Vec = Vec::new(); + for (c, column) in t.columns.iter().enumerate() { + let col: Vec> = t + .rows + .iter() + .map(|r| r.get(c).and_then(Option::as_ref)) + .collect(); + let numeric = col.iter().flatten().all(|v| v.parse::().is_ok()); + let nullable = col.iter().any(Option::is_none); + if numeric { + fields.push(Field::new(column, DataType::Int32, nullable)); + arrays.push(Arc::new(Int32Array::from( + col.iter() + .map(|v| v.and_then(|v| v.parse().ok())) + .collect::>(), + ))); + } else { + fields.push(Field::new(column, DataType::Utf8, nullable)); + arrays.push(Arc::new(StringArray::from( + col.iter() + .map(|v| v.map(|v| v.trim_matches(['"', '\'']))) + .collect::>(), + ))); + } + } + let schema = Arc::new(Schema::new(fields)); + let batch = RecordBatch::try_new(Arc::clone(&schema), arrays)?; + s.register_batches(name, schema, vec![batch])?; + Ok(()) +} + +/// The error class a golden expects, from its JSON error body. +fn expected_class(output: &str) -> Option { + output + .lines() + .find_map(|l| l.trim().strip_prefix("\"errorClass\" : \"")) + .map(|l| l.trim_end_matches(['"', ',']).to_owned()) +} + +/// Run every statement of `file`; return the statements that disagree with Spark. +fn run(file: &str) -> R> { + let mut views: Vec<(String, Table)> = Vec::new(); + let mut bad = Vec::new(); + for case in load(file)? { + let upper = case.query.to_ascii_uppercase(); + if upper.starts_with("CREATE") { + let name = case + .query + .split_whitespace() + .skip_while(|w| !w.eq_ignore_ascii_case("VIEW")) + .nth(1); + let (Some(name), Some((_, table))) = (name, extract_values(&case.query)) else { + bad.push(format!("{}\n unsupported setup", case.query)); + continue; + }; + views.retain(|(n, _)| n != name); + views.push((name.to_owned(), table)); + continue; + } + if let Some(name) = upper.strip_prefix("DROP VIEW ") { + views.retain(|(n, _)| !n.eq_ignore_ascii_case(name.trim())); + continue; + } + let s = Session::new()?; + for (name, t) in &views { + register(&s, t, name)?; + } + let mut query = case.query.clone(); + if let Some((text, table)) = extract_values(&query) { + register(&s, &table, &table.name)?; + query = text; + } + let result = block_on(async { + let df = s.sql(&query).await?; + let schema = df.schema().await?; + Ok::<_, KrishivError>((schema, df.collect().await?)) + }); + if case.schema == "struct<>" { + let class = expected_class(&case.output); + match (result, class) { + (Err(e), Some(c)) if e.error_class().name() == c => {} + (Err(e), c) => bad.push(format!( + "{}\n error class {} expected {c:?}", + case.query, + e.error_class().name() + )), + (Ok(_), c) => bad.push(format!("{}\n expected error {c:?}", case.query)), + } + continue; + } + match result { + Err(e) => bad.push(format!("{}\n error {e}", case.query)), + Ok((schema, batches)) => { + let (schema, mut rows) = render(&schema, &batches); + let mut want: Vec = case.output.lines().map(str::to_owned).collect(); + // Rows tied under `ORDER BY` come in no fixed order. + rows.sort(); + want.sort(); + if schema != case.schema || rows != want { + bad.push(format!( + "{}\n schema {schema} want {}\n rows {rows:?}\n want {want:?}", + case.query, case.schema + )); + } + } + } + } + Ok(bad) +} + +fn check(file: &str) -> R { + let bad = run(file)?; + if bad.is_empty() { + Ok(()) + } else { + Err(format!("{} mismatches:\n{}", bad.len(), bad.join("\n")).into()) + } +} + +#[test] +fn group_analytics_golden() -> R { + check("group-analytics.sql.out") +} + +#[test] +fn grouping_set_golden() -> R { + check("grouping_set.sql.out") +} diff --git a/crates/krishiv-session/tests/grouping_analytics_api.rs b/crates/krishiv-session/tests/grouping_analytics_api.rs new file mode 100644 index 00000000..9c638143 --- /dev/null +++ b/crates/krishiv-session/tests/grouping_analytics_api.rs @@ -0,0 +1,381 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Grouping analytics through the session: the DataFrame forms (`rollup`, +//! `cube`, `grouping_sets`) against the SQL forms, rolled-up columns' types, +//! `NULL` data against rolled-up `NULL`, and the errors of invalid uses. +//! The Spark goldens are in `grouping_analytics.rs`. +#![forbid(unsafe_code)] + +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Int8Type, Int32Type, Int64Type}; +use arrow_array::{Array, Int32Array, RecordBatch}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use futures::executor::block_on; +use krishiv_common::error::{ErrorClass, KrishivError}; +use krishiv_session::functions::{col, lit, sum}; +use krishiv_session::{DataFrame, Session}; + +type R = Result>; + +/// Tables `t(a, b, v)` (no nulls) and `n(a, v)` (with a real NULL key). +fn session() -> R { + let s = Session::new()?; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2])), + Arc::new(Int32Array::from(vec![1, 2, 1])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + )?; + s.register_batches("t", schema, vec![batch])?; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("v", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(1), None, None])), + Arc::new(Int32Array::from(vec![5, 7, 11])), + ], + )?; + s.register_batches("n", schema, vec![batch])?; + Ok(s) +} + +fn cell(c: &dyn Array, row: usize) -> String { + if c.is_null(row) { + return "NULL".into(); + } + match c.data_type() { + DataType::Int8 => c.as_primitive::().value(row).to_string(), + DataType::Int32 => c.as_primitive::().value(row).to_string(), + DataType::Int64 => c.as_primitive::().value(row).to_string(), + other => format!("<{other}>"), + } +} + +fn row_strings(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|b| { + (0..b.num_rows()).map(|r| { + b.columns() + .iter() + .map(|c| cell(c.as_ref(), r)) + .collect::>() + .join("\t") + }) + }) + .collect() +} + +/// The schema and the rows (tab-separated, sorted) of a DataFrame. +fn run(df: &DataFrame) -> Result<(SchemaRef, Vec), KrishivError> { + block_on(async { + let schema = df.schema().await?; + let mut rows = row_strings(&df.collect().await?); + rows.sort(); + Ok((schema, rows)) + }) +} + +fn sql(s: &Session, query: &str) -> Result<(SchemaRef, Vec), KrishivError> { + run(&block_on(s.sql(query))?) +} + +fn rows(s: &Session, query: &str) -> R> { + Ok(sql(s, query)?.1) +} + +fn sorted(rows: &[&str]) -> Vec { + let mut rows: Vec = rows.iter().map(|r| (*r).to_owned()).collect(); + rows.sort(); + rows +} + +fn error_class(s: &Session, query: &str) -> R { + match sql(s, query) { + Err(e) => Ok(e.error_class()), + Ok((_, rows)) => { + Err(format!("{query} should fail, it returned {} rows", rows.len()).into()) + } + } +} + +#[test] +fn dataframe_rollup_cube_and_grouping_sets_match_their_sql() -> R { + let s = session()?; + let sum_v = || [sum(col("v"))]; + let (schema, rollup) = run(&s.table("t").rollup([col("a"), col("b")]).agg(sum_v()))?; + assert_eq!( + rollup, + sorted(&[ + "1\t1\t10", + "1\t2\t20", + "2\t1\t30", + "1\tNULL\t30", + "2\tNULL\t30", + "NULL\tNULL\t60", + ]) + ); + let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, ["a", "b", "sum(v)"]); + assert_eq!( + rollup, + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY ROLLUP(a, b)")? + ); + + let (_, cube) = run(&s.table("t").cube([col("a"), col("b")]).agg(sum_v()))?; + assert_eq!( + cube, + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY CUBE(a, b)")? + ); + assert_eq!(cube.len(), 8); + + let sets = s + .table("t") + .grouping_sets( + [vec![col("a")], vec![col("b")], vec![]], + [col("a"), col("b")], + ) + .agg(sum_v()); + let (_, sets) = run(&sets)?; + assert_eq!( + sets, + sorted(&[ + "1\tNULL\t30", + "2\tNULL\t30", + "NULL\t1\t40", + "NULL\t2\t20", + "NULL\tNULL\t60", + ]) + ); + assert_eq!( + sets, + rows( + &s, + "SELECT a, b, sum(v) FROM t GROUP BY GROUPING SETS ((a), (b), ())" + )? + ); + Ok(()) +} + +#[test] +fn rolled_up_columns_become_nullable_and_grouping_is_a_tinyint() -> R { + let s = session()?; + let (schema, _) = sql( + &s, + "SELECT a, b, sum(v), grouping(a), grouping_id(a, b) FROM t GROUP BY ROLLUP(a, b)", + )?; + let fields: Vec<(String, DataType, bool)> = schema + .fields() + .iter() + .map(|f| (f.name().clone(), f.data_type().clone(), f.is_nullable())) + .collect(); + assert_eq!( + fields, + [ + ("a".to_owned(), DataType::Int32, true), + ("b".to_owned(), DataType::Int32, true), + ("sum(v)".to_owned(), DataType::Int64, true), + ("grouping(a)".to_owned(), DataType::Int8, false), + ("grouping_id(a, b)".to_owned(), DataType::Int64, false), + ] + ); + Ok(()) +} + +#[test] +fn grouping_tells_a_null_key_from_a_rolled_up_one() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, grouping(a), sum(v) FROM n GROUP BY ROLLUP(a)" + )?, + sorted(&["1\t0\t5", "NULL\t0\t18", "NULL\t1\t23"]) + ); + Ok(()) +} + +#[test] +fn grouping_id_sets_bit_n_minus_1_minus_position() -> R { + let s = session()?; + // `grouping_id(a, b)`: `a` is the high bit. + assert_eq!( + rows( + &s, + "SELECT a, b, grouping_id(a, b) FROM t GROUP BY GROUPING SETS ((a), (b), (a, b))" + )?, + sorted(&[ + "1\tNULL\t1", + "2\tNULL\t1", + "NULL\t1\t2", + "NULL\t2\t2", + "1\t1\t0", + "1\t2\t0", + "2\t1\t0", + ]) + ); + Ok(()) +} + +#[test] +fn mixed_plain_keys_and_rollup() -> R { + let s = session()?; + assert_eq!( + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY a, ROLLUP(b)")?, + sorted(&[ + "1\t1\t10", + "1\t2\t20", + "2\t1\t30", + "1\tNULL\t30", + "2\tNULL\t30" + ]) + ); + Ok(()) +} + +#[test] +fn operands_of_a_commutative_operator_are_one_grouping_expression() -> R { + let s = session()?; + // `b + a` is `a + b`: the select item reads the grouping copy, and the + // two sets are equal, so each group comes twice. + assert_eq!( + rows( + &s, + "SELECT b + a, sum(v) FROM t GROUP BY GROUPING SETS ((a + b), (b + a))" + )?, + sorted(&["2\t10", "2\t10", "3\t50", "3\t50"]) + ); + Ok(()) +} + +#[test] +fn empty_set_is_a_grand_total_even_over_no_rows() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT count(*), grouping_id() FROM t WHERE v > 100 GROUP BY GROUPING SETS (())" + )?, + ["0\t0"] + ); + assert_eq!( + rows( + &s, + "SELECT a, count(*) FROM t WHERE v > 100 GROUP BY ROLLUP(a)" + )?, + Vec::::new(), + "a rollup over no rows has no grand-total row, as in Spark" + ); + Ok(()) +} + +#[test] +fn having_and_order_by_read_grouping() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, sum(v) FROM t GROUP BY ROLLUP(a) HAVING grouping(a) = 1" + )?, + ["NULL\t60"] + ); + // `ORDER BY` follows the id: the grand total (id 1) comes first, descending. + let df = block_on( + s.sql("SELECT a, sum(v) FROM t GROUP BY ROLLUP(a) ORDER BY grouping_id(a) DESC, a"), + )?; + assert_eq!( + row_strings(&block_on(df.collect())?), + ["NULL\t60", "1\t30", "2\t30"] + ); + Ok(()) +} + +#[test] +fn invalid_uses_have_spark_error_classes() -> R { + let s = session()?; + let unsupported = ErrorClass::UnsupportedGroupingExpression; + // Outside a grouping aggregate, in any clause. + assert_eq!(error_class(&s, "SELECT grouping(a) FROM t")?, unsupported); + assert_eq!( + error_class(&s, "SELECT a FROM t GROUP BY a HAVING grouping(a) = 0")?, + unsupported + ); + assert_eq!( + error_class(&s, "SELECT count(*) FROM t ORDER BY grouping(a)")?, + unsupported + ); + // An argument that is no grouping expression, and a mismatched id list. + assert_eq!( + error_class(&s, "SELECT grouping(v) FROM t GROUP BY ROLLUP(a)")?, + ErrorClass::GroupingColumnMismatch + ); + assert_eq!( + error_class(&s, "SELECT grouping_id(b, a) FROM t GROUP BY ROLLUP(a, b)")?, + ErrorClass::GroupingIdColumnMismatch + ); + assert_eq!( + error_class(&s, "SELECT grouping(a, b) FROM t GROUP BY ROLLUP(a, b)")?, + ErrorClass::WrongNumArgsWithoutSuggestion + ); + // 64 grouping expressions do not fit the 64-bit id. + let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); + let query = format!( + "SELECT count(*) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ); + assert_eq!( + error_class(&s, &query)?, + ErrorClass::GroupingSizeLimitExceeded + ); + // A cube has 2^n sets: 17 expressions are too many, 16 still are not. + let cube = |n: usize| { + let keys: Vec = (0..n).map(|i| format!("a + {i}")).collect(); + format!( + "SELECT count(*) FROM t GROUP BY {} WITH CUBE", + keys.join(", ") + ) + }; + assert_eq!( + error_class(&s, &cube(17))?, + ErrorClass::GroupingSizeLimitExceeded + ); + let keys: Vec<_> = (0..17).map(|i| col("a").gt(lit(i))).collect(); + let df = s.table("t").cube(keys).agg([sum(col("v"))]); + assert_eq!( + run(&df).err().map(|e| e.error_class()), + Some(ErrorClass::GroupingSizeLimitExceeded) + ); + Ok(()) +} + +#[test] +fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, sum(v) FROM t GROUP BY GROUPING SETS ((a), (a))" + )?, + sorted(&["1\t30", "1\t30", "2\t30", "2\t30"]) + ); + // Spark reads the set's number where `HAVING` expects the id. + assert_eq!( + error_class( + &s, + "SELECT a FROM t GROUP BY GROUPING SETS ((a), (a)) HAVING grouping(a) = 0" + )?, + ErrorClass::UnsupportedGroupingExpression + ); + Ok(()) +} diff --git a/crates/krishiv-session/tests/scratch_rollup.rs b/crates/krishiv-session/tests/scratch_rollup.rs deleted file mode 100644 index f12c6053..00000000 --- a/crates/krishiv-session/tests/scratch_rollup.rs +++ /dev/null @@ -1,84 +0,0 @@ -//! scratch -#![cfg(test)] -use std::sync::Arc; - -use arrow_array::{Int32Array, RecordBatch, StringArray}; -use arrow_schema::{DataType as T, Field, Schema}; -use futures::executor::block_on; -use krishiv_session::Session; - -#[test] -fn scratch() { - let s = Session::new().unwrap(); - let schema = Arc::new(Schema::new(vec![ - Field::new("a", T::Int32, false), - Field::new("b", T::Int32, false), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(vec![1, 1, 2, 2, 3, 3])), - Arc::new(Int32Array::from(vec![1, 2, 1, 2, 1, 2])), - ], - ) - .unwrap(); - s.register_batches("td", schema, vec![batch]).unwrap(); - let schema = Arc::new(Schema::new(vec![ - Field::new("course", T::Utf8, false), - Field::new("year", T::Int32, false), - Field::new("earnings", T::Int32, false), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(StringArray::from(vec!["dotNET", "Java", "dotNET", "dotNET", "Java"])), - Arc::new(Int32Array::from(vec![2012, 2012, 2012, 2013, 2013])), - Arc::new(Int32Array::from(vec![10000, 20000, 5000, 48000, 30000])), - ], - ) - .unwrap(); - s.register_batches("cs", schema, vec![batch]).unwrap(); - for q in std::fs::read_to_string("/home/krishiv/krishiv-data/tmp/s-rollup/q.sql") - .unwrap() - .split(";\n") - .filter(|q| !q.trim().is_empty()) - { - match block_on(s.sql(q)) { - Err(e) => println!("ERR {q}\n {e}"), - Ok(df) => match block_on(df.collect()) { - Err(e) => println!("RUNERR {q}\n {e}"), - Ok(b) => { - println!("OK {q}"); - for b in &b { - println!("{}", arrow_cast_fmt(b)); - } - } - }, - } - } -} - -fn arrow_cast_fmt(b: &RecordBatch) -> String { - let mut out = format!("{:?}\n", b.schema().fields().iter().map(|f| format!("{}:{}{}", f.name(), f.data_type(), if f.is_nullable() {"?"} else {""})).collect::>()); - for r in 0..b.num_rows() { - let cells: Vec = b - .columns() - .iter() - .map(|c| { - if c.is_null(r) { - "NULL".to_owned() - } else { - let d = arrow_array::cast::AsArray::as_primitive_opt::(c) - .map(|a| a.value(r).to_string()) - .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) - .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) - .or_else(|| arrow_array::cast::AsArray::as_string_opt::(c).map(|a| a.value(r).to_owned())); - d.unwrap_or_else(|| "?".to_owned()) - } - }) - .collect(); - out.push_str(&cells.join("\t")); - out.push('\n'); - } - out -} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index 8daa1db3..f14c2e53 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -28,9 +28,9 @@ use krishiv_types::DecimalType; use crate::ast::{ Alias, BinaryOp, Expr, ExprKind, FromClause, GroupBy, GroupItem, GroupItemKind, GroupWith, - GroupingSet, JoinConstraint, JoinKind, Literal, - NullTreatment, NumKind, Query, Relation, RelationExt, Select, SetExpr, SetExprKind, - SetOperator, SetQuantifier, SortItem, Span, TableFactor, TableFactorKind, Truth, UnaryOp, + GroupingSet, JoinConstraint, JoinKind, Literal, NullTreatment, NumKind, Query, Relation, + RelationExt, Select, SetExpr, SetExprKind, SetOperator, SetQuantifier, SortItem, Span, + TableFactor, TableFactorKind, Truth, UnaryOp, }; use crate::validate::literal::{self, Number, TypedKind}; use crate::{ParseError, ParserOptions, parse_query}; @@ -452,10 +452,7 @@ impl Lowerer { (Some(GroupWith::Rollup), true) => return Ok(Grouping::Rollup(plain)), (Some(GroupWith::Cube), true) => return Ok(Grouping::Cube(plain)), (Some(GroupWith::GroupingSets(sets)), true) => { - let sets = sets - .iter() - .map(|s| self.set(s)) - .collect::>>()?; + let sets = sets.iter().map(|s| self.set(s)).collect::>>()?; return Ok(Grouping::GroupingSets { exprs: plain, sets }); } (Some(_), false) => { @@ -479,9 +476,8 @@ impl Lowerer { product = product .iter() .flat_map(|head| { - sets.iter().map(move |tail| { - head.iter().chain(tail).cloned().collect::>() - }) + sets.iter() + .map(move |tail| head.iter().chain(tail).cloned().collect::>()) }) .collect(); } @@ -545,7 +541,11 @@ impl Lowerer { }; chosen .into_iter() - .map(|c| c.into_iter().flat_map(|i| parts[i].iter().cloned()).collect()) + .map(|c| { + c.into_iter() + .flat_map(|i| parts[i].iter().cloned()) + .collect() + }) .collect() } }) diff --git a/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap b/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap index 2e9c27b9..ba9dd242 100644 --- a/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap +++ b/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap @@ -21,9 +21,9 @@ SELECT a FROM t LATERAL VIEW explode(b) x AS y SELECT row_number() OVER (ORDER BY a) FROM t => [KRISHIV_UNSUPPORTED_FEATURE] The feature window functions is not supported by Krishiv yet. SQLSTATE: 0A000 @ 7..37 SELECT a FROM t GROUP BY ROLLUP (a) - => [KRISHIV_UNSUPPORTED_FEATURE] The feature ROLLUP, CUBE and GROUPING SETS is not supported by Krishiv yet. SQLSTATE: 0A000 @ 25..35 + => Aggregate [a] GROUPING SETS ((a), ()), [a] | +- Read t | SELECT a FROM t GROUP BY a WITH CUBE - => [KRISHIV_UNSUPPORTED_FEATURE] The feature ROLLUP, CUBE and GROUPING SETS is not supported by Krishiv yet. SQLSTATE: 0A000 @ 16..36 + => Aggregate CUBE [a], [a] | +- Read t | SELECT * FROM t NATURAL JOIN u => [KRISHIV_UNSUPPORTED_FEATURE] The feature NATURAL JOIN is not supported by Krishiv yet. SQLSTATE: 0A000 @ 16..30 SELECT * FROM t, LATERAL (SELECT a) From a41e41eb2161828041829a9e11a3b15f7a68d771 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:57:22 +0200 Subject: [PATCH 3/9] Address review threads for ROLLUP, CUBE and GROUPING SETS Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-common/src/stack.rs | 2 +- crates/krishiv-optimizer/src/stats.rs | 100 +++++++++++++++++- crates/krishiv-optimizer/tests/it/rules.rs | 48 ++++++++- crates/krishiv-plan/src/analyzer/analytics.rs | 100 +++++++++++++++--- .../tests/grouping_analytics.rs | 34 +++++- .../tests/grouping_analytics_api.rs | 48 +++++++-- crates/krishiv-sql/src/lower.rs | 2 +- crates/krishiv-sql/tests/lower.rs | 22 ++++ 8 files changed, 321 insertions(+), 35 deletions(-) diff --git a/crates/krishiv-common/src/stack.rs b/crates/krishiv-common/src/stack.rs index e44a3097..2a0488ec 100644 --- a/crates/krishiv-common/src/stack.rs +++ b/crates/krishiv-common/src/stack.rs @@ -7,7 +7,7 @@ //! never needs more than `STEP` frames of the caller's stack. /// Levels a recursive pass may take on one stack segment. -pub const STEP: usize = 64; +pub const STEP: usize = 32; /// Runs `f` on a new thread with a 64 MiB stack (16 MiB in release /// builds), waiting for it; `None` when no thread can be started. A panic diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index d96d5abf..8e5c5210 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -241,11 +241,107 @@ impl<'a> Estimator<'a> { | PlanNode::ReplicateRows { input, .. } => Est::clone(&self.plan(input)), // Every input row once per projection. PlanNode::Expand { - input, projections, .. + input, + projections, + output, } => { let est = Est::clone(&self.plan(input)); let rows = est.rows * to_f64(u64::try_from(projections.len()).unwrap_or(u64::MAX)); - est.with_rows(rows) + let rows = rows.max(1.0); + if projections.is_empty() + || output.is_empty() + || projections.iter().any(|p| p.len() != output.len()) + { + return est.with_rows(rows); + } + let per = est.rows; + let mut cols = HashMap::with_capacity(output.len()); + for (j, out) in output.iter().enumerate() { + let mut null_rows = 0.0; + let mut col_id: Option = None; + let mut mixed = false; + let mut lits: Vec<&Literal> = Vec::new(); + let mut other = false; + for proj in projections.iter() { + let s = strip_alias(&proj[j]); + match &s.node { + ExprNode::Literal(Literal::Null(_)) => null_rows += per, + ExprNode::Literal(l) => { + if col_id.is_some() { + mixed = true; + } + if !lits.contains(&l) { + lits.push(l); + } + } + ExprNode::Resolved( + ResolvedExpr::Column(c) | ResolvedExpr::OuterColumn(c), + ) => match est.cols.get(&c.id) { + Some(col) => { + null_rows += col.nulls * per; + match col_id { + None => col_id = Some(c.id), + Some(id) if id == c.id => {} + Some(_) => mixed = true, + } + if !lits.is_empty() { + mixed = true; + } + } + None => other = true, + }, + _ => other = true, + } + } + let nulls = (null_rows / rows).clamp(0.0, 1.0); + let distinct_lits = to_f64(u64::try_from(lits.len()).unwrap_or(u64::MAX)); + let ndv = if other { + rows + } else if let Some(id) = col_id { + let base = if mixed { + // Different input columns (or columns plus + // literals): at most their sum. + let mut seen = Vec::new(); + let mut sum = 0.0; + for proj in projections.iter() { + let s = strip_alias(&proj[j]); + if let ExprNode::Resolved( + ResolvedExpr::Column(c) | ResolvedExpr::OuterColumn(c), + ) = &s.node + && !seen.contains(&c.id) + { + seen.push(c.id); + sum += est.cols.get(&c.id).map(|c| c.ndv).unwrap_or(per); + } + } + sum + distinct_lits + } else if lits.is_empty() { + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + } else { + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits + }; + base.clamp(1.0, rows) + } else if lits.is_empty() { + // Every projection is NULL. + 1.0 + } else { + distinct_lits.clamp(1.0, rows) + }; + let range = if !mixed && !other && lits.is_empty() { + col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) + } else { + None + }; + cols.insert( + out.id, + Col { + ndv: ndv.clamp(1.0, rows), + nulls, + range, + }, + ); + } + Est { rows, cols } } PlanNode::Union { inputs, .. } => { let rows = inputs.iter().map(|i| self.plan(i).rows).sum(); diff --git a/crates/krishiv-optimizer/tests/it/rules.rs b/crates/krishiv-optimizer/tests/it/rules.rs index f797e7a4..27189251 100644 --- a/crates/krishiv-optimizer/tests/it/rules.rs +++ b/crates/krishiv-optimizer/tests/it/rules.rs @@ -3,13 +3,14 @@ //! (when the result is deterministic) the same rows as the unoptimised plan //! on generated data. +use std::collections::HashMap; use std::sync::Arc; use krishiv_common::error::ErrorClass; -use krishiv_optimizer::{TopN, optimize}; +use krishiv_optimizer::{ColumnStats, Statistics, TableStats, TopN, optimize}; use krishiv_plan::{ Builtin, Expr, ExprNode, FunctionId, FunctionKind, Literal, LogicalPlan, Plan, PlanNode, - ResolvedExpr, ResolvedFunction, + ResolvedExpr, ResolvedFunction, Scan, }; use krishiv_types::DataType; @@ -476,9 +477,46 @@ fn an_expansion_has_a_copy_of_its_rows_per_projection() { }; let one_set = rows("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))"); let four_sets = rows("SELECT count(*) FROM t GROUP BY CUBE(a, b)"); - assert!( - (four_sets / one_set - 2.0).abs() < 1e-9, - "{four_sets} vs {one_set}" + assert_eq!(one_set, 2_000_000.0, "duplicate sets double the rows"); + assert_eq!(four_sets, 4_000_000.0, "a 2-cube quadruples the rows"); +} + +#[test] +fn expansion_copies_keep_their_input_ndv() { + #[derive(Debug)] + struct Stats(HashMap); + impl Statistics for Stats { + fn table(&self, scan: &Scan) -> Option { + let name = match &scan.source { + krishiv_plan::ReadSource::Table { name, .. } => name.last()?.clone(), + _ => return None, + }; + self.0.get(&name).cloned() + } + } + let stats = Stats(HashMap::from([( + "t".to_owned(), + TableStats { + rows: 1_000, + columns: HashMap::from([( + "a".to_owned(), + ColumnStats { + distinct: Some(10), + nulls: Some(0), + min: None, + max: None, + }, + )]), + }, + )])); + let ctx = support::ctx(); + let plan = support::resolve("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))", &ctx); + let optimized = optimize(&plan).unwrap(); + // Two projections of 1,000 rows; the grouping copy has 10 distinct + // values, the id one and the position two: 20 groups, not 2,000 rows. + assert_eq!( + krishiv_optimizer::estimated_rows(optimized.root(), &stats), + 20.0 ); } diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs index 8173c606..9e17591c 100644 --- a/crates/krishiv-plan/src/analyzer/analytics.rs +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -45,7 +45,7 @@ use crate::{ /// The most grouping expressions one `GROUP BY` takes: the grouping id is a /// `BIGINT` with one bit each. -const MAX_GROUPING: usize = 63; +const MAX_GROUPING: usize = 64; /// The most expressions of a `CUBE`: 2^16 grouping sets (the SQL front end /// bounds `CUBE (..)` items the same way). @@ -623,7 +623,14 @@ impl Resolver<'_> { ) -> Result, AnalysisError> { match &e.node { ExprNode::Resolved(ResolvedExpr::Function(f)) if f.kind == FunctionKind::Aggregate => { - return Ok(Arc::clone(e)); + // `grouping` calls inside aggregate arguments still resolve, + // but ordinary columns stay on the input (no copy). + let rebuilt = + Arc::clone(e).map_children(|c| self.rewrite(&c, an, false, env, memo))?; + if Arc::ptr_eq(&rebuilt, e) { + return Ok(Arc::clone(e)); + } + return self.retype(rebuilt, env); } ExprNode::Subquery { .. } => return Ok(Arc::clone(e)), ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some() => { @@ -736,26 +743,91 @@ impl Resolver<'_> { )); }; // Bit `n - 1 - j` of the id, as a `TINYINT`. + // With 64 expressions the id can be negative (bit 63 set), + // where signed `div`/`%` on `i64::MIN` reads -1 instead of 1: + // there the quotient is floored before `pmod(_, 2)` (bit 63 + // is the sign). Smaller ids stay non-negative and keep the + // plain `div`/`%` extraction. let shift = n - 1 - j; - // The divisors are literals that are not 0, so the id's - // bit is never null: `grouping` is not nullable (Spark). - let high = if shift == 0 { - id + let bit = if n < 64 { + // The divisors are literals that are not 0, so the id's + // bit is never null: `grouping` is not nullable (Spark). + let high = if shift == 0 { + id + } else { + let high = self.call( + "div", + vec![id, literal(Literal::Long(1_i64 << shift))], + env, + Origin::UNKNOWN, + )?; + not_null(&high) + }; + self.call( + "%", + vec![high, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )? + } else if shift >= 63 { + let cond = self.call( + "<", + vec![Arc::clone(&id), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + self.call( + "when", + vec![cond, literal(Literal::Long(1)), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )? } else { + let pow = 1_i64 << shift; let high = self.call( "div", - vec![id, literal(Literal::Long(1_i64 << shift))], + vec![Arc::clone(&id), literal(Literal::Long(pow))], + env, + Origin::UNKNOWN, + )?; + let high = not_null(&high); + let rem = self.call( + "%", + vec![Arc::clone(&id), literal(Literal::Long(pow))], env, Origin::UNKNOWN, )?; - not_null(&high) + let neg = self.call( + "<", + vec![Arc::clone(&id), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + let zero = self.call( + "==", + vec![rem, literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + let nonzero = self.call("not", vec![zero], env, Origin::UNKNOWN)?; + let adjust = self.call("and", vec![neg, nonzero], env, Origin::UNKNOWN)?; + let one = literal(Literal::Long(1)); + let decremented = + self.call("-", vec![Arc::clone(&high), one], env, Origin::UNKNOWN)?; + let floor = self.call( + "when", + vec![adjust, decremented, Arc::clone(&high)], + env, + Origin::UNKNOWN, + )?; + // `pmod` (not `%`) so a negative quotient reads 0/1. + self.call( + "pmod", + vec![floor, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )? }; - let bit = self.call( - "%", - vec![high, literal(Literal::Long(2))], - env, - Origin::UNKNOWN, - )?; let bit = not_null(&bit); let cast = Arc::new(Expr::new(ExprNode::Cast { expr: bit, diff --git a/crates/krishiv-session/tests/grouping_analytics.rs b/crates/krishiv-session/tests/grouping_analytics.rs index a544dce4..0920d710 100644 --- a/crates/krishiv-session/tests/grouping_analytics.rs +++ b/crates/krishiv-session/tests/grouping_analytics.rs @@ -28,11 +28,8 @@ struct Case { output: String, } -fn load(file: &str) -> R> { - let path = Path::new(env!("CARGO_MANIFEST_DIR")) - .join("../../tests/spark-golden/v4.2.0/results") - .join(file); - let text = std::fs::read_to_string(&path)?; +fn parse_cases(text: &str, file: &str) -> R> { + let total = text.split("-- !query\n").skip(1).count(); let cases: Vec = text .split("-- !query\n") .skip(1) @@ -49,9 +46,36 @@ fn load(file: &str) -> R> { if cases.is_empty() { return Err(format!("no statements in {file}").into()); } + if cases.len() != total { + return Err(format!( + "dropped {} of {total} statements in {file}", + total - cases.len() + ) + .into()); + } Ok(cases) } +fn load(file: &str) -> R> { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../tests/spark-golden/v4.2.0/results") + .join(file); + let text = std::fs::read_to_string(&path)?; + parse_cases(&text, file) +} + +#[test] +fn malformed_blocks_are_not_dropped_silently() { + let good = "-- !query\nSELECT 1\n-- !query schema\nstruct\n-- !query output\n1\n"; + assert_eq!(parse_cases(good, "good").unwrap().len(), 1); + let bad = + format!("{good}-- !query\nSELECT 2\n-- !query schema\nstruct\n-- !output\n2\n"); + assert!( + parse_cases(&bad, "bad").is_err(), + "a block without the schema/output separators must fail, not vanish" + ); +} + fn spark_type(t: &DataType) -> String { match t { DataType::Int8 => "tinyint".into(), diff --git a/crates/krishiv-session/tests/grouping_analytics_api.rs b/crates/krishiv-session/tests/grouping_analytics_api.rs index 9c638143..0259df48 100644 --- a/crates/krishiv-session/tests/grouping_analytics_api.rs +++ b/crates/krishiv-session/tests/grouping_analytics_api.rs @@ -328,14 +328,17 @@ fn invalid_uses_have_spark_error_classes() -> R { error_class(&s, "SELECT grouping(a, b) FROM t GROUP BY ROLLUP(a, b)")?, ErrorClass::WrongNumArgsWithoutSuggestion ); - // 64 grouping expressions do not fit the 64-bit id. - let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); - let query = format!( - "SELECT count(*) FROM t GROUP BY ROLLUP({})", - keys.join(", ") - ); + // 65 grouping expressions do not fit the 64-bit id; 64 still do. + let rollup = |n: usize| { + let keys: Vec = (0..n).map(|i| format!("a + {i}")).collect(); + format!( + "SELECT count(*) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ) + }; + block_on(s.sql(&rollup(64)))?; // Analyzer-only: 65 sets would still run. assert_eq!( - error_class(&s, &query)?, + error_class(&s, &rollup(65))?, ErrorClass::GroupingSizeLimitExceeded ); // A cube has 2^n sets: 17 expressions are too many, 16 still are not. @@ -346,6 +349,7 @@ fn invalid_uses_have_spark_error_classes() -> R { keys.join(", ") ) }; + block_on(s.sql(&cube(16)))?; // Analyzer-only: 65,536 sets would still run. assert_eq!( error_class(&s, &cube(17))?, ErrorClass::GroupingSizeLimitExceeded @@ -359,6 +363,26 @@ fn invalid_uses_have_spark_error_classes() -> R { Ok(()) } +#[test] +fn grouping_reads_the_top_bit_of_a_64_expression_id() -> R { + let s = session()?; + let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); + let query = format!( + "SELECT grouping(a + 0), grouping(a + 63) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ); + let got = rows(&s, &query)?; + assert!( + got.iter().any(|r| r == "1\t1"), + "the grand total (id -1) reads both bits as 1, got {got:?}" + ); + assert!( + !got.iter().any(|r| r.contains('-')), + "no negative bits, got {got:?}" + ); + Ok(()) +} + #[test] fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { let s = session()?; @@ -379,3 +403,13 @@ fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { ); Ok(()) } + +#[test] +fn grouping_calls_inside_aggregate_arguments_resolve() -> R { + let s = session()?; + assert_eq!( + rows(&s, "SELECT sum(grouping(a)) FROM t GROUP BY ROLLUP(a)")?, + sorted(&["0", "0", "3"]) + ); + Ok(()) +} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index f14c2e53..8f1fe348 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -508,7 +508,7 @@ impl Lowerer { GroupItemKind::GroupingSets(items) => { let mut sets = Vec::new(); for item in items { - sets.extend(self.item_sets(item, written)?); + sets.extend(self.deeper(|l| l.item_sets(item, &mut *written))?); } sets } diff --git a/crates/krishiv-sql/tests/lower.rs b/crates/krishiv-sql/tests/lower.rs index 33b469bb..013f8716 100644 --- a/crates/krishiv-sql/tests/lower.rs +++ b/crates/krishiv-sql/tests/lower.rs @@ -191,3 +191,25 @@ fn deep_queries_lower_on_a_small_stack_and_long_joins_hit_the_cap() { assert!(lower(&dotted).is_ok()); }); } + +#[test] +fn nested_grouping_sets_lower_on_a_small_stack() { + on_small_stack(|| { + // 900 nested `GROUPING SETS (...)`: the parser accepts it (below + // `MAX_DEPTH`), and lowering must hop stacks through the recursion. + let sql = format!( + "SELECT a FROM t GROUP BY {}a{}", + "GROUPING SETS (".repeat(900), + ")".repeat(900) + ); + assert!(lower(&sql).is_ok(), "nested GROUPING SETS lowers"); + // A single level nests the same way and keeps its sets. + let plan = lower("SELECT a FROM t GROUP BY GROUPING SETS (GROUPING SETS ((a)))") + .expect("nested sets lower"); + assert!( + plan.root().to_string().contains("GROUPING SETS"), + "keeps grouping sets: {}", + plan.root() + ); + }); +} From c997748ea516a76f82f25eeb5f54fc302ad75853 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:59:15 +0200 Subject: [PATCH 4/9] optimizer: name the IN-list literals in full (typos lint) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-optimizer/src/stats.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index 8e5c5210..cbbd57f2 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -260,7 +260,7 @@ impl<'a> Estimator<'a> { let mut null_rows = 0.0; let mut col_id: Option = None; let mut mixed = false; - let mut lits: Vec<&Literal> = Vec::new(); + let mut literals: Vec<&Literal> = Vec::new(); let mut other = false; for proj in projections.iter() { let s = strip_alias(&proj[j]); @@ -270,8 +270,8 @@ impl<'a> Estimator<'a> { if col_id.is_some() { mixed = true; } - if !lits.contains(&l) { - lits.push(l); + if !literals.contains(&l) { + literals.push(l); } } ExprNode::Resolved( @@ -284,7 +284,7 @@ impl<'a> Estimator<'a> { Some(id) if id == c.id => {} Some(_) => mixed = true, } - if !lits.is_empty() { + if !literals.is_empty() { mixed = true; } } @@ -294,7 +294,7 @@ impl<'a> Estimator<'a> { } } let nulls = (null_rows / rows).clamp(0.0, 1.0); - let distinct_lits = to_f64(u64::try_from(lits.len()).unwrap_or(u64::MAX)); + let distinct_lits = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); let ndv = if other { rows } else if let Some(id) = col_id { @@ -315,19 +315,19 @@ impl<'a> Estimator<'a> { } } sum + distinct_lits - } else if lits.is_empty() { + } else if literals.is_empty() { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) } else { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits }; base.clamp(1.0, rows) - } else if lits.is_empty() { + } else if literals.is_empty() { // Every projection is NULL. 1.0 } else { distinct_lits.clamp(1.0, rows) }; - let range = if !mixed && !other && lits.is_empty() { + let range = if !mixed && !other && literals.is_empty() { col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) } else { None From b3a96acfda20943ad6ee3f181d19d28e4bf9e2f2 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:00:24 +0200 Subject: [PATCH 5/9] optimizer: finish the literals rename (typos lint) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-optimizer/src/stats.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index cbbd57f2..255ab67b 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -294,7 +294,7 @@ impl<'a> Estimator<'a> { } } let nulls = (null_rows / rows).clamp(0.0, 1.0); - let distinct_lits = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); + let distinct_literals = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); let ndv = if other { rows } else if let Some(id) = col_id { @@ -314,18 +314,18 @@ impl<'a> Estimator<'a> { sum += est.cols.get(&c.id).map(|c| c.ndv).unwrap_or(per); } } - sum + distinct_lits + sum + distinct_literals } else if literals.is_empty() { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) } else { - est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_literals }; base.clamp(1.0, rows) } else if literals.is_empty() { // Every projection is NULL. 1.0 } else { - distinct_lits.clamp(1.0, rows) + distinct_literals.clamp(1.0, rows) }; let range = if !mixed && !other && literals.is_empty() { col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) From 80469d0ebe62b123ce31af2774ec7bb9c15987a5 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 23 Oct 2026 13:12:40 +0200 Subject: [PATCH 6/9] feat(plan-codec): carry the Expand node (format v5) PhysNode::Expand (ROLLUP, CUBE, GROUPING SETS) and the NullOf project item it uses made the codec's matches non-exhaustive. The proto gains ExpandNode (layout, projections, nullability, subqueries), ExpandProjection and ProjectNullOf. The writer format version becomes 5. A plan raises its min_reader_version to 5 only when it holds an Expand or a NullOf item, so other plans stay readable by older readers and an older reader never meets a node it does not know. The decoder refuses an expand without a projection and a projection whose item count differs from the nullability list. The v1 to v4 goldens still decode; the v5 goldens add three Expand shapes (a rollup with a grouping id, a single projection, a computed column beside a nulled one with a subquery). The TPC-DS ratchet shrinks by the four grouping-set queries (q05, q14, q18, q22) that now plan and round-trip. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-plan-codec/src/decode.rs | 35 +++++++++ crates/krishiv-plan-codec/src/encode.rs | 38 +++++++-- crates/krishiv-plan-codec/src/lib.rs | 2 +- .../src/proto/krishiv.plan.v1.pb.rs | 42 +++++++++- .../tests/common/inventory.rs | 65 ++++++++++++++++ .../tests/golden/v5/aggregate.bin | Bin 0 -> 1730 bytes .../tests/golden/v5/aggregate_final.bin | Bin 0 -> 848 bytes .../tests/golden/v5/exchange_broadcast.bin | Bin 0 -> 754 bytes .../tests/golden/v5/exchange_hash.bin | Bin 0 -> 758 bytes .../tests/golden/v5/exchange_range.bin | Bin 0 -> 764 bytes .../tests/golden/v5/exchange_round_robin.bin | Bin 0 -> 754 bytes .../tests/golden/v5/exchange_single.bin | Bin 0 -> 754 bytes .../tests/golden/v5/expand_rollup.bin | Bin 0 -> 887 bytes .../tests/golden/v5/expand_single.bin | Bin 0 -> 762 bytes .../tests/golden/v5/expand_subquery.bin | Bin 0 -> 952 bytes .../tests/golden/v5/expressions.bin | Bin 0 -> 1470 bytes .../tests/golden/v5/file_scan.bin | Bin 0 -> 495 bytes .../tests/golden/v5/filter.bin | Bin 0 -> 967 bytes .../tests/golden/v5/hash_join.bin | Bin 0 -> 1030 bytes .../tests/golden/v5/hash_join_anti.bin | Bin 0 -> 879 bytes .../tests/golden/v5/hash_join_partitioned.bin | Bin 0 -> 1030 bytes .../tests/golden/v5/limit.bin | Bin 0 -> 759 bytes .../tests/golden/v5/limit_offset.bin | Bin 0 -> 752 bytes .../tests/golden/v5/nested_loop_join.bin | Bin 0 -> 866 bytes .../tests/golden/v5/one_row.bin | Bin 0 -> 386 bytes .../tests/golden/v5/project.bin | Bin 0 -> 953 bytes .../golden/v5/project_without_nullable.bin | Bin 0 -> 392 bytes .../tests/golden/v5/replicate_rows.bin | Bin 0 -> 879 bytes .../tests/golden/v5/scan.bin | Bin 0 -> 742 bytes .../tests/golden/v5/sort.bin | Bin 0 -> 758 bytes .../tests/golden/v5/top_n.bin | Bin 0 -> 760 bytes .../tests/golden/v5/union.bin | Bin 0 -> 870 bytes .../tests/golden/v5/window_aggregates.bin | Bin 0 -> 1790 bytes .../tests/golden/v5/window_ranking.bin | Bin 0 -> 860 bytes .../tests/golden/v5/write_files.bin | Bin 0 -> 748 bytes crates/krishiv-plan-codec/tests/ratchet.rs | 6 +- crates/krishiv-plan-codec/tests/roundtrip.rs | 72 +++++++++++++++++- proto/krishiv/plan/v1/plan.proto | 28 ++++++- 38 files changed, 270 insertions(+), 18 deletions(-) create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/aggregate_final.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_range.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_rollup.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_single.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expressions.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/filter.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join_anti.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/limit.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/one_row.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/project.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/replicate_rows.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/sort.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/top_n.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/union.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/write_files.bin diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 552f70c8..80b6d2d0 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 b1fa442f..9c90308f 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -243,10 +243,29 @@ impl Encoder { input: self.node_index(input)?, }), PhysNode::Exchange(exchange) => Kind::Exchange(self.exchange(exchange)?), - PhysNode::Expand { .. } => { - return Err(CodecError::Unsupported( - "an Expand node (not in this format yet)".to_owned(), - )); + PhysNode::Expand { + input, + layout: l, + projections, + nullable, + subqueries, + } => { + // An older reader has no `Expand` node to read. + self.min_reader = self.min_reader.max(5); + Kind::Expand(pb::ExpandNode { + input: self.node_index(input)?, + layout: layout(l), + projections: projections + .iter() + .map(|items| { + Ok(pb::ExpandProjection { + items: items.iter().map(|i| self.item(i)).collect::>()?, + }) + }) + .collect::>()?, + nullable: nullable.clone(), + subqueries: self.subqueries(subqueries)?, + }) } }; Ok(pb::Node { kind: Some(kind) }) @@ -281,10 +300,13 @@ impl Encoder { name: name.clone(), expr: self.expr(expr)?, }), - ProjectItemPlan::NullOf { .. } => { - return Err(CodecError::Unsupported( - "a NullOf project item (not in this format yet)".to_owned(), - )); + ProjectItemPlan::NullOf { index, name } => { + // An older reader has no such item to read. + self.min_reader = self.min_reader.max(5); + Kind::NullOf(pb::ProjectNullOf { + index: *index as u64, + name: name.clone(), + }) } }), }) diff --git a/crates/krishiv-plan-codec/src/lib.rs b/crates/krishiv-plan-codec/src/lib.rs index a451f1ff..0bf83149 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 286966df..7d11c857 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 123fa625..2f673d7a 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -307,6 +307,34 @@ fn join(join_type: JoinType) -> HashJoinPlan { } } +/// `ROLLUP (a, b)` over a two-column input: the key columns are kept or +/// nulled and the last column is the grouping id. +fn expand_rollup() -> PhysNode { + let kept = |index, name: &str| ProjectItemPlan::Column { + index, + name: Some(name.to_owned()), + }; + let nulled = |index, name: &str| ProjectItemPlan::NullOf { + index, + name: name.to_owned(), + }; + let id = |n: i64| ProjectItemPlan::Expr { + name: "grouping_id".to_owned(), + expr: lit(Literal::Long(n)), + }; + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![ + vec![kept(0, "a"), kept(1, "b"), id(0)], + vec![kept(0, "a"), nulled(1, "b"), id(1)], + vec![nulled(0, "a"), nulled(1, "b"), id(3)], + ], + nullable: vec![true, true, false], + subqueries: Vec::new(), + } +} + /// One node of every variant, each as the root of its own plan. pub(crate) fn nodes() -> Vec<(&'static str, PhysNode)> { let keys = vec![SortKey::asc(0), SortKey::desc(1)]; @@ -383,6 +411,43 @@ pub(crate) fn nodes() -> Vec<(&'static str, PhysNode)> { subqueries: vec![subquery()], }, ), + ("expand_rollup", expand_rollup()), + ( + "expand_single", + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![vec![ProjectItemPlan::Column { + index: 0, + name: None, + }]], + nullable: vec![false], + subqueries: Vec::new(), + }, + ), + ( + "expand_subquery", + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![ + vec![ProjectItemPlan::Expr { + name: "shifted".to_owned(), + expr: call( + Builtin::Add, + vec![col(1, "a", DataType::Integer), lit(Literal::Integer(1))], + DataType::Integer, + ), + }], + vec![ProjectItemPlan::NullOf { + index: 0, + name: "shifted".to_owned(), + }], + ], + nullable: vec![true], + subqueries: vec![subquery()], + }, + ), ( "sort", PhysNode::Sort { diff --git a/crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin b/crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin new file mode 100644 index 0000000000000000000000000000000000000000..35a39060157d4948f6e7d071a38ec8e4588b88ea GIT binary patch literal 1730 zcmcgt&2G~`5Z+n;ZMG?OHwBEv!OE$M#3d)x15`bsiVLB~BD$)b#Hwpo_C^giE=c9Z zfd}9PcmxheJOUE0&;#!fX5zL=)bJNXXEiJD%#J_L&bMRX_&hx0@(mZBdM=;I)#^l- zM=uYx(OEH&UYzvxC^Hg_@}+%T#EBBptxTC78r7DLDUOm{-BWj^eNq;&T07NcGAu?_ z@?h!;R|F`~As0>rVMA~Q6)qalUdNK4J5oX=oixvrvMVG9?%M%urNUC=4iB1)F@b|I z`(kDcaJ7Nz93+G_k8~jH8tGr^6#TqS zl#BC6xm;)B%*ly57yA0yeCBh>h3nreC%h#&@wa!37+4|$xn)QE1_fC+YLJK zscdq*Y47+VI#TQSf_jF!jY{kN$z(E9J{h5{mf7;vbR_#F`QOjS{r1Ad7+W_ncCCrk zmyRtYXLS{~eDl-zj&4KRPm{h;kkya6uVCDj(&?pxp$P%HEot|P!YKR^7@cO=yBU9I z?I;}$)7(U-x{A_4UX)3^8wl_xpRM5a+sXG&A3wNlPv<&&Rhi!E++Uy@#08aLZh)116NHrAB& z61^;6n009ptoG6USZBFbIBsoI99uodk*!biNMHOmCf;CSTR-Hhdmw)?oadJ!%Tv+$jQ2;!sr zFmnwN0DlRK+n<4w*%GMc0K1I;Ro+Fwjr=CPB;6-v|83WG+{!gLMsy!!Ol^(nu*e}q zmdC?n-)`4=jKRVqQH0DxR~)X$5a`I9rzdi3xCmLX8-^I literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin new file mode 100644 index 0000000000000000000000000000000000000000..44186479a51888279684b067d8549bebdeb9267d GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M?28vloQ1D6FfHr8HNe6cD9Chz3K61gqG0XT{i&y{AA&gOqeB zcmQ62N1#CBVN&o8X2uRk5p?Y8`|SDsJnL0nDB}x>pQH-(OMHQ|?##5OZ%&LgRkOw* zFIHw-Sp=&?bl)|3p%wa_ZOap@7wFmMw5auceGlEUw#oJGxoL|{v+at5Ar$W@&|qSz zyaeKjk{UWPCdsSRK`<#Ibm*0JU9?Mu5@hH`a9Aowb2xrG6+$T5Lbx|91dz2)2#G^+ zN96;98cACus8U{^^;h&zM91{773(R@?YXC z^@DIhI3lq9y6<}y<9zCm$?#Dra&Ee$T-QxoO*I<*{xfX#c70- uva1_gZ>uVT1&BkbPYel}0{5FS&6)c%C-rgL;mowM74-|k;p=wGY= literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin new file mode 100644 index 0000000000000000000000000000000000000000..87676f17bdb8c6e75c1bcc692a3a0b666b25651d GIT binary patch literal 758 zcmbtSy>8S%5Z+n;?af_?vloQ1D6FfHr8HNe6cD9Chz3K61gqG0XT{i&y{AA&gOqeB zcmQ62N1#CBVN&o8X2uRk5p?Y6`|SDO*w@1c9vHo4wCH*K+Lwq0>BxZQL5dv3__)B}^@%j8pYpM;kjn2Cqr4Sc?x7jO+;ym56iG-q*L&dCYm z;Nq{r$KE31gW%;3GUj!p|L=hrs&>5_6?H*JwWixh8S%5Z+na=e@~=bM}HT7KL>cvXtg3l!8R55Te0QBEc&5-B~eqWbY}^(I6#V ziabDGAdf(S#KWZE9n6d!kRs^V)%V#mv-7*Fyg)|h62D62>o@osx7$-TRlh5)+c$TAL~cxo;G!+r)Q?g*Y&2&_j*&br9gv; zr1D~j#!6~vNf^g(69>U0h|r=}R8`(A6iSeR8^V5}9L?tF!$b(7Xba)qu@FGgE+HTe z!~-!U8=ne@LqbfTQ}PM15)-m+^}aEy2lqaPCE;kIZvCBJFK!Pa-2`eHO{gpx0=fzU zh#|R@|0;u%{veza4hby3>$;B9@_EXHasOo+Q3)g3ykyTz+zOX*b>GkD7`(WAb=Nw3 zaa{Jv2t)hgul9%BBII%Ka{E2_wIu)V*bK#3jM06SFRg}h_%%xnpB9K-TCCPK0$9w@ zP3zifdP3F|CEdD^&zvJ|;&oBk_{_9%v8w7O&t8T!=kB``e!1v=e!ICI&3#ju@^fp` z?R)Phjgk<@MO)Rj+LUDoa}fJd$wPPnsjKDF-E@}hAUr1vW&2}#n9R^IRmN%ueye{0 D;~ro$ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin new file mode 100644 index 0000000000000000000000000000000000000000..c4f0eabdd9a59858b29e2e6c03d64928e4e5353c GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M*I;vloQ1D6FfHC7P>H3KFG4hz3K61gqG0XT{i&y{AA&gA}?H zJOD4iBTz)j!=&IH%#0n7BIwxF_u2FNdDg4EP{tP$KT8$pSNIaQyED_CzCAJ4RLu&5 zyjYq|Wf80n(S6_Kg;wZywk=Pro}*`*)1uZ7^nG;C+9ub#=cX;z&890J521KVfd&&x z!MvKlpsSlf`d{yn#1wCsSrZZ7Q(${A%LuXLP#8n z9kEL>&I*YmLPDS{`H)zNDMh!2+JxO>`Ow3XaXeKw!QS8(H;1umeYF$IxkwZWE zscXUo;fTQYo4)T^jPt2KCc~4-Ct;s9FF7z155VGI@74v}gBO3ax*eFSxG&e_gmGx` z*U+Q3i1-kAxx*goIt)x;WO9 z75!@=pE$?bB{PF#pKMzc8s*jz`w(s1UEKVbw vlwIA}dQ(*q%t0JTC6C}4WUiKv_On^Khwzjlr1mFtFP))d>YUvU{C58WYo1@V literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin new file mode 100644 index 0000000000000000000000000000000000000000..b9c7afa741c7cac571b6fd349b8e0010bb9f5a35 GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M*Jk*$cv06xLP9Qktt!iV&qjhz3K61gqG0XT{i&y{AA&6Df2Q zJV0I`k3fOM!=&IH%#0n7BIwxF_u2FNdDg4EP{x-MFQf|eD}0ID-I-}m-<}w2s%C{j zUM$U~vItg(=)Q0ALM!w;+mLnz)-puxma zc?rZ5B{g(pOp@2BgJ4ob=+G)CCJc?;Gk5F=5YLODuht9g>Y|K2q0^Z5E6&t zp4g!nXNANOAt6wfd_=6ol%iWhZNl!ceCT1xIG(DTV0Z9~o2RjAeKn1x773(ZIdPHa71AHP2cw{#`)AAli|taldw;lmmHXh2Vn89ck2S~!HYjy-44uE+?Q)|!Z@_} zYv@s1M0^Om++mM(9oheTScZBmCg{E>mR3VGzM7?m4+}&uD_3h911x6fW=&%?T^wu5 zivG2bPn=_Il66_zSJfK?K`(7i_-`v uWmh-0-c(fta}Wno$s>3MnXBdFy=<25A{k~D>H(^r5aL4Mh*Xx^Nvs+>vNviG7ledV zZX9?3UVum7fW*V}z&nH)yDF-nCw4WTJwM;9zws(B^x`v#AEgTPb9{!|&C1lrheyVm zqFP{(W^=PHEP~Y`x^Jp9(+d5@*7=duWAtowoR#{vzJ>0pu2MZXF?F`A)=hT5O~qRZ zG?-W_FM+tHq=tr^Uh-n%AeadvH0b4JnblK;5@hH`u$L=Gb2@(27eXl7Lbw+w1dz8w z77~Wyh8R$et3twvEFq&Rd52JmK4rIBZo=V?-uAHM9Q4(BFl^)E{89Yw1DtC5iuf*i z=+K8_fF3J&jYB=atwmj}S9!TO%u^cNpQ8A>BA<)muG6(e!Po&12~0lAf9Yj)KgrI> z_Q}|Px!G)3AJ>y4Oxh=dPg;H2ykv_EJTun&SD&s;=asD$+|gxq06TSbm`StzVSAQIs(iCqSJKMf;QD^Za!d~8# zl`Ypr5y2S5o>X!l9>dJl^4`vDG#Mg1qzs*QnYuX{p<~LF!wy`AiLon~w<}xpe*8S%5Z+na=e@}tarS~R7KL>cvP5$gN-Sv^P3*2j6JAL<9_o-}o)w@*!zuj@^lKj}@;mI4hX zlFExA8Y`)xC1D)DNgM=|AVP~?QB`@fP$)qLZU}pYax|Nx_Y)z6qAi4b%R&H2yM%x^ z5ckA3+1M)}4hb=VyyPQdB_?Fu>SJS858-_dOTytq-S|7bUfk?Q!8uGdbrHEJQOrTa zkbIHv-;fLBHM9c*DN)BTp)UBv0B>*U@=2C zt!u040a;U&bmu}obq=+O*F|OHQ`5%9s;Zkjdl4!Kx@#vqyX<~meg8I^`=&DGr`D#^ zJFh2=k`Tv5Th+GOlw}BXJ$q8g19%RptL5X}be8NOJR=KL`(wJF%+N7a#%c$CtA7C; CJzw$w literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin b/crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin new file mode 100644 index 0000000000000000000000000000000000000000..1f9e9d7f857fbc0e6b84da5b9dd8dc000ed138bf GIT binary patch literal 952 zcmb_a&5G1O5bo+Elb)I#C)>MXY&?WHNyN*XWI-MFB;v)LR^UrvJN_I!s<1=i{VJv-B2gtW~Ub?w~f@o^ip3z1m+{W zg7=`QCNsjJ9$?rmi@D3QC^SPHn7;m{7z3nF)z{ZI3H4tqo~{eyA3m!%k%^bSlYeUr zx8Xa{G0~KWb~mOCSbg*2mv2`;zI^`F8;7>C<%`zMhUYeV zX3`k@MO)RbIxNc=#vpd2l2f5*)(q{9cB?u@zp^6IwQYC~$t*n!^>e|uT zt5PmqAr-YZ4jg)e3kMEdfD;l@5&bt17Y;peh=jzk=8boDm9)Jv*1UY?y?Jk(@AE9P zfd=;}zekxXKI9+pR&gASZ+?0Zl~I}>ad#LWMw7JUq!hrdTlp{+j5|d+P7X@Z<#w6h zj5Bdwta0mjoDW6&b~KJh`J{+nuR}qD5kaDWGCL%J$EY9$F1+xgp2Z> z4;ZI}g5{I_gqcExgBwl5Fc?0DVQFg^1Xn9K0WBCyMjK;l7ohcVhB$E--9XEXCPo`| zZm8U|^twigi>poc#O>66?wCA`5FvkIO=C}=FjyteC$6Z*N;~uQe@e-Vyoe=lI z*R%!aH4Pxpv<;!A9k{4z7rL7E;FXZ{zM?OZUy1z~=Y5=4aH{nuRaGgu)qYf%hxK1& z9F*R{M-5tAWmF55{9m8Vb5#vBRNI(txyn*?Rn8FStnJ2gZJ(WO`06uI!>XTIUV-cX z`?*9|W8iV?{rIpHB%NV&D9C36&TW|-jmm(K!7{gGo|gizP8lUBe(eUmY+e;nI7+fI zyd4!`a+KxcczDfc1ghUmfBK;M`_~`7b-HepMd_WQl&v%OMg~2fdr6VyWj0AupLB__ zPZ_;J-XqcybaO>6_d1-sjS)O;5nAmnb2G|@%59Q`8#A~en}Tz^%B%se`CL^2B5h&% z6_Jjx!l#uf@eXWagQ~V{iAMjaLN$nVh5Z()G1c(su51YJUG(5MxhpS-jcWl1)v!l@ zO#Z$1=AG}Z|Md80F+Z#=TCy?v{rfb{!;^0w)@436jmlLg-}vXR>#E#nIz6dXQSWw4 zBp2*+hphwgsRXg~1!=J9!yo_r#YT_ytYCb=*qWfs@iB41^c0)VraZ6l`f0_^XuST? zB6M2g3mXgft0=CE|F8+R@Ab~9i(X7_LC2?EvbIV3FXK)}w!1Cn2|J&Z@i?HYDcAN_ j@x42Y3pYxWsECK?ZiPgy^aOVf@(h2NeG_w+RPX--6Rp<0 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin b/crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..50724761fef2212a8852bd0822410443c91392fa GIT binary patch literal 495 zcmZWmy-ve05I(zU!k}@Ff^nG_H`@2uFMYYNOMBWIYj@6NLCuVv(i>%s~Tj~B`5M0i#t|J?m zpw<38+EnY(nM{*VHPU!6j>5f(u|vNhm}17~qN zt!KUj00I-hSuk|~pNL1qDT3yUJkNKq(tQ9Xby10GsHRj4EA#jse zc#?|tk|snnVFLlMS~!cM_i!$~!J{(@r^r&5;#}%h{^>Mb2FA&od>v?D1<(i#vMtD9 z?bM}n9`&taqp~%9#qK*&R(tVaXnn>y`0BkRvC0i2YvK{UF7Sg&R~bEw$5!8&RO5Ff c(?6FU-2Mz^w|weO_m|lHj+tDGsfgf@kZayi!-y6>yK6I#W9*rpZ_Jy3KF&klvD9 zaifLwVy;J0DBhwmiXSBoaFYP=7QCXW@@6JMa1pu@-!G)22-Z)>j4_FuG51_E#?jhC z2~mga65BzHqC(UWN{oV0;u>nn#)xk9xiPti^q#|_aWIx=!EWzoXLqzvO#6u|dC0ZZ zjlu&(qv9Tdm77;#1D{g%opsm#3~0 r*Z0y%vJ3n+B3Nw$T}~$8m?|T;!#D1~l6#xbU&*~-1N}@c{U7e{Zj^Eb literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin b/crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..ddfde900465a97dd93e490ca1dcf1663622258d3 GIT binary patch literal 1030 zcmbtSzi-n(6ux)1Fix424Oz#t!VovQ7>< zUe0$I6PLtAZX*x(N3jyCe5E+oZZ_|YHn(Hdl+EDGkJhqSI=c2;k!)pgw7x6V;BQwh zFWjAHq;ui*D{XH(J@9Gpoc@U|&CMtDee?}!SLoc#GXe%Sjv2R-NnX@HV}!q0kw(`K8^h?|HWpr(M*|~sb=z_BN5{8s6IV|_i&?L2UWNYuPD)~punFdq zWU2(swnsS?^lb=WMd@s=JW4~HMpP6^;ICCO&G4fUcF!0K>Ce+#`RB6oaZE*-#E)H0 zoy|AHU4PhIef{#eF|cJWvv-w>IydGQd4UVfw8{&Wud>Xg14>4Wvr`;XWD0htAN7JR W(1(cNYrmoWpa+J`V;wf=b@&$mJbirt literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin b/crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin new file mode 100644 index 0000000000000000000000000000000000000000..048d0341b790566fb4f1ddffe0fdfd5ede05a9fa GIT binary patch literal 1030 zcmbtSzi-n(6ux)1Fix424Oz#t!VovQ7>< zUe0$I6PLtAZX*x(N3jyCe5E+oZZ_|YHn(Hdl+EDGkJhqSI=c2;k!)pgw7x6V;BQwh zFWjAHq;ui*D{XH(J?hioIsFq`nwwAP`{*0!I=?-cOf)A#7qRnumcYdTk8{=e!<#AsEy zc|Gnc0qOKp_60m>ac+gl)Nv8ZgS@9r|`2YFw^k#js%m*IXN7nhHdSvHAw*9IwMpXI+! z!2Ny_E(k{i_Fs{VJMw;Nf@%L`8d3KVEjHOR6Hmb>_;OPhat$^>Ag`nk1E)8FGv zRnrF-@}YOEO}eQXm!6w0t=5fg%i?LQAlm)#;t!v8zrKI_I$DILG4=b-<+B^VPM%~j wj;pS*uG!Xg3=0s4QpqED0=ci{{eyg-?IS#*40ZdHx}D9@Gflx^53a+%0R$>&D*ylh literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin b/crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin new file mode 100644 index 0000000000000000000000000000000000000000..d55e972c842f15a43345d50a9d77725d3eeb3e44 GIT binary patch literal 752 zcmbtSyKdA#6rDSM&E9N?GaH1lD6FlJr8HZi6d_855DkVB2}ZGZ*NU+tdscytCQ|4q z_<;OCK7j&>k4eEln0xJj6hX&~&avp?FJy1`|u= zB@j=P)Xl4hYF zgiFE^f$cYa-?J3wQ+G^;FOyHQK208SU?yHZOMksv7jO?AykvE|U9RH3T$2;VEsMXl zJZg)GkAR0e%(1Q``+twhP>;m~-515uYN*Cvv()fmf#_xBYHeeH#SGo7X{@GaV@+An zO$+(NIo2jwm$gmKO_!9bx@n92d88ohzdO;V&;8GD*H`{LFtw>Zb~f9-^JcO*jc`(S tbz|#IRYfodaUhjEf@hGqT0Y*(X6Y`%Aw@{-Pw0L+L&wxPyB+xL{smTIUrGP~ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin b/crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..0f49b3f684db38f274f53e354754e1257d3ef418 GIT binary patch literal 866 zcmbtSy>8S%5Z+na=k?^m`SyY^7KL>cvXtg3lmenu2+?4u5E86n-<=g>NA{jVx->|k zqu>E}0Um(@iHAwSI~+5<=tKaCf?a)|J%8Wqj;*{<#+MS`OXch5_zZWuQ`4TjJ~q}= z%>wNqj3~xgA#p@V2$Uu75i2pI=uTgouzOS9YgjUlhU&^6_kMBpD3)5pU#-_`?#%hr43qxJG@w%qX!DXiGx2TsB%f~9`K-Z< zZ=kMk&sD6;H92830rAHKHoZl}6XE6ddsEku{lAklq%jN7eOk<|hHCR_<{IA35WTEi zENu)ho1mLDjn(uQSyNW@;e~wQ9BGp*%i1Pqrc262-LysiBvKHrzd6yxhxL!IUp^0} zzNt-h-q~#T#xIn`X@o)9)s3xJRTaS$#DP@u5FSJ3YI%Pzo1|lehZLc1+tBTFf{v+k cb~~_j|5M!Crv5JO`Pii!N#-+}ZkuRfS z$fPWS7>I9weL3G?Xu70#ne5`b>}iMvPKdx!O4eyEX`i!E@*6{?ntc=Pv+wYq&+%sn z$=RH%uumOuL;}Yq@zm=*qi))Du%e$wzqaC`HfFi%Hrr4MQ6C$%8pmMsrVYO9r%FOz qEUv6~+inQ%-cPRE4C7tfT$frD@oP~&y~Ll-kN4~>_k-^rr|=8UmMiH1 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/project.bin b/crates/krishiv-plan-codec/tests/golden/v5/project.bin new file mode 100644 index 0000000000000000000000000000000000000000..cb7da8071df616756d098deac40e248a474db82e GIT binary patch literal 953 zcmb_ay>8S%5Z+na=e<0FQtK5)YGtcQ|J3=%fH0>}tk4Ki|xFW>gSKb0+bvRH1&1k8q|d1mN|Ub~n_;V3{41f(z2w>NhP_Fqfh?JNXtN;K2 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin b/crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin new file mode 100644 index 0000000000000000000000000000000000000000..91115ce8affabf703ee8dab06fc19e701b060b4c GIT binary patch literal 392 zcmZWiF;2uV5FGn_i7^Nw-xX0rVO>d?iV{&Gs;jTDl-ymUD3ORIpreG2f(P&coAdHKh zB_Urz#gI{11Thfb{`ykB!O(O`A2Qj+ciB@BGc1X~SxnY(&S{^rk@Fiv#hQEr?UV2H zpU?4U2+7%;t8h#ma6|&fM)B0kBcpEXZLp%9M!&Vy=wD2oh#f@ zpuvPvc@cyoB{j69jH2hUgJ5DrXwfUGDsK{n5~S}2uwN)gvpak>7D6c6Lb#VG1dz5z z;uHGfn%E&9pYjOPtBM9ZSG(TQoJVo}vi zo;?oeao0bb;N8dduWw(!45mXGqXjO~W|AQPEbltyX0jz!b#3RB{H7Aa%98 ey_ZhnU4#eZp==x1^>~7gsWLV@urU+kuKOE#TV@6T literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/scan.bin b/crates/krishiv-plan-codec/tests/golden/v5/scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..8d48e76e5683a6f6a4c8cde8f70b6fc8c4904fb2 GIT binary patch literal 742 zcmbtSy>8S%5Z+na=e@~=ID0`Di^94JSxR#iN&!(SglI67NU(~1cUH_gviB6|Xpllj z!2|FDJOTw050ip-Ff(>QilAdg-)GPNj=d^~WO6C-gH)k@j?ZwnJ2&mwt5f4l-K;Sz z%9YvH4#DXN{WncfYK4=|wbiN93k+OyR$6^m-$DPpZ3?}2VcK%jY`gMs2qilTG?+xH zAcbVAq=t@+Y5F4b5KM*$9R`)PWxG@;K}LQIN0st4hm+ScA%r3o!oOxAfUE;TL>!4* zVvk~+6%ogTlt5YX9mWTvjd{lPD;9%Y~8zXo#qAB0Q7F@f!u zec$r{&ZnN34o{{D4V#d%$bp&o&^++-&AO0#u<$|E^^#o0eYvI}jEjoD79D4cm=|a9 z2OaA=vj6u38S1f|qW`p9ISuvrYE~NFE)j#gT5nteu$-fxH;vPDS)8eA`oBUx@J_Tz zH8S%5Z+na=e@~=ID0`Di^94JSxR#ix&oq95u%Ahi3F?IcW1@ek-eutM}w4f zDR=;0fJdM}%EP4K9n6d!kRs^V(f8Riv-7vR$_r$4De<#ZzCOWY+-}cJbN2SsSX0(3 z^s{_vHl;V(^mXyPZy7hN^ySRN3O%|wS2-;N;Kn%&X z{8s~9?uKwlI3Te8rt3N`$or`V#{HLRM8!uWykyTzJOnrJ<$j#cIe2m7>TYO`;=CM_ z5r)CVUxN>|MaT!i%k5=|Yf1m#12dFkK1TOdzO)+3;nyrRe3&D8X|Y<{2w*-%H?3={ z>G4=olyu)hK64JWiPuGC;|tTq#j2{CJbM|^Id?yt@Y7ZI^ZU22qnU3iQ+{l1x_$5E zq)`&$xM-`|R-3X6VFqF$l{|nWNL?+T?xoXY7vVX1sM;Ub!(@t%sWLV@@Z08S%5Z+na=e@}#arS~R7KL>cvP5%Llp;h|Aw+|rM1oc9yR%~K$lg;(mnKru zQSbn~0FOX{#KWZE9n6d!kRs^V(f8Riv-7vR$_r$4E%8b!U%$pzxZR$a=Jef(v8Jq7 z=x6!TY)Xq@b%5@NI?J`fQEQvx#OfJ(wm!`({YXDV_pGTiy?bt&d|hwa{AurswiIYE zkyKs`(O5|hEh*#pP2wP!1QA;FimJ++xk3pta6>pOl%v@ly`Kmn6p0Y-9SZ@Z?Gpmx zKKI3&ab+9e+oD={JORv#O)c?j=wSW*@fb?fi-c5(YG+F4Ar3_-gJ0thjN z>rA8Lz`kAj!m%MjO+{=Y|ND93z^?#p~>HI&1zS!(z=NA%KSwYCw!e2Q*b z*H+W>v8E{L#)Uj^7TUz?qO$S1Y2#v5)lHtg2{X4$MqnYqGPIz%?|uF{{mwz BU%UVS literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/union.bin b/crates/krishiv-plan-codec/tests/golden/v5/union.bin new file mode 100644 index 0000000000000000000000000000000000000000..32f0afadd3ea26360418e4dff225fd2f263f5917 GIT binary patch literal 870 zcmbtS!EVz)5Zztd@n(`xXVa>##lgx6St^mj1qDH?oDkxI;fMs=*lDb|cJOXgxaERW zZXEaket=KlfW#60qlf-Om~~R53OF&VdFz?knf04l?szP?VDO%CS3HMj&`}Fn&0kHV zmU(#!Zjw&rGS@)0@W4JRlT>hMC|zX}E&AZ-a-J6AK#m>@w3XPk?N z;sLnWi#N+uiv!qK_eO)k$#ODJ^*|1yyVs#T(c4mQsqC~!lf#A7vwzX9Nsl{xwX(ZE ze;Bzs3fvP1Ed1ZliaxTx5jOO`qh6pMqZ<3=dc8IT&3WVmVZ9q@AlU{kBdkUPrfG)o zr<-}MsbR!4if@~2vP@l*lb~)IP5x-~=5FGf#v8HgwaHUx|L=e#))+Ryewt3Tp!w!g zP6d5C1aRW)bfyDJhd6+^EVaPLs%4hr`y%YVH5M|QWrYsUr3!H}WtApRd``XfH_Kmr zSYLhp^10D>Wg+u-O2?fWBa4H`hi0aVQWwiS_i3M!5##KbKBloP*qz RJG2+|z>-B`+!nod{{p@Ne6#=n literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin b/crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin new file mode 100644 index 0000000000000000000000000000000000000000..c9a55711561e6751d820d9fdadc9e50e00a4abb3 GIT binary patch literal 1790 zcmcgt&2G~`5Z+lk$!42UcU!<%94wzuBrZ82AE4@~LR>IC7PYI|Nvs;XvNuXW98jTB zxpCkDcmW=P0}>C@1Md)Ky=j!v@{=O5t66z>Jf3g9of)%+V|%BRpHS8k&-hcmTph~t zV0T|CnH2-xij%$^Wr~9m4!2(vaUvLRR;o<*l?Zr46$eQ!?ugsmJ}iq^tc+!u42w~f z+^<#M5(5F!qpabB*Je~eg`2j&-L*K#F6U73Mw;hI*<+jnb!->HlvzTny}b@0gyAH_ zp00!d?$%KpTpe2a~_$c@my^AX&9n>z>w!S_#=JhueZZcqi=DF=G%^;d2H)bhodAM%$oM(CGBA>uF$U1U<+%o7$agMSMIgvSha` z##QhZ!i-T52s&Pv1|9m;aZ5kzp{Zc#T$4#rv|dwbKpY_mT7y`cf{^(j30{XoW*HDW zzmNXbm3LvBgQf30qclfZiInqoPm(e-yW=*KQL zz{yw3ee-_u{mbW1&7dW7nZ2x3w0z+Y5_w&hx6>*wR6fcw7Xlz5Wpo3!AhHF$w;ru^ YS2;XH4SJiYZg$tWCG%JxTX25-11?3X4gdfE literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin b/crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin new file mode 100644 index 0000000000000000000000000000000000000000..f8a20245a49e3d9fb349b5b10d66cb20427aadae GIT binary patch literal 860 zcmbtSv5wO~5ZzgClFeK$@of%aRuq;iWC;nSLRY{k6+$!^N+ejtPGZH_k-gCY6%EqS zQSbr$0G~jC#K&C0KZIH5Aa{a}UA^a>H#@VM@yL!Z^f|{5T=vv6e2Rz7+SIF;Ghl0>*HxM+iQUH5`OK;aJGNS7rMjhVqFC2eszzs~&X(1t$?mo;eIP-B z(Oh;y(0$1jG-UL{=dp)iVnk@Llb2;yPb6~Sz6jtXm!5KN{c6A%le8HVm&_QDbxiPy zeKuqx^0`|+aX<(Obc^36mTW-YgElsF@-Dp1Vab>d1;0T9Xi65!NIOALOW`fvB^2O3>fXNs|QdL&bh+2~u zG*yh>_om8(%bdD;W|}ZxlvSOjj{^z*_PZCn|G546?dz9r)H9_i-ZnNlyf&YSjsxuH uOwMJRlFf_Q!Q29;0W<)X5(FmVW{2=xufY literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/write_files.bin b/crates/krishiv-plan-codec/tests/golden/v5/write_files.bin new file mode 100644 index 0000000000000000000000000000000000000000..0e6927a4e89736bd79b3c5c33d9214f733087356 GIT binary patch literal 748 zcmbtSy>8S%5Z+na=e@~=ID0`Di^94JS)#cLrGO|ELNpjkBv{41I}_uL>^%iK8l=!s z@Bq93k3fOM!=&IH%#0n7BIwxF_u2FN?bxe=NG2B&KS&koXZRF%yED_CzC1C`)OL+w zQLfCkb_h;K=)bW=sTEE-*H$M^FEDWSv~2WkeGC1w))sp2+_dG!ZoBf{5K49wXfTOX zK?=!KNevwt)AV`fA(#viIt;3&Dchw&2{Q6yIINVXIh?$j2_Y11A^a;A0?0ZbM8uJ} zA@(T7SrKteNC}iB?+`07qv+01o3eW>A9^@4j%VsJ+#men@?j!hixjd?@?Wwn@`G?e zI3lq9qVIc_;e6_h>F{KlkmQ6mFF7z1ADv}CU#| Vec { diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index 62958449..b9f9fd44 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -282,11 +282,10 @@ fn an_offset_limit_needs_a_v3_reader() { fn a_window_needs_a_v4_reader() { // A `Window` node sets `min_reader_version` to 4: an older reader has // no such node. Plans without one stay readable by older readers. - assert_eq!(FORMAT_VERSION, 4); assert_eq!(message("window_ranking").min_reader_version, 4); assert_eq!(message("window_aggregates").min_reader_version, 4); for (name, plan) in cases() { - if !name.starts_with("window") { + if !name.starts_with("window") && !name.starts_with("expand") { let m = to_message(&plan).unwrap(); assert!(m.min_reader_version <= 3, "{name} needs a newer reader"); } @@ -337,6 +336,75 @@ fn a_corrupt_window_is_refused() { ); } +#[test] +fn an_expand_needs_a_v5_reader() { + // An `Expand` node (and the `NullOf` item inside it) sets + // `min_reader_version` to 5: an older reader has neither. Plans + // without one keep the older minimum, and a v4 reader still reads them. + assert_eq!(FORMAT_VERSION, 5); + for name in ["expand_rollup", "expand_single", "expand_subquery"] { + assert_eq!(message(name).min_reader_version, 5, "{name}"); + } + for (name, plan) in cases() { + if !name.starts_with("expand") { + let m = to_message(&plan).unwrap(); + assert!(m.min_reader_version <= 4, "{name} needs a newer reader"); + } + } +} + +#[test] +fn an_expand_plan_survives_the_wire_with_every_field() { + for name in ["expand_rollup", "expand_single", "expand_subquery"] { + let (_, plan) = cases().into_iter().find(|(n, _)| n == name).unwrap(); + assert_eq!(round_trip(&plan), plan, "{name}"); + } +} + +#[test] +fn a_loose_null_of_item_still_needs_a_v5_reader() { + let plan = plan_of(PhysNode::Project { + input: Box::new(PhysNode::OneRow), + layout: Vec::new(), + items: vec![krishiv_exec::physical::ProjectItemPlan::NullOf { + index: 0, + name: "x".to_owned(), + }], + nullable: None, + subqueries: Vec::new(), + }); + assert_eq!(to_message(&plan).unwrap().min_reader_version, 5); + assert_eq!(round_trip(&plan), plan); +} + +#[test] +fn a_corrupt_expand_is_refused() { + use pb::node::Kind; + let edit = |f: &dyn Fn(&mut pb::ExpandNode)| { + let mut m = message("expand_rollup"); + for node in &mut m.nodes { + if let Some(Kind::Expand(e)) = &mut node.kind { + f(e); + } + } + m + }; + corrupt( + &edit(&|e| e.projections.clear()), + "an expand without a projection", + ); + corrupt( + &edit(&|e| { + e.nullable.pop(); + }), + "do not match the nullability", + ); + corrupt( + &edit(&|e| e.projections[0].items[0].kind = None), + "a project item kind", + ); +} + #[test] fn a_newer_writer_that_stays_compatible_is_read() { // N reads N+1 when the writer only added fields: unknown fields (here diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 5945829d..dfa05676 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -19,7 +19,7 @@ syntax = "proto3"; package krishiv.plan.v1; message PhysicalPlan { - // The writer's format version (4). + // The writer's format version (5). uint32 format_version = 1; // The oldest reader version that understands every field in use. uint32 min_reader_version = 2; @@ -69,6 +69,9 @@ message Node { ExchangeNode exchange = 14; // v4: a plan with a window needs a v4 reader. WindowNode window = 15; + // v5: a plan with an expand (ROLLUP, CUBE, GROUPING SETS) needs a v5 + // reader. + ExpandNode expand = 16; } } @@ -98,6 +101,8 @@ message ProjectItem { oneof kind { ProjectColumn column = 1; ProjectExpr expr = 2; + // v5: only inside an expand; a plan holding one needs a v5 reader. + ProjectNullOf null_of = 3; } } @@ -111,6 +116,12 @@ message ProjectExpr { uint32 expr = 2; } +// v5: a column of nulls in the Arrow type of an input column. +message ProjectNullOf { + uint64 index = 1; + string name = 2; +} + message SubqueryPlan { uint64 id = 1; SubqueryValueKind kind = 2; @@ -139,6 +150,21 @@ message ReplicateRowsNode { repeated SubqueryPlan subqueries = 4; } +// v5: each input row once per projection (ROLLUP, CUBE, GROUPING SETS). +message ExpandNode { + uint32 input = 1; + repeated uint64 layout = 2; + // The projections: the same output columns, each computed one way. + repeated ExpandProjection projections = 3; + // Whether each output column is nullable. + repeated bool nullable = 4; + repeated SubqueryPlan subqueries = 5; +} + +message ExpandProjection { + repeated ProjectItem items = 1; +} + message SortKey { uint64 column = 1; bool descending = 2; From 504b4824ef8210f9301f9f41a482e67b5cffb6ee Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:17:43 +0200 Subject: [PATCH 7/9] plan-codec: feature header instead of sequential format versions A plan now lists the named features it uses (sorted string ids) and its wire epoch. A reader checks the epoch and the features before it reads a node, names any it lacks, and refuses a header that does not match the plan's constructs, so the header can be read on its own (features_of). - features.rs: the sorted registry (limit_offset, window), feature_of with exhaustive matches over node kinds, project items, join types, partitionings, window and aggregate functions, and ReaderFeatures. - format_version is frozen at 5; min_reader_version keeps working for readers that predate features (new features use level 6). Plans without an epoch decode by their content. - Goldens: base/ and one directory per feature, byte-exact; v1-v4 frozen. - The decoder maps protobuf enums without wildcards. - plan-protogen --check requires a compatibility-class comment on every schema element added from now on (baseline.txt lists the others). Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .github/workflows/ci.yml | 4 +- crates/krishiv-plan-codec/CLAUDE.md | 28 +- crates/krishiv-plan-codec/src/decode.rs | 347 ++++++++++++----- crates/krishiv-plan-codec/src/encode.rs | 59 ++- crates/krishiv-plan-codec/src/error.rs | 18 +- crates/krishiv-plan-codec/src/features.rs | 312 +++++++++++++++ crates/krishiv-plan-codec/src/lib.rs | 52 ++- .../src/proto/krishiv.plan.v1.pb.rs | 37 +- crates/krishiv-plan-codec/tests/features.rs | 251 ++++++++++++ crates/krishiv-plan-codec/tests/golden.rs | 310 ++++++++++++--- .../tests/golden/base/aggregate.bin | Bin 0 -> 1732 bytes .../tests/golden/base/aggregate_final.bin | Bin 0 -> 850 bytes .../tests/golden/base/exchange_broadcast.bin | Bin 0 -> 756 bytes .../tests/golden/base/exchange_hash.bin | Bin 0 -> 760 bytes .../tests/golden/base/exchange_range.bin | Bin 0 -> 766 bytes .../golden/base/exchange_round_robin.bin | Bin 0 -> 756 bytes .../tests/golden/base/exchange_single.bin | Bin 0 -> 756 bytes .../tests/golden/base/expressions.bin | Bin 0 -> 1472 bytes .../tests/golden/base/file_scan.bin | Bin 0 -> 497 bytes .../tests/golden/base/filter.bin | Bin 0 -> 969 bytes .../tests/golden/base/hash_join.bin | Bin 0 -> 1032 bytes .../tests/golden/base/hash_join_anti.bin | Bin 0 -> 881 bytes .../golden/base/hash_join_partitioned.bin | Bin 0 -> 1032 bytes .../tests/golden/base/limit.bin | Bin 0 -> 761 bytes .../tests/golden/base/nested_loop_join.bin | Bin 0 -> 868 bytes .../tests/golden/base/one_row.bin | Bin 0 -> 388 bytes .../tests/golden/base/project.bin | Bin 0 -> 955 bytes .../golden/base/project_without_nullable.bin | Bin 0 -> 394 bytes .../tests/golden/base/replicate_rows.bin | Bin 0 -> 881 bytes .../tests/golden/base/scan.bin | Bin 0 -> 744 bytes .../tests/golden/base/sort.bin | Bin 0 -> 760 bytes .../tests/golden/base/top_n.bin | Bin 0 -> 762 bytes .../tests/golden/base/union.bin | Bin 0 -> 872 bytes .../tests/golden/base/write_files.bin | Bin 0 -> 750 bytes .../golden/limit_offset/limit_offset.bin | Bin 0 -> 768 bytes .../tests/golden/window/window_aggregates.bin | Bin 0 -> 1800 bytes .../tests/golden/window/window_ranking.bin | Bin 0 -> 870 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 83 ++-- proto/krishiv/plan/v1/plan.proto | 52 ++- tools/plan-protogen/Cargo.lock | 1 + tools/plan-protogen/Cargo.toml | 1 + tools/plan-protogen/baseline.txt | 364 ++++++++++++++++++ tools/plan-protogen/src/classes.rs | 307 +++++++++++++++ tools/plan-protogen/src/main.rs | 31 ++ 44 files changed, 2001 insertions(+), 256 deletions(-) create mode 100644 crates/krishiv-plan-codec/src/features.rs create mode 100644 crates/krishiv-plan-codec/tests/features.rs create mode 100644 crates/krishiv-plan-codec/tests/golden/base/aggregate.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/aggregate_final.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/exchange_broadcast.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/exchange_hash.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/exchange_range.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/exchange_round_robin.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/exchange_single.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/expressions.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/file_scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/filter.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/hash_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/hash_join_anti.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/hash_join_partitioned.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/limit.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/nested_loop_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/one_row.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/project.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/project_without_nullable.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/replicate_rows.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/sort.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/top_n.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/union.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/base/write_files.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/limit_offset/limit_offset.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/window/window_aggregates.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/window/window_ranking.bin create mode 100644 tools/plan-protogen/baseline.txt create mode 100644 tools/plan-protogen/src/classes.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4d3c0afc..819288ff 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -317,7 +317,8 @@ jobs: workspaces: tools/connect-protogen - run: cargo run --locked --manifest-path tools/connect-protogen/Cargo.toml -- --check - # The generated `krishiv.plan.v1` types match proto/krishiv/plan/v1 (P2-T173). + # The generated `krishiv.plan.v1` types match proto/krishiv/plan/v1 (P2-T173), + # and every schema element added since ADR 0040 declares its compatibility class. # `tools/plan-protogen` is a workspace of its own, like connect-protogen. plan-protos: needs: changes @@ -338,6 +339,7 @@ jobs: with: workspaces: tools/plan-protogen - run: cargo run --locked --manifest-path tools/plan-protogen/Cargo.toml -- --check + - run: cargo test --locked --manifest-path tools/plan-protogen/Cargo.toml # The cargo-fuzz targets (`fuzz/`, a workspace of its own) only run nightly, # so a change to an API they call would break them unseen. Build all of them diff --git a/crates/krishiv-plan-codec/CLAUDE.md b/crates/krishiv-plan-codec/CLAUDE.md index 2081ab8f..b9cd29a4 100644 --- a/crates/krishiv-plan-codec/CLAUDE.md +++ b/crates/krishiv-plan-codec/CLAUDE.md @@ -10,22 +10,32 @@ enforcement of design P2-T171 invariant I1 (plans are pure data) and the wire fo ## Invariants - `proto/krishiv/plan/v1/plan.proto` is the schema. `src/proto/krishiv.plan.v1.pb.rs` is **generated, never edited**: `cargo run --locked --manifest-path tools/plan-protogen/Cargo.toml -- --write` (CI runs `--check`; no `protoc`). -- Compatibility is N to N-1 both ways (ADR 0026): add fields with new numbers only, `reserved` the number and name - of anything removed, never reuse an enum value. A change an older reader would silently misread raises - `MIN_READER_VERSION` (and `FORMAT_VERSION`) and adds `tests/golden/v` (see `tests/golden.rs`); old golden - directories are never edited. +- Compatibility is by named features (ADR 0040, `src/features.rs`): the header lists the features a plan uses + (derived from the plan alone), and a reader decodes every plan of its epoch whose features it knows. Add fields + with new numbers only, `reserved` the number and name of anything removed, never reuse an enum value. +- A construct an older reader would misread or reject unhelpfully needs a feature: one **alphabetically sorted** + registry line, the arms of `feature_of` (no wildcards, so new variants do not compile until classified), a + `// feature: ` comment on each new proto element (`ignorable` for what an older reader may skip; + `plan-protogen --check` enforces it against `tools/plan-protogen/baseline.txt`), and `tests/golden//`. + New features take legacy level `POST_MIGRATION_READER_VERSION` (6). +- `format_version` is frozen at 5 and `min_reader_version` is kept for readers that predate features; plans + without an `epoch` are such older plans and decode by their content. `tests/golden/v1`-`v4` are frozen and + decode-only; `base/` and the feature directories are byte-exact against the encoder forever. +- Function behaviour is not a codec matter: functions travel by name and registry id (registry version). - Flat tables with child indices, no recursion in the encoder or decoder; the decoder bounds size, counts and depth (`Limits`), and every reference must point to an earlier entry. - A plan carries `secret(scope, key)` references only (ADR 0038): an inline secret or a URL with credentials fails to encode, and error messages never name a value. -- A new `PhysNode` variant fails to compile in `encode.rs`, `decode.rs` and `tests/common/inventory.rs` - (`variant_of`), and the inventory test fails until a case of it exists. Adding a case changes the golden files: - they only grow for the current version. +- A new `PhysNode` variant fails to compile in `encode.rs`, `decode.rs`, `features.rs` and + `tests/roundtrip.rs` (`variant_of`), and the inventory test fails until a case of it exists; a new case gets its + golden file with `KRISHIV_BLESS_GOLDEN=1` (then rerun the tests). The decoder maps protobuf enums without + wildcards (`enum_map!`). - `#![forbid(unsafe_code)]`. ## ADRs (in `../krishiv-internal/adr/`) -0003, 0026, 0038 +0003, 0026, 0038, 0040 ## Test `cargo nextest run -p krishiv-plan-codec --no-tests=pass` · tests: `roundtrip` (inventory, secrets, corrupt input), -`golden` (compatibility bytes), `ratchet` (TPC-H 22 and TPC-DS 99 planned plans). +`features` (header, reader checks), `golden` (compatibility bytes, simulated older readers, header honesty), +`ratchet` (TPC-H 22 and TPC-DS 99 planned plans). diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 552f70c8..391aa010 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -4,7 +4,7 @@ //! entry may be used once (nodes) and counts must agree. The tables are //! flat, so nothing here recurses and protobuf nesting stays shallow. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use krishiv_common::secret::{OptionValue, SecretValue}; @@ -31,25 +31,34 @@ use krishiv_types::{DataType, DecimalType}; use crate::arrow::{schema_from, type_from}; use crate::error::{CodecError, Result, corrupt}; +use crate::features::{Feature, ReaderFeatures, legacy_level, plan_features}; use crate::proto as pb; -use crate::{FORMAT_VERSION, Limits}; +use crate::{EPOCH, FORMAT_VERSION, Limits}; -/// Decodes `krishiv.plan.v1` bytes under the default [`Limits`]. +/// Decodes `krishiv.plan.v1` bytes under the default [`Limits`], with every +/// feature this build implements. /// /// # Errors -/// [`CodecError::Corrupt`] for bytes that are not a valid plan, -/// [`CodecError::TooNew`] for a plan that needs a newer reader, +/// [`CodecError::Corrupt`] for bytes that are not a valid plan (a header +/// that does not match the plan's constructs included), +/// [`CodecError::UnsupportedEpoch`] and [`CodecError::UnsupportedFeatures`] +/// for a plan this reader cannot read, [`CodecError::TooNew`] for a plan +/// written before the feature header that needs a newer reader, /// [`CodecError::Limit`] for a plan beyond the limits. pub fn decode(bytes: &[u8]) -> Result { - decode_with(bytes, &Limits::default()) + decode_with(bytes, &Limits::default(), &ReaderFeatures::all()) } -/// [`decode`] under `limits`. +/// [`decode`] under `limits`, as a reader that implements `reader`. /// /// # Errors /// As [`decode`]. -pub fn decode_with(bytes: &[u8], limits: &Limits) -> Result { - use prost::Message as _; +pub fn decode_with(bytes: &[u8], limits: &Limits, reader: &ReaderFeatures) -> Result { + check_size(bytes, limits)?; + from_message(&parse(bytes)?, limits, reader) +} + +fn check_size(bytes: &[u8], limits: &Limits) -> Result<()> { if bytes.len() > limits.max_bytes { return Err(CodecError::Limit(format!( "{} bytes (at most {})", @@ -57,30 +66,164 @@ pub fn decode_with(bytes: &[u8], limits: &Limits) -> Result { limits.max_bytes ))); } + Ok(()) +} + +fn parse(bytes: &[u8]) -> Result { // The messages are flat tables: protobuf nesting stays under a dozen // levels, whatever the plan's depth, so prost's recursion is bounded. #[expect( clippy::disallowed_methods, reason = "the plan format has constant nesting depth (flat tables)" )] - let message = pb::PhysicalPlan::decode(bytes).map_err(|_| corrupt("malformed protobuf"))?; - from_message(&message, limits) + M::decode(bytes).map_err(|_| corrupt("malformed protobuf")) } -/// [`decode_with`] from a parsed message. +/// The features an encoded plan needs, from its header. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PlanFeatures { + /// The plan's wire epoch. + pub epoch: u32, + /// The ids of the features it uses, sorted. + pub features: Vec, +} + +impl PlanFeatures { + /// The features `reader` lacks to read the plan. + #[must_use] + pub fn missing(&self, reader: &ReaderFeatures) -> Vec { + reader.missing(self.features.iter().map(String::as_str)) + } +} + +/// The epoch and features of an encoded plan, read from its header without +/// decoding the plan (the header is checked against the plan on decode, so +/// it can be trusted). A plan written before the header existed is decoded +/// under the default [`Limits`] to find its features. /// /// # Errors -/// As [`decode`]. -pub fn from_message(m: &pb::PhysicalPlan, limits: &Limits) -> Result { - if m.format_version == 0 { +/// [`CodecError::Corrupt`] for a malformed header, +/// [`CodecError::UnsupportedEpoch`] for an epoch this build does not know, +/// and for an older plan, as [`decode`]. +pub fn features_of(bytes: &[u8]) -> Result { + check_size(bytes, &Limits::default())?; + let h: pb::PlanHeader = parse(bytes)?; + let features = match header(h.format_version, h.min_reader_version, h.epoch, &h.features)? { + Header::Current => h.features, + Header::Legacy => plan_features(&decode(bytes)?) + .iter() + .map(|f| f.id().to_owned()) + .collect(), + }; + Ok(PlanFeatures { + epoch: EPOCH, + features, + }) +} + +/// What a valid header says about the rest of the plan. +enum Header { + /// A plan with a feature header: the ids are well formed and sorted. + Current, + /// A plan written before the feature header: its features follow from + /// its content. + Legacy, +} + +/// The longest feature id a reader accepts. +const MAX_FEATURE_ID: usize = 64; + +fn header(format_version: u32, min_reader: u32, epoch: u32, ids: &[String]) -> Result
{ + if format_version == 0 { return Err(corrupt("no format version")); } - if m.min_reader_version > FORMAT_VERSION { - return Err(CodecError::TooNew { - required: m.min_reader_version, - supported: FORMAT_VERSION, + if epoch == 0 { + if !ids.is_empty() { + return Err(corrupt("a feature list without an epoch")); + } + if min_reader > FORMAT_VERSION { + return Err(CodecError::TooNew { + required: min_reader, + supported: FORMAT_VERSION, + }); + } + return Ok(Header::Legacy); + } + if epoch != EPOCH { + return Err(CodecError::UnsupportedEpoch { + epoch, + supported: EPOCH, }); } + for id in ids { + let well_formed = id.len() <= MAX_FEATURE_ID + && id.starts_with(|c: char| c.is_ascii_lowercase()) + && id + .chars() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_'); + if !well_formed { + return Err(corrupt("a malformed feature id")); + } + } + if ids.windows(2).any(|w| w[0] >= w[1]) { + return Err(corrupt("feature ids that are not sorted or repeat")); + } + Ok(Header::Current) +} + +/// [`decode_with`] from a parsed message. +/// +/// # Errors +/// As [`decode`]. +pub fn from_message( + m: &pb::PhysicalPlan, + limits: &Limits, + reader: &ReaderFeatures, +) -> Result { + let kind = header(m.format_version, m.min_reader_version, m.epoch, &m.features)?; + if let Header::Current = kind { + // Before any node is read. + let missing = reader.missing(m.features.iter().map(String::as_str)); + if !missing.is_empty() { + return Err(CodecError::UnsupportedFeatures { missing }); + } + } + let plan = plan(m, limits)?; + let used = plan_features(&plan); + match kind { + Header::Current => { + let declared: BTreeSet = m + .features + .iter() + .filter_map(|id| Feature::from_id(id)) + .collect(); + if declared != used { + return Err(corrupt( + "the feature header does not match the plan's constructs", + )); + } + if m.min_reader_version != legacy_level(&used) { + return Err(corrupt("a reader version that does not match the features")); + } + } + Header::Legacy => { + if used.iter().any(|f| f.legacy_level() > m.min_reader_version) { + return Err(corrupt("a construct newer than the plan's reader version")); + } + let missing: Vec = used + .iter() + .filter(|f| !reader.contains(**f)) + .map(|f| f.id().to_owned()) + .collect(); + if !missing.is_empty() { + return Err(CodecError::UnsupportedFeatures { missing }); + } + } + } + Ok(plan) +} + +fn plan(m: &pb::PhysicalPlan, limits: &Limits) -> Result { for (what, n, max) in [ ("nodes", m.nodes.len(), limits.max_nodes), ("explain nodes", m.explain.len(), limits.max_nodes), @@ -186,7 +329,7 @@ fn expr_children(m: &pb::Expr) -> impl Iterator + '_ { Some(Kind::Cast(c)) => (Some(c.expr), &[]), Some(Kind::Alias(a)) => (Some(a.expr), &[]), Some(Kind::Function(f)) => (f.filter, &f.args), - _ => (None, &[]), + Some(Kind::Literal(_) | Kind::Column(_) | Kind::OuterColumn(_)) | None => (None, &[]), }; one.into_iter().chain(many.iter().copied()) } @@ -219,11 +362,13 @@ fn required<'a, T>(field: &'a Option, what: &str) -> Result<&'a T> { .ok_or_else(|| corrupt(format!("{what} is missing"))) } -macro_rules! enum_of { - ($ty:ty, $value:expr, $what:literal) => { +/// Maps a protobuf enum value with an exhaustive match: a value added to the +/// schema does not compile here until it is mapped. +macro_rules! enum_map { + ($ty:ty, $value:expr, $what:literal, { $($from:ident => $to:expr),+ $(,)? }) => { match <$ty>::try_from($value) { + $(Ok(<$ty>::$from) => Ok($to),)+ Ok(<$ty>::Unspecified) | Err(_) => Err(corrupt(concat!("an unknown ", $what))), - Ok(v) => Ok(v), } }; } @@ -342,10 +487,10 @@ impl Decoder { .map(|s| { Ok(SubqueryPlan { id: SubqueryId(s.id), - kind: match enum_of!(pb::SubqueryValueKind, s.kind, "subquery kind")? { - pb::SubqueryValueKind::Exists => SubqueryValueKind::Exists, - _ => SubqueryValueKind::Scalar, - }, + kind: enum_map!(pb::SubqueryValueKind, s.kind, "subquery kind", { + Scalar => SubqueryValueKind::Scalar, + Exists => SubqueryValueKind::Exists, + })?, data_type: type_from(&s.data_type)?, plan: self.take(s.plan, self.slots.len())?, }) @@ -460,11 +605,11 @@ impl Decoder { Kind::Cast(c) => ExprNode::Cast { expr: self.expr_ref(c.expr)?, to: spark_type(required(&c.to, "a cast type")?)?, - mode: match enum_of!(pb::CastMode, c.mode, "cast mode")? { - pb::CastMode::Legacy => CastMode::Legacy, - pb::CastMode::Ansi => CastMode::Ansi, - _ => CastMode::Try, - }, + mode: enum_map!(pb::CastMode, c.mode, "cast mode", { + Legacy => CastMode::Legacy, + Ansi => CastMode::Ansi, + Try => CastMode::Try, + })?, }, Kind::Alias(a) => ExprNode::Resolved(ResolvedExpr::Alias { expr: self.expr_ref(a.expr)?, @@ -496,12 +641,12 @@ impl Decoder { }; Ok(ResolvedFunction { id, - kind: match enum_of!(pb::FunctionKind, f.kind, "function kind")? { - pb::FunctionKind::Scalar => FunctionKind::Scalar, - pb::FunctionKind::Aggregate => FunctionKind::Aggregate, - pb::FunctionKind::Window => FunctionKind::Window, - _ => FunctionKind::Generator, - }, + kind: enum_map!(pb::FunctionKind, f.kind, "function kind", { + Scalar => FunctionKind::Scalar, + Aggregate => FunctionKind::Aggregate, + Window => FunctionKind::Window, + Generator => FunctionKind::Generator, + })?, args: self.exprs_of(&f.args)?, distinct: f.distinct, filter: f.filter.map(|i| self.expr_ref(i)).transpose()?, @@ -538,10 +683,10 @@ fn aggregate(a: &pb::AggregateNode) -> Result { .map(aggregate_expr) .collect::>()?, distinct: a.distinct.clone(), - output: match enum_of!(pb::AggregateOutput, a.output, "aggregate output")? { - pb::AggregateOutput::Final => AggregateOutput::Final, - _ => AggregateOutput::Partial, - }, + output: enum_map!(pb::AggregateOutput, a.output, "aggregate output", { + Final => AggregateOutput::Final, + Partial => AggregateOutput::Partial, + })?, }) } @@ -575,19 +720,19 @@ fn window_function(f: &pb::WindowFunctionPlan) -> Result { }; let frame = required(&f.frame, "a window frame")?; let bound = |b: i32| { - Ok(match enum_of!(pb::FrameBound, b, "window frame bound")? { - pb::FrameBound::UnboundedPreceding => FrameBound::UnboundedPreceding, - pb::FrameBound::CurrentRow => FrameBound::CurrentRow, - _ => FrameBound::UnboundedFollowing, + enum_map!(pb::FrameBound, b, "window frame bound", { + UnboundedPreceding => FrameBound::UnboundedPreceding, + CurrentRow => FrameBound::CurrentRow, + UnboundedFollowing => FrameBound::UnboundedFollowing, }) }; Ok(WindowFunctionPlan { function, frame: WindowFrame { - units: match enum_of!(pb::FrameUnits, frame.units, "window frame unit")? { - pb::FrameUnits::Rows => FrameUnits::Rows, - _ => FrameUnits::Range, - }, + units: enum_map!(pb::FrameUnits, frame.units, "window frame unit", { + Rows => FrameUnits::Rows, + Range => FrameUnits::Range, + })?, lower: bound(frame.lower)?, upper: bound(frame.upper)?, }, @@ -599,39 +744,41 @@ fn window_function(f: &pb::WindowFunctionPlan) -> Result { fn aggregate_expr(a: &pb::AggregateExpr) -> Result { use pb::AggregateFunction as F; let ignore = a.ignore_nulls; - let function = match enum_of!(F, a.function, "aggregate function")? { - F::First => AggregateFunction::First { - ignore_nulls: ignore, - }, - F::Last => AggregateFunction::Last { - ignore_nulls: ignore, - }, - _ if ignore => return Err(corrupt("`ignore_nulls` on a function without it")), - F::Sum => AggregateFunction::Sum, - F::Avg => AggregateFunction::Avg, - F::Min => AggregateFunction::Min, - F::Max => AggregateFunction::Max, - _ => AggregateFunction::Count, - }; + let function = enum_map!(F, a.function, "aggregate function", { + Count => AggregateFunction::Count, + Sum => AggregateFunction::Sum, + Avg => AggregateFunction::Avg, + Min => AggregateFunction::Min, + Max => AggregateFunction::Max, + First => AggregateFunction::First { ignore_nulls: ignore }, + Last => AggregateFunction::Last { ignore_nulls: ignore }, + })?; + let takes_ignore = matches!( + function, + AggregateFunction::First { .. } | AggregateFunction::Last { .. } + ); + if ignore && !takes_ignore { + return Err(corrupt("`ignore_nulls` on a function without it")); + } Ok(AggregateExpr { function, args: usizes(&a.args)?, - input: match enum_of!(pb::AggregateInput, a.input, "aggregate input")? { - pb::AggregateInput::Raw => AggInput::Raw, - _ => AggInput::State, - }, + input: enum_map!(pb::AggregateInput, a.input, "aggregate input", { + Raw => AggInput::Raw, + State => AggInput::State, + })?, arg_type: type_from(&a.arg_type)?, name: a.name.clone(), }) } fn join_type(t: i32) -> Result { - Ok(match enum_of!(pb::JoinType, t, "join type")? { - pb::JoinType::Inner => JoinType::Inner, - pb::JoinType::LeftOuter => JoinType::LeftOuter, - pb::JoinType::LeftSemi => JoinType::LeftSemi, - pb::JoinType::LeftAnti => JoinType::LeftAnti, - _ => JoinType::NullAwareAnti, + enum_map!(pb::JoinType, t, "join type", { + Inner => JoinType::Inner, + LeftOuter => JoinType::LeftOuter, + LeftSemi => JoinType::LeftSemi, + LeftAnti => JoinType::LeftAnti, + NullAwareAnti => JoinType::NullAwareAnti, }) } @@ -720,26 +867,26 @@ fn option(m: &pb::OptionValue) -> Result { fn spark_type(t: &pb::DataType) -> Result { use pb::data_type::Kind; - let kind = enum_of!(Kind, t.kind, "type")?; - if kind != Kind::Decimal && (t.precision != 0 || t.scale != 0) { + let data_type = enum_map!(Kind, t.kind, "type", { + Null => DataType::Null, + Boolean => DataType::Boolean, + Byte => DataType::Byte, + Short => DataType::Short, + Integer => DataType::Integer, + Long => DataType::Long, + Float => DataType::Float, + Double => DataType::Double, + String => DataType::String, + Binary => DataType::Binary, + Decimal => DataType::Decimal(decimal(t)?), + Date => DataType::Date, + Timestamp => DataType::Timestamp, + TimestampNtz => DataType::TimestampNtz, + })?; + if !matches!(data_type, DataType::Decimal(_)) && (t.precision != 0 || t.scale != 0) { return Err(corrupt("a precision or scale on a type that has none")); } - Ok(match kind { - Kind::Boolean => DataType::Boolean, - Kind::Byte => DataType::Byte, - Kind::Short => DataType::Short, - Kind::Integer => DataType::Integer, - Kind::Long => DataType::Long, - Kind::Float => DataType::Float, - Kind::Double => DataType::Double, - Kind::String => DataType::String, - Kind::Binary => DataType::Binary, - Kind::Decimal => DataType::Decimal(decimal(t)?), - Kind::Date => DataType::Date, - Kind::Timestamp => DataType::Timestamp, - Kind::TimestampNtz => DataType::TimestampNtz, - _ => DataType::Null, - }) + Ok(data_type) } fn decimal(t: &pb::DataType) -> Result { @@ -797,17 +944,17 @@ fn literal(l: &pb::Literal) -> Result { } fn year_month_field(v: i32) -> Result { - Ok(match enum_of!(pb::YearMonthField, v, "interval field")? { - pb::YearMonthField::Year => YearMonthField::Year, - _ => YearMonthField::Month, + enum_map!(pb::YearMonthField, v, "interval field", { + Year => YearMonthField::Year, + Month => YearMonthField::Month, }) } fn day_time_field(v: i32) -> Result { - Ok(match enum_of!(pb::DayTimeField, v, "interval field")? { - pb::DayTimeField::Day => DayTimeField::Day, - pb::DayTimeField::Hour => DayTimeField::Hour, - pb::DayTimeField::Minute => DayTimeField::Minute, - _ => DayTimeField::Second, + enum_map!(pb::DayTimeField, v, "interval field", { + Day => DayTimeField::Day, + Hour => DayTimeField::Hour, + Minute => DayTimeField::Minute, + Second => DayTimeField::Second, }) } diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index 43ad12dd..a483c798 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -26,8 +26,9 @@ use krishiv_types::DataType; use crate::arrow::{schema_bytes, type_bytes}; use crate::error::{CodecError, Result}; +use crate::features::{legacy_level, plan_features}; use crate::proto as pb; -use crate::{FORMAT_VERSION, MIN_READER_VERSION}; +use crate::{EPOCH, FORMAT_VERSION}; /// Encodes `plan` as `krishiv.plan.v1` bytes. Deterministic: the same plan /// always gives the same bytes. @@ -43,11 +44,12 @@ pub fn encode(plan: &PhysicalPlan) -> Result> { } /// [`encode`] without the final serialization. +/// +/// # Errors +/// As [`encode`]. pub fn to_message(plan: &PhysicalPlan) -> Result { - let mut e = Encoder { - min_reader: MIN_READER_VERSION, - ..Encoder::default() - }; + let features = plan_features(plan); + let mut e = Encoder::default(); let order = post_order(plan.root(), PhysNode::children); for node in &order { let message = e.node(node)?; @@ -73,7 +75,7 @@ pub fn to_message(plan: &PhysicalPlan) -> Result { } Ok(pb::PhysicalPlan { format_version: FORMAT_VERSION, - min_reader_version: e.min_reader, + min_reader_version: legacy_level(&features), nodes: e.table, root, schema: schema_bytes(plan.schema()), @@ -82,6 +84,8 @@ pub fn to_message(plan: &PhysicalPlan) -> Result { warnings: plan.warnings().to_vec(), env: Some(env(plan.env())), exprs: e.exprs, + features: features.iter().map(|f| f.id().to_owned()).collect(), + epoch: EPOCH, }) } @@ -125,9 +129,6 @@ struct Encoder { /// The index of each written expression, by `Arc` identity, so a /// shared subexpression is written once. expr_ids: HashMap<*const Expr, u32>, - /// The oldest reader version that understands every field written so - /// far (a `Limit` with a nonzero offset needs version 3, a `Window` version 4). - min_reader: u32, } fn layout(layout: &[AttrId]) -> Vec { @@ -148,20 +149,16 @@ impl Encoder { fn node(&mut self, node: &PhysNode) -> Result { use pb::node::Kind; let kind = match node { - PhysNode::Window { input, window } => { - // An older reader has no `Window` node to read. - self.min_reader = self.min_reader.max(4); - Kind::Window(pb::WindowNode { - input: self.node_index(input)?, - partition_by: indexes(&window.partition_by), - order_by: window.order_by.iter().map(sort_key).collect(), - functions: window - .functions - .iter() - .map(window_function) - .collect::>()?, - }) - } + PhysNode::Window { input, window } => Kind::Window(pb::WindowNode { + input: self.node_index(input)?, + partition_by: indexes(&window.partition_by), + order_by: window.order_by.iter().map(sort_key).collect(), + functions: window + .functions + .iter() + .map(window_function) + .collect::>()?, + }), PhysNode::Scan(scan) => Kind::Scan(self.scan(scan)?), PhysNode::OneRow => Kind::OneRow(pb::Empty {}), PhysNode::Filter { @@ -197,17 +194,11 @@ impl Encoder { input, offset, limit, - } => { - if *offset != 0 { - // An older reader would silently ignore the offset. - self.min_reader = self.min_reader.max(3); - } - Kind::Limit(pb::LimitNode { - input: self.node_index(input)?, - limit: *limit, - offset: *offset, - }) - } + } => Kind::Limit(pb::LimitNode { + input: self.node_index(input)?, + limit: *limit, + offset: *offset, + }), PhysNode::ReplicateRows { input, layout: l, diff --git a/crates/krishiv-plan-codec/src/error.rs b/crates/krishiv-plan-codec/src/error.rs index 1eab60ba..51e28e03 100644 --- a/crates/krishiv-plan-codec/src/error.rs +++ b/crates/krishiv-plan-codec/src/error.rs @@ -15,7 +15,8 @@ pub enum CodecError { /// protobuf, a reference out of range, a count that does not match. #[error("corrupt plan: {0}")] Corrupt(String), - /// The plan uses features of a newer format than this reader knows. + /// A plan written before the feature header that needs a newer reader + /// version than the frozen format version. #[error("the plan needs reader version {required}, this reader implements {supported}")] TooNew { /// The writer's `min_reader_version`. @@ -23,6 +24,21 @@ pub enum CodecError { /// The version this reader implements. supported: u32, }, + /// The plan is of a wire epoch this reader does not implement. + #[error("the plan is of epoch {epoch}, this reader implements epoch {supported}")] + UnsupportedEpoch { + /// The plan's epoch. + epoch: u32, + /// The epoch this reader implements. + supported: u32, + }, + /// The plan uses features this reader does not implement (ids unknown to + /// it, or left out of its [`ReaderFeatures`](crate::ReaderFeatures)). + #[error("the plan uses features [{}] that this reader does not implement", missing.join(", "))] + UnsupportedFeatures { + /// The ids of the missing features, sorted. + missing: Vec, + }, /// The plan is larger or deeper than the decoder's [`Limits`](crate::Limits). #[error("the plan exceeds a decode limit: {0}")] Limit(String), diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs new file mode 100644 index 00000000..31dd6bcf --- /dev/null +++ b/crates/krishiv-plan-codec/src/features.rs @@ -0,0 +1,312 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Plan features (ADR 0040): what a reader must implement to read a plan. +//! +//! A construct needs a feature when a reader without it would misread it or +//! reject it with an unhelpful error: a new node kind, a new project item or +//! enum value, a new field whose absence changes results. Everything else is +//! *ignorable* and needs no feature: an older reader skips the unknown field +//! (a hint, an explain detail) and still computes the same result. There is +//! no third class. +//! +//! The header of an encoded plan lists exactly the features the plan uses, +//! computed from the plan alone ([`plan_features`]), so the same plan always +//! gets the same header wherever it is encoded. A reader checks the header +//! before it reads a node and refuses a plan that uses a feature it lacks, +//! naming it ([`CodecError::UnsupportedFeatures`](crate::CodecError)). +//! +//! # Adding a feature +//! 1. One line in the registry below, at its **alphabetical** place (two +//! changes adding features then touch different lines). Ids are +//! `snake_case`, never renamed and never reused; a retired id moves to +//! [`RETIRED`], which happens only at an epoch change. +//! 2. The arms of [`feature_of`] for the new constructs: its matches have no +//! wildcard, so a new node kind, project item, join type, partitioning, +//! window function or aggregate function does not compile until it is +//! classified. +//! 3. A golden directory `tests/golden//` with the cases that use it. +//! 4. A `// feature: ` annotation on each new proto element +//! (`tools/plan-protogen --check` enforces it). +//! +//! A new feature's legacy level is [`POST_MIGRATION_READER_VERSION`]. +//! +//! Semantic drift inside an existing construct (a function that computes +//! something else under the same name) is not a codec matter: functions +//! travel by name and registry id, and the function registry's version +//! covers them. + +use std::collections::BTreeSet; +use std::fmt; + +use krishiv_exec::JoinType; +use krishiv_exec::aggregate::AggregateFunction; +use krishiv_exec::physical::{Partitioning, PhysNode, PhysicalPlan, ProjectItemPlan}; +use krishiv_exec::window::WindowFunction; + +/// The `min_reader_version` of a plan that uses no feature (the epoch-1 +/// baseline): the format of version 2. +pub const MIN_READER_VERSION: u32 = 2; + +/// The legacy `min_reader_version` of every feature added after the format +/// version was frozen: above every frozen format version, so a reader that +/// predates features refuses the plan instead of misreading it. +pub const POST_MIGRATION_READER_VERSION: u32 = 6; + +/// Whether the planner can avoid a feature, which decides what happens while +/// some workers lack it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum FeatureKind { + /// The planner has an equivalent plan without it: it masks the feature + /// while a live worker lacks it. + Optimization, + /// No alternative plan exists: such a plan runs only on workers that + /// have the feature. + Semantic, +} + +macro_rules! registry { + ($(($id:ident, $variant:ident, $legacy:expr, $kind:ident, $meaning:literal),)*) => { + /// A plan feature. The order is the registry's (alphabetical by id). + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub enum Feature { + $(#[doc = $meaning] $variant,)* + } + + impl Feature { + /// Every feature this reader implements, in registry order. + pub const ALL: &[Self] = &[$(Self::$variant),*]; + + /// The stable id written in a plan's header. + #[must_use] + pub const fn id(self) -> &'static str { + match self { + $(Self::$variant => stringify!($id),)* + } + } + + /// The `min_reader_version` a plan using this feature carries, for + /// readers that predate features. + #[must_use] + pub const fn legacy_level(self) -> u32 { + match self { + $(Self::$variant => $legacy,)* + } + } + + /// Whether the planner can avoid it. + #[must_use] + pub const fn kind(self) -> FeatureKind { + match self { + $(Self::$variant => FeatureKind::$kind,)* + } + } + + /// What it covers, in one line. + #[must_use] + pub const fn meaning(self) -> &'static str { + match self { + $(Self::$variant => $meaning,)* + } + } + } + }; +} + +// The registry: one line per feature, kept alphabetically sorted by id. +// (id, variant, legacy level, kind, meaning) +registry! { + (limit_offset, LimitOffset, 3, Semantic, "A `Limit` that skips leading rows (a nonzero offset)."), + (window, Window, 4, Semantic, "The `Window` node and its functions."), +} + +/// Ids retired at an epoch change: never reused. +pub const RETIRED: &[&str] = &[]; + +impl Feature { + /// The feature with id `id`, if this reader implements it. + #[must_use] + pub fn from_id(id: &str) -> Option { + Self::ALL.iter().copied().find(|f| f.id() == id) + } +} + +impl fmt::Display for Feature { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.id()) + } +} + +/// A construct of a plan that may need a feature. +#[derive(Debug, Clone, Copy)] +#[non_exhaustive] +pub enum Construct<'a> { + /// A plan node (not its children). + Node(&'a PhysNode), + /// A project item. + Item(&'a ProjectItemPlan), + /// A join type. + Join(JoinType), + /// An exchange's partitioning. + Partitioning(&'a Partitioning), + /// A window function. + Window(&'a WindowFunction), + /// An aggregate function (of an aggregate or a window). + Aggregate(AggregateFunction), +} + +/// The feature `construct` needs, or `None` for the epoch-1 baseline. +#[must_use] +pub fn feature_of(construct: Construct<'_>) -> Option { + match construct { + Construct::Node(node) => match node { + PhysNode::Limit { offset, .. } => (*offset != 0).then_some(Feature::LimitOffset), + PhysNode::Window { .. } => Some(Feature::Window), + PhysNode::Scan(_) + | PhysNode::OneRow + | PhysNode::Filter { .. } + | PhysNode::Project { .. } + | PhysNode::ReplicateRows { .. } + | PhysNode::Sort { .. } + | PhysNode::TopN { .. } + | PhysNode::Aggregate { .. } + | PhysNode::HashJoin(_) + | PhysNode::NestedLoopJoin(_) + | PhysNode::Union { .. } + | PhysNode::WriteFiles { .. } + | PhysNode::Exchange(_) => None, + }, + Construct::Item(item) => match item { + ProjectItemPlan::Column { .. } | ProjectItemPlan::Expr { .. } => None, + }, + Construct::Join(join) => match join { + JoinType::Inner + | JoinType::LeftOuter + | JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::NullAwareAnti => None, + }, + Construct::Partitioning(p) => match p { + Partitioning::Hash { .. } + | Partitioning::Single + | Partitioning::Broadcast + | Partitioning::RoundRobin + | Partitioning::Range { .. } => None, + }, + // Every window function came with the `Window` node. + Construct::Window(w) => match w { + WindowFunction::RowNumber + | WindowFunction::Rank + | WindowFunction::DenseRank + | WindowFunction::PercentRank + | WindowFunction::CumeDist + | WindowFunction::Ntile(_) + | WindowFunction::Aggregate(_) => Some(Feature::Window), + }, + Construct::Aggregate(a) => match a { + AggregateFunction::Count + | AggregateFunction::Sum + | AggregateFunction::Avg + | AggregateFunction::Min + | AggregateFunction::Max + | AggregateFunction::First { .. } + | AggregateFunction::Last { .. } => None, + }, + } +} + +/// Calls `f` with `node` and each construct inside it (not its children). +fn constructs_of<'a>(node: &'a PhysNode, f: &mut impl FnMut(Construct<'a>)) { + f(Construct::Node(node)); + match node { + PhysNode::Project { items, .. } => items.iter().for_each(|i| f(Construct::Item(i))), + PhysNode::Window { window, .. } => { + for function in &window.functions { + f(Construct::Window(&function.function)); + if let WindowFunction::Aggregate(a) = &function.function { + f(Construct::Aggregate(a.function)); + } + } + } + PhysNode::Aggregate { aggregate, .. } => aggregate + .aggregates + .iter() + .for_each(|a| f(Construct::Aggregate(a.function))), + PhysNode::HashJoin(j) => f(Construct::Join(j.join_type)), + PhysNode::NestedLoopJoin(j) => f(Construct::Join(j.join_type)), + PhysNode::Exchange(e) => f(Construct::Partitioning(&e.partitioning)), + PhysNode::Scan(_) + | PhysNode::OneRow + | PhysNode::Filter { .. } + | PhysNode::Limit { .. } + | PhysNode::ReplicateRows { .. } + | PhysNode::Sort { .. } + | PhysNode::TopN { .. } + | PhysNode::Union { .. } + | PhysNode::WriteFiles { .. } => {} + } +} + +/// The features `plan` uses: a pure function of its content. Subquery plans +/// count; nothing recurses. +#[must_use] +pub fn plan_features(plan: &PhysicalPlan) -> BTreeSet { + let mut out = BTreeSet::new(); + let mut stack = vec![plan.root()]; + while let Some(node) = stack.pop() { + constructs_of(node, &mut |c| out.extend(feature_of(c))); + stack.extend(node.children()); + } + out +} + +/// The legacy `min_reader_version` of a plan that uses `features`. +pub(crate) fn legacy_level(features: &BTreeSet) -> u32 { + features + .iter() + .map(|f| f.legacy_level()) + .fold(MIN_READER_VERSION, u32::max) +} + +/// The features a reader implements: by default every compiled-in one. A +/// smaller set simulates an older or partial reader. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ReaderFeatures { + features: BTreeSet, +} + +impl Default for ReaderFeatures { + fn default() -> Self { + Self::all() + } +} + +impl ReaderFeatures { + /// Every feature this build implements. + #[must_use] + pub fn all() -> Self { + Self::only(Feature::ALL.iter().copied()) + } + + /// Exactly `features`. + #[must_use] + pub fn only(features: impl IntoIterator) -> Self { + Self { + features: features.into_iter().collect(), + } + } + + /// Whether the reader implements `feature`. + #[must_use] + pub fn contains(&self, feature: Feature) -> bool { + self.features.contains(&feature) + } + + /// The ids of `ids` this reader does not implement, unknown ids + /// included, in their order. + #[must_use] + pub fn missing<'a>(&self, ids: impl IntoIterator) -> Vec { + ids.into_iter() + .filter(|id| Feature::from_id(id).is_none_or(|f| !self.contains(f))) + .map(str::to_owned) + .collect() + } +} diff --git a/crates/krishiv-plan-codec/src/lib.rs b/crates/krishiv-plan-codec/src/lib.rs index a451f1ff..bde8c5a0 100644 --- a/crates/krishiv-plan-codec/src/lib.rs +++ b/crates/krishiv-plan-codec/src/lib.rs @@ -15,11 +15,23 @@ //! generated from it by `tools/plan-protogen` (CI checks they are //! current). Fields are only added with new numbers; removed ones are //! `reserved`. -//! - Compatibility is N to N-1 in both directions (ADR 0026): unknown -//! fields are skipped; a writer that relies on something an older reader -//! would ignore raises `min_reader_version`, and a reader older than that -//! rejects the plan ([`CodecError::TooNew`]). The golden plans under -//! `tests/golden/` pin the bytes of every released format version. +//! - Compatibility is by named features (ADR 0040, [`features`]): the +//! header lists the [`Feature`]s a plan uses, and a reader decodes every +//! plan of its epoch whose features it knows. A reader checks the epoch +//! and the features before it reads a node +//! ([`CodecError::UnsupportedEpoch`], [`CodecError::UnsupportedFeatures`]), +//! and a plan that uses a construct its header does not declare (or +//! declares one it does not use) is corrupt, so [`features_of`] can trust +//! the header without decoding the plan. Unknown fields are skipped: an +//! addition that an older reader may ignore needs no feature. +//! - Plans written before the header (format versions 1 to 5, no epoch) +//! still decode: their features follow from their content, and their +//! `min_reader_version` gate still applies ([`CodecError::TooNew`]). New +//! plans keep writing `format_version` 5 and a `min_reader_version` that +//! makes such older readers refuse what they cannot read. +//! - The golden plans under `tests/golden/` pin the bytes: `base/` for plans +//! without features, one directory per feature, and the frozen `v` +//! directories of the sequential versions. //! - Nodes, expressions and explain nodes are flat tables in post-order //! with child indices, so nothing recurses (deep plans are safe) and an //! expression shared through `Arc` is stored once. @@ -39,6 +51,16 @@ //! [`decode_with`] bounds the size, the table sizes and the plan and expression depth //! ([`Limits`]), requires every reference to point to an earlier entry, //! every node to be used once and counts to agree. +//! +//! # Plans in other records +//! The header is derived from the plan alone, never from the cluster, so a +//! fragment's bytes are the same wherever it runs. A record that stores +//! plans to run later stores their [`PlanFeatures`] with them: checkpoint +//! manifests (P4-T011) and job records (P5-T002) do not exist yet and take +//! it when they are built, so a restore or a handoff is refused before it +//! reaches a reader that lacks a feature. A cache of plans keyed by query +//! text must include the reader capabilities the plan was built for in its +//! key; a cache keyed by the encoded bytes needs nothing more. #![forbid(unsafe_code)] @@ -46,19 +68,27 @@ mod arrow; mod decode; mod encode; mod error; +pub mod features; /// The generated protobuf types of `krishiv.plan.v1`. pub mod proto; -pub use decode::{decode, decode_with, from_message}; +pub use decode::{PlanFeatures, decode, decode_with, features_of, from_message}; pub use encode::{encode, to_message}; pub use error::{CodecError, Result}; +pub use features::{ + Feature, FeatureKind, MIN_READER_VERSION, POST_MIGRATION_READER_VERSION, ReaderFeatures, + plan_features, +}; -/// The format version this crate writes. -pub const FORMAT_VERSION: u32 = 4; +/// The `format_version` this crate writes: frozen at the last sequential +/// version. Readers that predate features compare `min_reader_version` +/// with their own format version. +pub const FORMAT_VERSION: u32 = 5; -/// The `min_reader_version` this crate writes: readers older than this -/// cannot understand the plan. -pub const MIN_READER_VERSION: u32 = 2; +/// The wire epoch this crate reads and writes. It changes only for a +/// wire-incompatible restructure, which is also the only time features fold +/// into the baseline or retire. +pub const EPOCH: u32 = 1; /// What [`decode_with`] accepts. #[derive(Debug, Clone, PartialEq, Eq)] 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 286966df..787ab2b6 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,10 +3,13 @@ // This file is @generated by prost-build. #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalPlan { - /// The writer's format version (4). + /// Frozen at 5 (the last sequential format version): readers that predate + /// features check it with `min_reader_version`. #[prost(uint32, tag = "1")] pub format_version: u32, - /// The oldest reader version that understands every field in use. + /// For readers that predate features: 2 for a plan without features, else + /// the highest legacy level of its features (3 limit_offset, 4 window, + /// 6 for a feature added after the format version was frozen). #[prost(uint32, tag = "2")] pub min_reader_version: u32, /// Plan nodes in post-order; a node's inputs are earlier entries. @@ -29,6 +32,26 @@ pub struct PhysicalPlan { /// Expressions in post-order; an expression's inputs are earlier entries. #[prost(message, repeated, tag = "10")] pub exprs: ::prost::alloc::vec::Vec, + /// The ids of the features the plan uses, sorted and without duplicates. + #[prost(string, repeated, tag = "11")] + pub features: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + /// The wire epoch, written explicitly (1). Absent marks a plan written + /// before the feature header existed: its features follow from its + /// content. + #[prost(uint32, tag = "12")] + pub epoch: u32, +} +/// The header fields of a `PhysicalPlan`, read without decoding its tables. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct PlanHeader { + #[prost(uint32, tag = "1")] + pub format_version: u32, + #[prost(uint32, tag = "2")] + pub min_reader_version: u32, + #[prost(string, repeated, tag = "11")] + pub features: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(uint32, tag = "12")] + pub epoch: u32, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ExprEnv { @@ -91,7 +114,7 @@ pub mod node { WriteFiles(super::WriteFilesNode), #[prost(message, tag = "14")] Exchange(super::ExchangeNode), - /// v4: a plan with a window needs a v4 reader. + /// feature: window #[prost(message, tag = "15")] Window(super::WindowNode), } @@ -175,7 +198,8 @@ pub struct LimitNode { pub input: u32, #[prost(uint64, tag = "2")] pub limit: u64, - /// v3: the number of leading rows to skip. v2 bytes leave this as zero. + /// feature: limit_offset (when nonzero) + /// The number of leading rows to skip. v2 bytes leave this as zero. #[prost(uint64, tag = "3")] pub offset: u64, } @@ -215,8 +239,9 @@ pub struct TopNNode { #[prost(uint64, tag = "3")] pub limit: u64, } -/// v4: window functions over an input sorted by the partitioning columns, -/// then the ordering. +/// feature: window +/// Window functions over an input sorted by the partitioning columns, then +/// the ordering. #[derive(Clone, PartialEq, ::prost::Message)] pub struct WindowNode { #[prost(uint32, tag = "1")] diff --git a/crates/krishiv-plan-codec/tests/features.rs b/crates/krishiv-plan-codec/tests/features.rs new file mode 100644 index 00000000..4f74770a --- /dev/null +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -0,0 +1,251 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The feature header (ADR 0040): what the encoder declares, how a reader +//! checks it before reading a node, and how corrupt headers are refused. +#![expect(clippy::unwrap_used, reason = "test code")] + +#[path = "common/inventory.rs"] +mod inventory; + +use krishiv_exec::physical::PhysNode; +use krishiv_plan_codec::features::{Construct, feature_of}; +use krishiv_plan_codec::proto as pb; +use krishiv_plan_codec::{ + CodecError, EPOCH, Feature, FeatureKind, Limits, PlanFeatures, ReaderFeatures, decode, + decode_with, encode, features_of, plan_features, to_message, +}; +use prost::Message as _; + +use inventory::{cases, plan_of}; + +fn message(name: &str) -> pb::PhysicalPlan { + let (_, plan) = cases().into_iter().find(|(n, _)| n == name).unwrap(); + to_message(&plan).unwrap() +} + +fn corrupt(m: &pb::PhysicalPlan, expected: &str) { + let err = decode(&m.encode_to_vec()).unwrap_err(); + assert_eq!(err, CodecError::Corrupt(expected.to_owned())); +} + +#[test] +fn the_header_lists_exactly_the_features_in_use() { + let m = message("window_ranking"); + assert_eq!(m.features, ["window"]); + assert_eq!(m.epoch, EPOCH); + assert_eq!(message("limit_offset").features, ["limit_offset"]); + assert!(message("limit").features.is_empty()); + // An offset of zero is the baseline `Limit`. + let zero = plan_of(PhysNode::Limit { + input: Box::new(PhysNode::OneRow), + offset: 0, + limit: 1, + }); + assert!(plan_features(&zero).is_empty()); + // Every registered feature occurs in some inventory case. + for f in Feature::ALL { + assert!( + cases().iter().any(|(_, p)| plan_features(p).contains(f)), + "{f}: no inventory case uses it" + ); + } +} + +#[test] +fn a_feature_in_a_subquery_plan_counts() { + let (_, window) = cases() + .into_iter() + .find(|(n, _)| n == "window_ranking") + .unwrap(); + let plan = plan_of(PhysNode::Filter { + input: Box::new(PhysNode::OneRow), + layout: Vec::new(), + predicate: inventory::predicate(), + subqueries: vec![krishiv_exec::physical::SubqueryPlan { + plan: window.root().clone(), + ..inventory::subquery() + }], + }); + assert_eq!( + plan_features(&plan).into_iter().collect::>(), + [Feature::Window] + ); + assert_eq!(decode(&encode(&plan).unwrap()).unwrap(), plan); +} + +#[test] +fn constructs_map_to_one_feature_or_the_baseline() { + let (_, window) = cases() + .into_iter() + .find(|(n, _)| n == "window_ranking") + .unwrap(); + assert_eq!( + feature_of(Construct::Node(window.root())), + Some(Feature::Window) + ); + assert_eq!(feature_of(Construct::Node(&PhysNode::OneRow)), None); + assert_eq!( + feature_of(Construct::Join(krishiv_exec::JoinType::LeftSemi)), + None + ); + assert_eq!( + feature_of(Construct::Partitioning( + &krishiv_exec::physical::Partitioning::Single + )), + None + ); + assert_eq!( + feature_of(Construct::Window( + &krishiv_exec::window::WindowFunction::RowNumber + )), + Some(Feature::Window) + ); + assert_eq!( + feature_of(Construct::Aggregate( + krishiv_exec::aggregate::AggregateFunction::Sum + )), + None + ); + assert_eq!(Feature::Window.kind(), FeatureKind::Semantic); + assert_eq!(Feature::LimitOffset.to_string(), "limit_offset"); + assert_eq!(Feature::from_id("no_such_feature"), None); +} + +#[test] +fn an_unknown_feature_is_named_before_any_node_is_read() { + let mut m = message("scan"); + m.features = vec!["zeta_future".to_owned()]; + // A node the reader would refuse: the feature check comes first. + m.nodes[0].kind = None; + let err = decode(&m.encode_to_vec()).unwrap_err(); + assert_eq!( + err, + CodecError::UnsupportedFeatures { + missing: vec!["zeta_future".to_owned()] + } + ); + assert_eq!( + err.to_string(), + "the plan uses features [zeta_future] that this reader does not implement" + ); +} + +#[test] +fn a_reader_without_a_feature_refuses_the_plan_naming_it() { + let bytes = encode( + &cases() + .into_iter() + .find(|(n, _)| n == "window_ranking") + .unwrap() + .1, + ) + .unwrap(); + let reader = ReaderFeatures::only([Feature::LimitOffset]); + assert!(!reader.contains(Feature::Window)); + let err = decode_with(&bytes, &Limits::default(), &reader).unwrap_err(); + assert_eq!( + err, + CodecError::UnsupportedFeatures { + missing: vec!["window".to_owned()] + } + ); + assert_eq!(ReaderFeatures::default(), ReaderFeatures::all()); +} + +#[test] +fn an_unknown_epoch_is_refused() { + let mut m = message("scan"); + m.epoch = EPOCH + 1; + let err = decode(&m.encode_to_vec()).unwrap_err(); + assert_eq!( + err, + CodecError::UnsupportedEpoch { + epoch: EPOCH + 1, + supported: EPOCH + } + ); + assert_eq!( + err.to_string(), + "the plan is of epoch 2, this reader implements epoch 1" + ); +} + +#[test] +fn malformed_headers_are_corrupt() { + for bad in ["", "Window", "1window", "win-dow", &"x".repeat(65)] { + let mut m = message("window_ranking"); + m.features = vec![bad.to_owned()]; + corrupt(&m, "a malformed feature id"); + } + let mut m = message("window_ranking"); + m.features = vec!["window".to_owned(), "window".to_owned()]; + corrupt(&m, "feature ids that are not sorted or repeat"); + let mut m = message("window_ranking"); + m.features = vec!["window".to_owned(), "limit_offset".to_owned()]; + corrupt(&m, "feature ids that are not sorted or repeat"); + // The legacy reader version must follow from the features. + let mut m = message("window_ranking"); + m.min_reader_version = 2; + corrupt(&m, "a reader version that does not match the features"); + let mut m = message("scan"); + m.format_version = 0; + corrupt(&m, "no format version"); +} + +#[test] +fn features_of_reads_the_header_only() { + let mut m = message("window_ranking"); + let header = features_of(&m.encode_to_vec()).unwrap(); + assert_eq!( + header, + PlanFeatures { + epoch: EPOCH, + features: vec!["window".to_owned()] + } + ); + assert_eq!( + header.missing(&ReaderFeatures::only([])), + ["window".to_owned()] + ); + assert!(header.missing(&ReaderFeatures::all()).is_empty()); + // The tables are not decoded: a broken node does not matter. + m.nodes[0].kind = None; + assert_eq!(features_of(&m.encode_to_vec()).unwrap(), header); + // An unknown id is reported, not refused. + m.features = vec!["zeta_future".to_owned()]; + let header = features_of(&m.encode_to_vec()).unwrap(); + assert_eq!(header.features, ["zeta_future"]); + assert_eq!( + header.missing(&ReaderFeatures::all()), + ["zeta_future".to_owned()] + ); +} + +#[test] +fn features_of_a_legacy_plan_follow_from_its_content() { + let mut m = message("limit_offset"); + m.epoch = 0; + m.features.clear(); + let header = features_of(&m.encode_to_vec()).unwrap(); + assert_eq!(header.features, ["limit_offset"]); + assert_eq!(header.epoch, EPOCH); +} + +#[test] +fn features_of_refuses_what_it_cannot_interpret() { + let mut m = message("scan"); + m.epoch = 9; + assert!(matches!( + features_of(&m.encode_to_vec()).unwrap_err(), + CodecError::UnsupportedEpoch { epoch: 9, .. } + )); + assert_eq!( + features_of(&[0xff, 0xff, 0xff, 0xff]).unwrap_err(), + CodecError::Corrupt("malformed protobuf".to_owned()) + ); + let mut m = message("scan"); + m.features = vec!["b".to_owned(), "a".to_owned()]; + assert_eq!( + features_of(&m.encode_to_vec()).unwrap_err(), + CodecError::Corrupt("feature ids that are not sorted or repeat".to_owned()) + ); +} diff --git a/crates/krishiv-plan-codec/tests/golden.rs b/crates/krishiv-plan-codec/tests/golden.rs index 44a87002..e266973e 100644 --- a/crates/krishiv-plan-codec/tests/golden.rs +++ b/crates/krishiv-plan-codec/tests/golden.rs @@ -1,83 +1,285 @@ // SPDX-License-Identifier: Apache-2.0 -//! Golden `krishiv.plan.v1` bytes (compatibility fixture, ADR 0026). +//! Golden `krishiv.plan.v1` bytes (compatibility fixture, ADR 0040). //! -//! `tests/golden/v/.bin` holds the bytes the codec wrote for each -//! plan of `common/inventory.rs` while format version N was current. The -//! plans are built by hand, not by the planner, so they change only when -//! the format does. Two rules keep plans decodable across releases: +//! The plans are the hand-built cases of `common/inventory.rs`, so they +//! change only when the format does. The directories under `tests/golden/`: //! -//! - every golden file of every version still decodes, with the current -//! decoder, to the plan its case builds (N reads N-1); -//! - the current version's files are byte-for-byte what the encoder writes -//! today, so a format change cannot slip in without a new version. +//! - `base/` holds the cases that use no feature, and `/` the cases +//! that use that feature (a case that uses several sits in one of their +//! directories). These files are byte for byte what the encoder writes +//! today, forever: a plan that does not use a feature encodes the same +//! before and after the feature exists, so a change to them is either an +//! undeclared feature or an epoch change. +//! - `v1` to `v4` are the frozen files of the sequential format versions. +//! They are never written again and must keep decoding to their plans. //! -//! When the format changes (a new field, a new node kind), keep the old -//! `v` directory untouched, raise `FORMAT_VERSION`, and add `v` -//! with `KRISHIV_BLESS_GOLDEN=1 cargo test -p krishiv-plan-codec --test -//! golden`. A case added to the inventory later gets files in the current -//! version only. -#![expect(clippy::unwrap_used, reason = "test code")] +//! A new feature adds its directory and edits no existing file: run +//! `KRISHIV_BLESS_GOLDEN=1 cargo test -p krishiv-plan-codec --test golden`, +//! which writes a file for every case that has none. +#![expect(clippy::unwrap_used, clippy::panic, reason = "test code")] #[path = "common/inventory.rs"] mod inventory; -use std::path::PathBuf; +use std::collections::BTreeSet; +use std::path::{Path, PathBuf}; -use krishiv_plan_codec::{FORMAT_VERSION, decode, encode}; +use krishiv_exec::physical::PhysicalPlan; +use krishiv_plan_codec::features::RETIRED; +use krishiv_plan_codec::proto as pb; +use krishiv_plan_codec::{ + CodecError, Feature, Limits, ReaderFeatures, decode, decode_with, encode, features_of, + plan_features, +}; +use prost::Message as _; + +/// The parsed message of golden bytes (flat tables: bounded nesting). +#[expect( + clippy::disallowed_methods, + reason = "test code over the codec's own goldens" +)] +fn message(bytes: &[u8]) -> pb::PhysicalPlan { + pb::PhysicalPlan::decode(bytes).unwrap() +} fn root() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/golden") } -fn version_dir(n: u32) -> PathBuf { - root().join(format!("v{n}")) +/// What a golden directory holds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Dir { + Base, + Feature(Feature), + Legacy(u32), } -#[test] -fn golden_files_of_every_version_decode_to_their_plans() { - if std::env::var_os("KRISHIV_BLESS_GOLDEN").is_some() { - std::fs::create_dir_all(version_dir(FORMAT_VERSION)).unwrap(); - for (name, plan) in inventory::cases() { - let path = version_dir(FORMAT_VERSION).join(format!("{name}.bin")); - std::fs::write(path, encode(&plan).unwrap()).unwrap(); - } +fn dir_kind(name: &str) -> Dir { + if name == "base" { + return Dir::Base; + } + if let Some(f) = Feature::from_id(name) { + return Dir::Feature(f); } - let cases = inventory::cases(); - let mut versions = Vec::new(); + let n = name.strip_prefix('v').and_then(|n| n.parse().ok()); + Dir::Legacy(n.unwrap_or_else(|| { + panic!("tests/golden/{name} is not `base`, a registered feature or a legacy `v`") + })) +} + +/// Every golden file: its directory, case name, path and bytes. +fn goldens() -> Vec<(Dir, String, PathBuf, Vec)> { + let mut out = Vec::new(); for entry in std::fs::read_dir(root()).unwrap() { let dir = entry.unwrap().path(); - let n: u32 = dir - .file_name() - .and_then(|n| n.to_str()) - .and_then(|n| n.strip_prefix('v')) - .and_then(|n| n.parse().ok()) - .unwrap_or_else(|| panic!("{} is not a v directory", dir.display())); - versions.push(n); - let mut files = 0; + let kind = dir_kind(dir.file_name().unwrap().to_str().unwrap()); for file in std::fs::read_dir(&dir).unwrap() { let path = file.unwrap().path(); let name = path.file_stem().unwrap().to_str().unwrap().to_owned(); let bytes = std::fs::read(&path).unwrap(); - let (_, plan) = cases - .iter() - .find(|(c, _)| *c == name) - .unwrap_or_else(|| panic!("{}: no such case in the inventory", path.display())); - let decoded = decode(&bytes).unwrap_or_else(|e| panic!("v{n} {name}: {e}")); - assert!(decoded == *plan, "v{n} {name}: decodes to another plan"); - if n == FORMAT_VERSION { - assert!( - encode(plan).unwrap() == bytes, - "{name}: the encoder no longer writes the v{n} bytes; a format change needs a new version" + out.push((kind, name, path, bytes)); + } + } + out.sort_by(|a, b| a.2.cmp(&b.2)); + out +} + +fn case(name: &str, path: &Path) -> PhysicalPlan { + inventory::cases() + .into_iter() + .find(|(c, _)| c == name) + .unwrap_or_else(|| panic!("{}: no such case in the inventory", path.display())) + .1 +} + +/// Writes a file for every inventory case that has none outside the legacy +/// directories: in `base/`, or in the directory of its last feature. +fn bless() { + let have: BTreeSet = goldens() + .into_iter() + .filter(|(d, ..)| !matches!(d, Dir::Legacy(_))) + .map(|(_, name, ..)| name) + .collect(); + for (name, plan) in inventory::cases() { + if have.contains(&name) { + continue; + } + let dir = plan_features(&plan) + .last() + .map_or("base", |f| f.id()) + .to_owned(); + std::fs::create_dir_all(root().join(&dir)).unwrap(); + let path = root().join(dir).join(format!("{name}.bin")); + std::fs::write(path, encode(&plan).unwrap()).unwrap(); + } +} + +#[test] +fn every_golden_decodes_to_its_plan_and_current_ones_are_exact() { + if std::env::var_os("KRISHIV_BLESS_GOLDEN").is_some() { + bless(); + } + let mut current = BTreeSet::new(); + let mut legacy = BTreeSet::new(); + for (dir, name, path, bytes) in goldens() { + let plan = case(&name, &path); + let decoded = decode(&bytes).unwrap_or_else(|e| panic!("{}: {e}", path.display())); + assert!( + decoded == plan, + "{}: decodes to another plan", + path.display() + ); + let features = plan_features(&plan); + match dir { + Dir::Legacy(n) => { + legacy.insert(n); + continue; + } + Dir::Base => assert!(features.is_empty(), "{}: uses features", path.display()), + Dir::Feature(f) => assert!( + features.contains(&f), + "{}: does not use {f}", + path.display() + ), + } + assert!( + encode(&plan).unwrap() == bytes, + "{}: the encoder no longer writes these bytes; a format change needs a feature", + path.display() + ); + let header = features_of(&bytes).unwrap().features; + let expected: Vec = features.iter().map(|f| f.id().to_owned()).collect(); + assert_eq!(header, expected, "{}", path.display()); + assert!( + current.insert(name.clone()), + "{name}: two current golden files" + ); + } + for (name, _) in inventory::cases() { + assert!(current.contains(&name), "{name}: no golden file (bless it)"); + } + assert_eq!( + legacy, + BTreeSet::from([1, 2, 3, 4]), + "the legacy directories" + ); +} + +#[test] +fn the_registry_is_sorted_and_every_feature_has_goldens() { + let ids: Vec<&str> = Feature::ALL.iter().map(|f| f.id()).collect(); + assert!(ids.windows(2).all(|w| w[0] < w[1]), "not sorted: {ids:?}"); + for f in Feature::ALL { + let id = f.id(); + assert!( + id.starts_with(|c: char| c.is_ascii_lowercase()) + && id + .chars() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_'), + "{id} is not snake_case" + ); + assert!(!RETIRED.contains(&id), "{id} is retired"); + assert_eq!(Feature::from_id(id), Some(*f)); + assert!( + (3..=krishiv_plan_codec::POST_MIGRATION_READER_VERSION).contains(&f.legacy_level()), + "{id}: legacy level {}", + f.legacy_level() + ); + assert!(!f.meaning().is_empty(), "{id}: no meaning"); + let dir = root().join(id); + let files = std::fs::read_dir(&dir).map_or(0, Iterator::count); + assert!(files > 0, "{id}: no golden files in {}", dir.display()); + } + // `goldens` panics on a directory that is none of the three kinds. + assert!(!goldens().is_empty()); +} + +/// The reader feature sets to simulate: every subset of the registry while +/// it is small, then the empty and full sets, each feature alone and each +/// feature left out. +fn reader_sets() -> Vec> { + let all = Feature::ALL; + if all.len() <= 8 { + return (0..1_u32 << all.len()) + .map(|mask| { + (0..all.len()) + .filter(|i| mask & (1 << i) != 0) + .map(|i| all[i]) + .collect() + }) + .collect(); + } + let full: BTreeSet = all.iter().copied().collect(); + let mut sets = vec![BTreeSet::new(), full.clone()]; + for f in all { + sets.push(BTreeSet::from([*f])); + let mut without = full.clone(); + without.remove(f); + sets.push(without); + } + sets +} + +#[test] +fn an_older_reader_reads_exactly_the_plans_whose_features_it_has() { + for (_, name, path, bytes) in goldens() { + let plan = case(&name, &path); + let needs = plan_features(&plan); + for set in reader_sets() { + let reader = ReaderFeatures::only(set.iter().copied()); + let result = decode_with(&bytes, &Limits::default(), &reader); + let missing: Vec = needs.difference(&set).map(|f| f.id().to_owned()).collect(); + if missing.is_empty() { + assert!(result.unwrap() == plan, "{}", path.display()); + } else { + assert_eq!( + result.unwrap_err(), + CodecError::UnsupportedFeatures { missing }, + "{} with {set:?}", + path.display() ); } - files += 1; } - if n == FORMAT_VERSION { - assert_eq!(files, cases.len(), "a case without a golden file in v{n}"); + } +} + +#[test] +fn a_header_that_hides_or_invents_a_feature_is_corrupt() { + let mut checked = 0; + for (dir, _, path, bytes) in goldens() { + if matches!(dir, Dir::Legacy(_)) { + continue; + } + let m = message(&bytes); + let mut edited = Vec::new(); + // Each declared feature stripped. + for i in 0..m.features.len() { + let mut e = m.clone(); + e.features.remove(i); + edited.push(e); + } + // Each unused feature added at its sorted place. + for f in Feature::ALL { + if !m.features.iter().any(|id| id == f.id()) { + let mut e = m.clone(); + e.features.push(f.id().to_owned()); + e.features.sort(); + edited.push(e); + } + } + for e in edited { + let err = decode(&e.encode_to_vec()).unwrap_err(); + assert_eq!( + err, + CodecError::Corrupt( + "the feature header does not match the plan's constructs".to_owned() + ), + "{} with {:?}", + path.display(), + e.features + ); + checked += 1; } } - assert!( - versions.contains(&FORMAT_VERSION), - "no golden files for the current format version" - ); + assert!(checked > 0); } diff --git a/crates/krishiv-plan-codec/tests/golden/base/aggregate.bin b/crates/krishiv-plan-codec/tests/golden/base/aggregate.bin new file mode 100644 index 0000000000000000000000000000000000000000..bc656f46c9764f2d9da4ceb30a900bb289ed7c87 GIT binary patch literal 1732 zcmcgt&2G~`5Z+lk@n)M+cT>Pv9ITwGNL+G4JwVkHs<>c!EF!DgO{}_hWpC7QG4VU|8T)yGLQP1Twxmup+ z;^^gpHag9R(uv|jJx+}Tqg-hp=TWSLY?h`-4vgx`hRKiOOx;s=rF~N5ky<;|MLf#K zW&EIW`6~hxXrGHl0Denw1tl(8!CudjpnFn6DH}j4^?O zG5ca>3~;r9>f-FOTWk$|w9CcWLk&=I7r&3QU~TlSRIve>>+tFw1{XVRaqg^Fc6R>Q z|K{v|brfgZe<2`%-rh3!6oW88-$Pzz2p1|2@M#Oc0GPiAHC5k&i>q zBy^?%bE2wnRPs|z>f9#&NLV%6{;yB(&cf=|C}2?~8EYa+69)-l%_9wjT_gQVoxGpd ziE?TDQ7+e+ICFBM&V{~yHlO)iGSPHzmJ`mBocP;2Mhq;Ff!wkKeglMW9>>Knj=C*6 z@Jcp0-n4go5gn;@d_g@!-A1ML{&YH>DW43`R?BSpYC4krlKk)I<9>T#B8;t@7`xWQ z>PyEKlC!*uTfX^ed`G_}?WggfQIOV;`cT27FJ&W4h9lzx^gGfH^V}%>5g47M*t;2j zXzeH+jFQX*r@9Q1VU`zhwCf3Qr=KnF_1o$9Pai)tyN=Fu`l>YH>fB!-^n0FcC1sYI sY@DVZbb)Pg!MEWFgtp?F8)2unF5wY+5ber!yVsGH<2faA_XE50r<5HUIzs literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/aggregate_final.bin b/crates/krishiv-plan-codec/tests/golden/base/aggregate_final.bin new file mode 100644 index 0000000000000000000000000000000000000000..d29681eaeed67235f1138fb505b183c0f447e838 GIT binary patch literal 850 zcmbVKv2N5r5Zzh(&U=#!arS~R7KL?HWGRg+lmenuD5Aln1d&zjyR%~K$lg<+qd^KC z1s}i<@Cg)1d`t@d!OYkpQbelEYDPOVJMVeMqei|AFC~7K%G0m$6>hfYra60eYOE>i zC3;!DFzeDHSnZ?xzRq&3aNOFaIJJ6?BU_*4m42)rp?lucncliEO}?tvZNAs1!VLu) zOeobTf^edwhL)U3bR0VfCPsu7M@3cTO`=eO^xXjV3*~4|hbL1ZgrXyadu@dP@^(ml zvc7mAwkXH1e6j&)L`tvZ6S7K7DZA0vMjRf>dkjm?;Z$9D+kISIJr5U0S@>2#1o82G zn7M`sfWL&r?ax5TYzb6zfL+G_D(@oTMt+lCk{*z<|EB9YZsi&rBf1YVrnbg(SmY2Q z%j03PZ@23_#$e%*C_?6;auwrpO-4G5RQxy6fm;N3HP+tL-Tzk@!bm3QzRVX^LpfZ{ zLc@myag-Lzl??%s8MR~)X$5a`I9rzbMfxky@W5@si literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/exchange_broadcast.bin b/crates/krishiv-plan-codec/tests/golden/base/exchange_broadcast.bin new file mode 100644 index 0000000000000000000000000000000000000000..aef2747401badd667b1f012150fab951e61b7cdc GIT binary patch literal 756 zcmbtSy>8S%5Z+n;?M?28vloQ1D6FfHC7P>H3W!o6M1!G3BCFVUXT{i&y{ABzCQ|4q zcmQ62N1#CBVN&o8X2uRk5p?Y8`|SDsJnL0nDB~-Mmr@1#6~4q-cV^ntw92&=UFAr$W@&|qSz zyaeKjk{UWPCdup6K`<#Ibm*0JU9?Mu5@hH`a8N2ob2vVp3LzA2A>11l0?67Ygv6n^ zBX%jqSs`&mNC=cA9}p`srRdI3o3MKDra&Ee$T-QxoO*I<*{xfX#c70- wva1_gZ>uVT1&9NwPYel}0{5S~zk)c%C-rgL;mowM74-|l1h1sf1x4FCWD literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/exchange_hash.bin b/crates/krishiv-plan-codec/tests/golden/base/exchange_hash.bin new file mode 100644 index 0000000000000000000000000000000000000000..bd37e63ade735c3e21e12551b551df9633340d07 GIT binary patch literal 760 zcmbtSy>8S%5Z+n;?af_?vloQ1D6FfHC7P>H3W!o6M1!G3BCFVUXT{i&y{ABzCQ|4q zcmQ62N1#CBVN&o8X2uRk5p?Y6`|SDO*c}yW``T5JHg%;oh(iK-xYbBo4(L zu}eOV3W*~^LZB%5fLMtcd8b2e!sao3s9{MtnyKqxZ?KE&r}3;sHI03AQ3^;P{Uraj z0dDt$a78#Iu>Pj+dv3__)B}^@%j8pYpM;kjn2Cqr4Sc?x7jO+;ym56iG-q*L&dCYm z;Nq{r$KE31gW%;3GUj!p|L=hrs&>5_6?H*JwWj}(OccPF~|*#G?Y^^3m8S%5Z+na=e@~=bM}HT7KL>cvP5$gNyR%~K$lg<+OA{${ zDR=;0fJdM};$c$o4razqkRs^V)%V#mv-7*Fyg)`*5-+9l^%1_o?e@$xr|(aUHD$d* zKg*Y9Q(6S819U&sS*{h1TH6#SR?pG1^=V$|$NCYvXHA{y>A7k0b-iixXT2%fQlP;^ zQh6~%VHCw!M%@RNjRLS8-J(Qi<|vOH-VZ)6DmuFfUbf7 zVo0v#zslgGKM7Za0|LwMx~}82e4a94+<%!yRKkchFWEB_x58z7x$oz53|?Hmx^11k zI4=8SgrR-$SNlV55%M^Ax&0peT9W^FY=&Yi#^}DzmsUeL{F6ZxlmGw# literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/exchange_round_robin.bin b/crates/krishiv-plan-codec/tests/golden/base/exchange_round_robin.bin new file mode 100644 index 0000000000000000000000000000000000000000..24fa2fe1505a61826e63185e906864a705d73748 GIT binary patch literal 756 zcmbtSy>8S%5Z+n;?M*I;vloQ1D6FfHr8HNe6cD9Chz3K6L{_oy&Wf=kdryIm1}W)M z@Bq93k3fOM!=&IH%#0n7BIwxF_u2FNdDg4EP{x-MKS>qnSNIaQyK~c?y*)M7RLu&5 zyjYq|Wf80n(S6tCg;wZywk=Pso}*`*v!d1y^nG;C+a}k$7p5)N&8903hETkvK!b^; z@)C$AN^0oHm?W=L2f?I>(4kk>bY|K2q0^p5E6%C zN96;98c9vus8U{&C^)5zS@c9TqKZQ%YTWp z)DOZX;fTQYo4)T^jPt2KCc~4-Ct;s9FF7z155VG|@74v}gBO3ax*eFSxG&e_gmGx` z*U+Q3i1-kAxx*goIt)x;WO9 z75!@=A3Mj|BYUvU{C1zfFDuw!761SM literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/exchange_single.bin b/crates/krishiv-plan-codec/tests/golden/base/exchange_single.bin new file mode 100644 index 0000000000000000000000000000000000000000..0821723a9a742efa1836b86977d1d8881df5f681 GIT binary patch literal 756 zcmbtSy>8S%5Z+n;?M*Jk*$cv06xLP963ta81&LB2M1!G3BCFVUXT{i&y{AA&gA}?H zJOD4iBTz)j!=&IH%#0n7BIwxF_u2FNdDg4EP{tP$KT8$p*Z2yzyED_CzB@72RLu&5 zyjYq|Wf80n(S6_Kg;wZywk=Pro}*`*)1uZ7^#gRz+9uaK=cX;z&891!454^Sfd&&x z!MvKlpsSlg8foCn#1w&R0yGH3*p|f5J1)*AtVmP zJ+VVE&I*YmLPDS{`G{DFDMh!2+JxO>`Ow3XaWqvo!S3J}HwUq5eKn1x773)6@?YXC zbxpV+91_@m+xI<-aX$6OWOy?9B<$1XB?o5W0a*O&-MWB#@ZyhFw*zw(_vMO@)wbUu)lJmPha|lPSIf#9!PoUN%d25uQF zRVkOQkc!$H2M)c#g#!mJzzGSdi2fUh3x^&!L_*?N^Ts>7N?LA=H7}oeZ{8c{`#j5R zpuv60?@{K85BU4MRh&fQo1Yv;Wt8T}+#SY8(IhQ7DFtxrRz8dc<4#eIlfzPUxn1Tr z<4jx>Yuq{+=R?uH9gX8rJ}KhY>QK;NM35+;%nnK5F)B!b3orbzXL1ttI41?SlPrtJ z1I8(#VEJS}VWv>w;Gk(32E)fNENu;g;A#aYpao;eXk$$60<=EP5GU@U8)%u)#Au_= z4V8PAUe_pbajnUoxSiT>Jb62KK)YX1fv*uT;c%Mv!Pc||j-~;)nzn&y+5u0~E^tkI z(9m=lnwqXaOVjJH6nMhrX1djf_8hkW<~V@P9JfKtaR-*?xC<+D+=JCQUWN;^6XHJj znzrDgrU3++wjtEC121XXg|4POcsV4!FX?~Duf%?g^FGe2IMw>&s;ZRSYCkH>!}_l> z4odIfqXw<5GOC42{;$sGxvGX5s%=cSTxF@cDrbmu)^_8$w$IKseDxWqVb#wpufX;H z{ahlfG4QzcUVKyvlFl$X66DhX=eA6aM`b|BV3}Jo&r5+R!U&$W2(9*(xfx|c<>_E!i0T_FbCh;px{8>oOmiM&+uL5B~o1hAKCjPETu9)VrM! z$p!n&Ve3GADnTrLP8w|b@P|Ksw$WofD;OUzwk9ZZd`w(0J;Ua+DbH)Xepaz_8n3^w z2%XjV(#FF53X1FEKWu{Sd%X+lq8F1}(D7-PtZkD1i@4K~?QV;C!p;=N#7ydC1pU|>HF?Z_IICTi)xemiM$a)9jQaAh?`fT}L)B zL2LK-XkD#JXFN$l)kx$1C<=GR#t!_3V2T-^i>hXRU2w%x9O|u#WTvwyvdmfe{-T&2 zx1RX|00>L~XU^0Cd?FqZCkUF)^E}_eO7{Vn)I}w#p{|BW!3!p02#a|6t8-}wg}_Z> z?nx@zOPUbTgf#@fa^WnB-ov@{29M4toFYqIigT%3`lr)$85k#T@^zqr6+j~}$hIJZ zwOyCaS=6(N4a?T_6uWOrS?$LCf%O?{O1vuEI25Ve zIPd_x0FQtK5)ab@@34%MvZ8>*0c_1AGvCZNp0R}&a&^w}4HtoYjE``$J<-kavm-4}J1OM!lCn&QaFF?yyx&MSFS-az-HsWZ8Cs+)XKFWY>#gVcrq z30iUC#h?a)OK8a$#7`0jK_`gNqE}Q^-b@5?;GrAAULhQ5VRblUj0xI|xtE$TAZv#t zBptGAY>O}(6_SofViJntw@3>%By^+8jjg*+?>Gz@2SafYYT0@<*>^ROI!}xW?5dFlachpX1@))ugxMme7oFk_jxVZ|9VN9 zI3@#hALcV7q3nP4Ov1|vqL&u)g;4;L5xQwz8%Yn==%S=I#`qoQKBYnS+<@?*&W zQrQqld69^FgVD7b11l#WLZ3DFiwFwrVGu48A7+ctyw|{ z!$pXB90>u;YUmC|hnynw*wd#Rj4ry5j;H8pjEu};ccz^6b$K#gVkk`7v+U4bDC^{~ z<>h>rF>zU3;x_Vde;6yV!k3Hl&3gUbaC;|Kby*KKezcav;_VONYlt4g%G;nWLxdGG=w94COncgS8Qk zpL<+;f9>m|C#{RGUuk>W>VQv!=k!l(X>LBDAE0ld>-^4mJl5>!^GFEy%Uy>GGF8E) zjh0ad&A4Xk<;gi)V`$SXGep)p=_!q?XCm~8O5`_{CjTa`R!y5(uE|`1`Twdb5u;J% z=8d?k1f-Ks*%k1h!MPPCy}t4Q8mM5H=Stv9S29WQaT0pLXbb80lT7*hvhY#bd=N+1 zT?Wqhqv4L8jgLOOdsA7pWhT>yg$n0RXdZfj%d1I|7aF%N3j%GEA%oMYuUr9&gXF^59I+`i)9|ky>qpUXZfm#k4ngE z2#}yWE~-9wb-^VRGVdj=Y)+O~x=@!3DR^@F?D(7x>KXg7*vn#Ix4Lh0GRD3^psEwNr(?Vmr4 zTpfGeP&(jNdz?YQfemPX3_cG&^1m@RJin7(ke-kl|K(=0F-*;UB8Pst8&f0JHCje^ zX^k=6Gt58T%(G1mBc@|w*KLz!>YAL0blY+EN5{8sGuKSOh*hp_UP1oBYnS+<@?*&W zQrQqld69^FgVD7b11l#WLZ3DFiwFwrVGu48A7+ctyw|{ z!$pXB90>u;YUmC|hnynw*wd#Rj4ry5j;H8pjEu};ccz^6b$K#gVkk`7v+U4bDC^{~ z<>h>rF>zU3;x_Vde;6yV!k3Hl&3gUbaC;|Kby*KKezcav;_VONYlt4g%G;nWLxdGG=w94COncgS8Qk zpL<+;f9>m|C#{RGUuk>W>QJ8s&*`7o(%gJPKS19^*ZH0Cc&ypc=aCTZm%9!XWU7Ko z8!e*_nsLq6%ae1q#?Yo&W{9kH(o-5&&qU}GmB?=@P5w<>t(rEoT$8y1^Z!*>B1WUi z%^Pu72}mcOvMb<0gL5lPdVS>qG*H1X&y~QJu4Izp<0SNg(H7G0Czi6B1Sx`!S-oe^%)XhOwUrmh=tAPTRK9+PZ*kI|n&#x=v9YGC zm*{8t!mLY+V0D1*r#j2E!a-}B;@IjLdbU2vEB#bILHD$&Gre_YntWBS+kCf6MI!|o zOeB>TLo`%ULrcyuexEo9CP9Q2y`rk}X0A|z4BQa*3gu`{M@M5JgrX&cd&fcmc{_xF zI1rD-7Ug(XKpYZc0_~E|h?N*qcGPKO4)^k|h9&24tZx17E-r3gMw4U|Z0-$WNWRN| zoq+rOBwP><2<(3#8F%FU)CA-1$uywu16pjdVL$-#hYEttjT3(Tw)u5^bvc;%rZVN1)~1sOzfKw@ yAr6bSs%^C{%MfND_N0;r@Cs5_%jY}kG}%Vjrwn!blX{#?(J@uVVF&KRNAMffL1>Hs literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/nested_loop_join.bin b/crates/krishiv-plan-codec/tests/golden/base/nested_loop_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..f9883d5d63defa06459b85aa33b8c15a94144327 GIT binary patch literal 868 zcmbtSy>8S%5Z+na=k?^m`SyY^7KL>cvP5$gN&!(SglI5S2#Ks>-<=g>NA{iqU7ARt zqu>E}0Um(@iHAwSI~+5<=tKaCf?a)|J%8Wqj;*{<#+MQ=r1JF(e2%-_scBE%92;w@ zW`TZQ%+0E@2v&#azHRbCD;#vTEsw39qGy|vqSkly9du9ICfB1g(-zBS)fEr=P`sl+ zgNdc`5{LsOHFRVI$;;G1FexH*=#_O{v@?YgWavh)Un)m)I6fK*Arx&P+-nvB$l4=> z#G$w)Mik?$kT@bF1j>^4h?N*pbf>RP*u5$5H7ps2Lv`hkd%w7P97`?Y@vBh?kuwd@ z0qJNNU?S6x@^2^MzTXL#gaZQGuh;7}cjkO*hDrZq8qg^Qw0X&%nfNw*l2145eAeK_ zH&EBN=PK6anw+qifcRqqo8BVgiSTm!y{YTS{@=+N(wGJ4J}c%{L$!G|a}Doih+bAM zmNo{MP0-Dn#%lVDtSKw{@IpRx4z)>^Wo?r)(eRYI%Pzo1|leM--uM+tBTFf{v+k db~~_j|5M!Crv5JO`PVSd^5hC=n&1y7~$%C3hDoN+e;cjFT^V%@2hbfc%Ko}Q0OG3Vc ziXo%22x1_<|M#VQgQ4k?-et0j@3N;NW;i7RCox&aIj4QfM$T^x6>IVhv`@a{J)h&x z5R$VwSK*L4;D`i{jpC`72S(l2>tIDYjec##V`X@^2W tUe2zqciW~5?!gbP*>wGVUESnbWaIZNe|{Z*zdk>b^VD~~eHy}@_yZ_7EJ^?X literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/project.bin b/crates/krishiv-plan-codec/tests/golden/base/project.bin new file mode 100644 index 0000000000000000000000000000000000000000..cd09e7b73f13a2d2d85da5bf1f46821269090f9d GIT binary patch literal 955 zcmb_ay>8S%5Z+na=e@avID0{u6@_&bGNQQ(MFLK#5TfCrL?Wx$xifL>$lg<+qW~#% z6d+OZ0z3i|NIXmm-r<<>MJEO5U{^EV`T1tXGoyk?nlp)Sr3&>Ee2kmziEWOaA6RG0 zdWK<^AKFFf5S)(CKd!S}E9|wdDGr<-Vc_bcywcb8HS|xKI@4RHw#n!9qRnr2CNof= z!5XQ81WaE^4J`@%!N!zi*@#rkDB(A~2s| z^bs_bWPFqg%mDj3Z1Tz;=b51=dhOyk3P%BYARv9AzP-6iu>V@}WK|je@w((hCSLkp z{;f0IhaW^|L=z&GpDmY5?gaZ&5KOu+vmSM&N0*Q6I5W=0o%wJX7jg_f+);Hg4MuTX z#$-gRiHJWYvdS%Do)90uyDMBv^1n{fQjBpQ{YUws(@?IS_E5viF=8++W^-o%#zXX{ zb?r25wX;P@izwtBZ%^A~UQ{kQwQW+&s=CRuhcT_d^1Bzmez*Mj^~>kpD72L=U$t(! zd1Y*c?KMNAASL!7jF#! literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/project_without_nullable.bin b/crates/krishiv-plan-codec/tests/golden/base/project_without_nullable.bin new file mode 100644 index 0000000000000000000000000000000000000000..05188a2ce11e26073011accd7417f431f29b0467 GIT binary patch literal 394 zcmZWiF;2uV5FGn*i7^Nw-xX0rVO>d?iV{&Gs;eoml-ymUD3ORIpreG2f(P&co=zP&%_HN?>8ZYQ|kO55ZcDh zl9Df_V#ug0g6N5Fe|;I>U?{qz_L=PByX>inIhI7=Bqpmk7qm~=DEN($#F~5~?UV2L zpU?7V2+7%;o8XW-V2K2djpC?R2RgZKHr|MC8r{Z-hf3@Dy4~)4DMYo<$$A*P$?C?t xwi`XPk{Vi4M$xm_K`=2QwCEL8l{bk(3DS21I4G2(*&QB@g%FCi5bh-k0i^Ab z_=LW=CU(fjr+mVIBqE_x@)n^IWAbiw+=$IRy^CQ&A~-t?vV{{9{r{_ACCeNL`rbD2m+PW(y!t!A$9C&@X)>-CyzC8>GqXjO~W|AQPEbltyX0jz!b!RRB{FnA$7I9 gy_ZhnU4;AOp==x1^>~7gsWLV@urU+kj(Y^Z0Y*1wX8-^I literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/scan.bin b/crates/krishiv-plan-codec/tests/golden/base/scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..3a1a78f470cf3efb7bd7d6945a06ffd17623ab66 GIT binary patch literal 744 zcmbtSy>8S%5Z+na=e@~=ID0`Di^94JS)#cLrGO|ELNpjkB(jQqcUH_gviB6|Xpllj z!2|FDJOTw050ip-Ff(>QilAdg-)GPNj=d^~WO6C-gH)k@fzNTbJ2&mw>r>-Q-K;Sz z%9YvH4#DXN{WncfYK4=|wbiN93k+OyR$6^m-$DPpZ3?}2VcK%jY`gNo5K49wXfTOX zK?=!KNevwt)AVKLA(#viIt(gn%XX!4* zVvk~+6%ogTlt5YX9mc&4tx{lPD;9%rBAzXo#qAB0Q75rOSj zec$r{&ZnN34o{{D4V#d%$bp&o&^++-&AO0#u<$|E^^#o0eYvI}jEjoD79D4cm=|a9 z2OaA=vj6u38S1f|qW`R1ISuvrYE~NFE)j#gT5nteu$-fxH;vPDS)8eA`oBUx^p3Sj zH8S%5Z+na=e@~=ID0`Di^94JS)#cLT@j*G2+?3Dk;p3c-B~eqWbY}^rHPbu z6nTKW0FOX{#KWZE9n6d!kRs^V(f8Riv-7vR$_r$4De*!oU%$mSxZR$b=H$b%v8Jq7 z=x6!TY)Xq@b%5^2I?J`fQEQvx*yn)od|hwa{8{gcwiIYE zkyKs`(O5|hEh*#pUE&~^1QA;FimJ++xk3pta6?!W%F*nOjwV6~MIwZI$3g&UdxU^E z5D&x-`FK}A91>yz?UGN3m6(ustB;M@JcRc-EGdT*b?fi;c5(Y6nk-Pu5VWfxfEbc5 z@?Q;bxgUf}!U2Kx_g&X3;U#-!;vu+!>-%v&=itSStGl5&it}Sce+@p=79k%5FSnN=t|k3{56n=G`54{T`O<1AhhMYQ@M(_drNwG(BY^o7-L$T) zrpIGVQPO=2`P@0wCSDhnjn7OQ7ptmn^6XVe=iGgF!t<-{=eMsnqnU3iQ+{r3x_$5E zq)`&$xM-`|R-3X6VFqF$l{|o#kh)qv-Akv*F2X)}sM;Ub!(@t%sWLV@@Y{R@zW{!$ BU+(|_ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/top_n.bin b/crates/krishiv-plan-codec/tests/golden/base/top_n.bin new file mode 100644 index 0000000000000000000000000000000000000000..2c91e369170692a98d67b170295583bcf61d0212 GIT binary patch literal 762 zcmbtSy>8S%5Z+na=e@}#arS~R7KL>cvXtg3l!8Q8MTiDNi9}Yh@6L*`BYRJQjs_`o z6nOw%fJdM}%EP4K9n6d!kRs^V(f8Riv-7vR$_r$4CGnF~zJ7%-al5@R&H39iV@+AF z(9iOv*_0N+>Hyt$b(U*|qt-UXnbkA&Y<-?r`jLK!?nP5)diT;a`MTb;`9bfBwiIYE zkyKs`(O5|hEh*#pb>bkH1QA;FimJ++xk3pta6>pOl%v@lolb-hibM$aj)efy_6Y%T zAa=wq`FK}A91>yz?UIj)m6(ustB;M@JcRc-EGdhLy7l*ZySP1yb`}#YL(r~*078r* z`6B<-0vEg?ToH~5tiS2Hj!W`>s)BL#|9fPHa?Ho*zQ~tWLpl7KrH1!&L@zB?Ya0Q~r|712 zZ8bd~Yl@O?T*#-+LYsJ9R5rdeZCtFXy2-QWA)R#h!wElr?tXs%_H{J#O=ZdttxdP@ zy`3~lLL3)uRoiM)mLbeQ97-jR;RI4w%P0HkG}%LVMjop6$MqnYqGPIz%?|uFpTaLA CbzmI; literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/union.bin b/crates/krishiv-plan-codec/tests/golden/base/union.bin new file mode 100644 index 0000000000000000000000000000000000000000..b0514d8d1cabdbd0ed131f1dda9f8e76f4aa5381 GIT binary patch literal 872 zcmbtS!EVz)5Zztd@n(`xXVa>##lgx6S!yGN3krf(IU&Ra!x4#YV<)lV+QD9{aLWa$ z+&J(7`~aW80f{60M-TmnFzcj96>wr!^VT!7GwV0A+;W(A$>2TXws-;0p;Inpv3Na^ zN~ZY?>?oefDpf$0aKJpzqgZfgl&VN3O7y`}`6AB5f!GIgS>%z}xsXLX&#N*%tV3Rl zQ$eN2xaCu?$(f)f8cqLYU;vc?04>2vvMer!92jMe>C#cc4WYx{X`2wjaS>u}kr0a3 zE~_a?orq0kt&`IX(R{p9M{V}f`&oO3Q7 ziHBfgFWywK5(ltf-X9GHXVqj8tAQK@M>nB8(c4mQmB}oNqr;_C^MBEJ^?Bq3e!XjGAlU{kEv!Zjx@ns5 zr`vh9uAxOYif@~&vvggZ6QOPzP5x-~=5FHZ#%nR_waF{d{@(#ftTAkY`7E9)LDS8r zoC^AO2w;WDY_2>?hd6*R&y~Q(s$`Pl`y%XtF&5IFCzi)jV$y67n(_#l7$bPw T@6cY*14Cw!b{q7@eM)};xH)|l literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/base/write_files.bin b/crates/krishiv-plan-codec/tests/golden/base/write_files.bin new file mode 100644 index 0000000000000000000000000000000000000000..fa1fb659fdc12126d032bbca142c2abe7bcbae7e GIT binary patch literal 750 zcmbtSy>8S%5Z+na=e@~=ID0`Di^94JSxR#iN)e(|2+_o$L?Wx$cV}Yk$lg;(mnKr^ zD0l!~fJdM};$c$o4rayig{Z{dVkCULd0ziPuv3`WTOJvpqM>+1pcNO<6C| z&+>&?mlnb50NwX>mTQHh);7hd)id;LeU?}Hv3`W^c~fV4>%uhos$RGGZXb#^6lgG! zR9+0xSV;{n8RPhM;vkp=5nA+$s>++WLJ2Z(L)a^nqd6R%OoR}Mwh-8HW>f>u>jdar-=y??epAC;2Z~7WqNA zAsi6ce$#aw%Wyt*#<+hnjYx7to0sgFiI2{*pYPZC+=CY%SKW=uRos_rGQu#f_-ouj zTZBA5FSp-AT}$@=ekMab=3{hUr!-qMdmln&FjR59Tbkn-Fny!pBMM-~I z$Y;)>Hu0*cYkO1wr@S39mkOKfitbGMf3OGUa7!)6IMLC5@60$3E_yvCkU#-qNEDCEYWGT&7CwD;4v`waWE=`-Sn@yMR520vFfd&&v z<;4(9mDJFYF^yj(4uVM#p+m2%>%3hmlpq5)go9E!n#0k_ObDT93gO?l8x?j_m(EDnmV%Q*@u_h1F1vzoyXec8Tbv{+NF=)XJR`;Yz4Z(qMm7QU%X^{%t&_N_OQMoEa% zva1_gZ>lPU1&9Nw95gt>7)c%C-CUbO5ow3`2-|pA5syr|4aZ?ms HZcpGBeF0?_ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/window/window_aggregates.bin b/crates/krishiv-plan-codec/tests/golden/window/window_aggregates.bin new file mode 100644 index 0000000000000000000000000000000000000000..bb298e4f027b1c2543276fc74e43660a66505bd9 GIT binary patch literal 1800 zcmcgt&2G~`5Z+lk&SslZcU!<%94wzuBrZ82AE4@~LR>ICp`vzGJBd}-uI!Cc5C>GK zRBjx20A7Gc;DE%#^uRlWSv!qV8h%nFb~P*Sj>q%Ow=-kbuub=j@>9y1;u(L+m&+qr z9PI8(CDXjmn^D}8!&Grl!sg~-9>s!jt5ij@uY}JVDnE!baYx+d=24MHVr3+Yc#sdv z_%IXQpi zwgWDx1HxorgL}$eoiGM|8N6WNL1+6qXtWJF1&vM*vzj#wji3wJrls9EM#x8lJdJmY zd{p{xAxIf@fuQ4sY0#lh95?l|9-0b<&NZ17Me8+{HHc#bK`RhT6A&`%$NuYJz>FHi z&hMkYHTN!R=V0i2A;d!G7EV1hgr5Ou7HG(dVODc64EC5=d;b^ionLc*5jG85mDg#| zjlT-|Y68wM&+4Fcn6rUj4Vvnp_lL#w;xlGp`d#np7p;IU)EsCKtqMfN!HG`4=6}Ud z)xee!vR;y^xQHYF9Rvl|05wdHE?xu+GhrRZML&)vtnx zzcU_>XL5jhcvg?u9H^6DpI=nq$IEtqtl9`|>jp@!G{Cy@$=tLrr*_wGoD#p$ZE^ES z+*1P5`cd`-9CtZygh_v(T!8KxH^V$v0)HY(CMouPLLV5LLVANFQ{G6HUeeEUyy(Xc zHo);$!+G<5{Qb-4PnO@5nM_}nDqOzw2MOJd!&^z2lwMj}5puer}DDEXqfF@B`)Mt!DrL literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/window/window_ranking.bin b/crates/krishiv-plan-codec/tests/golden/window/window_ranking.bin new file mode 100644 index 0000000000000000000000000000000000000000..0b6145978678bce25aa70103a04e443f27f8884b GIT binary patch literal 870 zcmbtSv5wO~5ZzgCoXuP=@of%aEDFmNvV??Ep)25&3LzQ{6++@xv6EOac4Ti9prS!K zIto63AK()xkocG@_=hm-9OO>Wv8(sI^JZsOGalLTL~zdW1D9R(44>j*v(okQ)l?f@ zRCDa6*-WnsgJ6`0;&qi~N}}7CI-eR9VaHUwuO%$uTO4V?q>ugc2o9u4u3JxSF z&;gg75Q3iM3K}wc;q%x*&@m!3*vZQ>tH%;Ka8LMfl1oQfcW~Bcj7i#ziA!b-$T}u? z#2y>4A^GerkJu-K1iHm<6HC@7??D?IT6q`V<}hST`tqWC)avZwVNkzYDGn^)J_JH? z5269GiEB?jPk3EM19jD2M#ErlJ84ZqbVXQ^Fo^d8pQh2^p4ewBO$hNP{;x7^Q9lUh zgj0gGUv9TsTb8{~m0;L@ExA;pOWQ8qvL)LpTjuBMab26kF54b?*;*TA^V*n{uxl^- ztG&ConQz-|muUAct|9w>e~YFZ;~t79*~};?c1Jx^@MetIN%Hx^1OVd^ilnNHq7gMZ zFKDV5zwb pb::PhysicalPlan { + m.epoch = 0; + m.features.clear(); + m +} + #[test] -fn a_plan_that_needs_a_newer_reader_is_refused() { - let mut m = message("scan"); - m.format_version = FORMAT_VERSION + 1; +fn a_legacy_plan_that_needs_a_newer_reader_is_refused() { + let mut m = legacy(message("scan")); m.min_reader_version = FORMAT_VERSION + 1; let err = decode(&m.encode_to_vec()).unwrap_err(); assert_eq!( @@ -269,26 +276,39 @@ fn a_plan_that_needs_a_newer_reader_is_refused() { } #[test] -fn an_offset_limit_needs_a_v3_reader() { - // A `Limit` with a nonzero offset sets `min_reader_version` to 3: an - // older reader would silently ignore the skipped rows. - assert_eq!(message("limit").min_reader_version, MIN_READER_VERSION); - assert_eq!(message("limit_offset").min_reader_version, 3); +fn readers_that_predate_features_still_refuse_what_they_cannot_read() { + // The frozen format version and the legacy `min_reader_version` gate: + // an offset `Limit` needs a v3 reader (an older one would ignore the + // skipped rows), a `Window` a v4 reader, a plan without features v2. + assert_eq!(FORMAT_VERSION, 5); + for (name, plan) in cases() { + let m = to_message(&plan).unwrap(); + assert_eq!(m.format_version, FORMAT_VERSION, "{name}"); + let expected = match name.as_str() { + "limit_offset" => 3, + "window_ranking" | "window_aggregates" => 4, + _ => MIN_READER_VERSION, + }; + assert_eq!(m.min_reader_version, expected, "{name}"); + } } #[test] -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") { - let m = to_message(&plan).unwrap(); - assert!(m.min_reader_version <= 3, "{name} needs a newer reader"); - } +fn a_legacy_plan_decodes_by_its_content() { + // Written before the header: no epoch and no feature list. + for name in ["scan", "limit_offset", "window_ranking"] { + let (_, plan) = cases().into_iter().find(|(n, _)| n == name).unwrap(); + let bytes = legacy(to_message(&plan).unwrap()).encode_to_vec(); + assert_eq!(decode(&bytes).unwrap(), plan, "{name}"); } + // A legacy plan whose content is newer than its reader version. + let mut m = legacy(message("window_ranking")); + m.min_reader_version = 3; + corrupt(&m, "a construct newer than the plan's reader version"); + // A feature list needs an epoch. + let mut m = message("window_ranking"); + m.epoch = 0; + corrupt(&m, "a feature list without an epoch"); } #[test] @@ -337,11 +357,10 @@ fn a_corrupt_window_is_refused() { #[test] fn a_newer_writer_that_stays_compatible_is_read() { - // N reads N+1 when the writer only added fields: unknown fields (here - // number 999, a varint) are skipped, and the version is higher. + // A writer that only added ignorable fields: unknown fields (here + // number 999, a varint) are skipped. let (_, plan) = cases().remove(0); - let mut m = to_message(&plan).unwrap(); - m.format_version = FORMAT_VERSION + 1; + let m = to_message(&plan).unwrap(); let mut bytes = m.encode_to_vec(); bytes.extend_from_slice(&[0xB8, 0x3E, 0x07]); assert_eq!(decode(&bytes).unwrap(), plan); @@ -509,7 +528,7 @@ fn limits_bound_size_and_depth() { ..Limits::default() }; assert!(matches!( - decode_with(&bytes, &small).unwrap_err(), + decode_with(&bytes, &small, &ReaderFeatures::all()).unwrap_err(), CodecError::Limit(m) if m.contains("bytes") )); let few = Limits { @@ -517,7 +536,7 @@ fn limits_bound_size_and_depth() { ..Limits::default() }; assert!(matches!( - decode_with(&bytes, &few).unwrap_err(), + decode_with(&bytes, &few, &ReaderFeatures::all()).unwrap_err(), CodecError::Limit(m) if m.contains("expressions") )); @@ -542,7 +561,10 @@ fn limits_bound_size_and_depth() { max_depth: 2000, ..Limits::default() }; - assert_eq!(decode_with(&bytes, &roomy).unwrap(), deep); + assert_eq!( + decode_with(&bytes, &roomy, &ReaderFeatures::all()).unwrap(), + deep + ); } #[test] @@ -567,12 +589,15 @@ fn limits_bound_expression_depth() { max_depth: 51, ..Limits::default() }; - assert_eq!(decode_with(&bytes, &exact).unwrap(), plan); + assert_eq!( + decode_with(&bytes, &exact, &ReaderFeatures::all()).unwrap(), + plan + ); let tight = Limits { max_depth: 50, ..Limits::default() }; - let err = decode_with(&bytes, &tight).unwrap_err(); + let err = decode_with(&bytes, &tight, &ReaderFeatures::all()).unwrap_err(); assert!( matches!(&err, CodecError::Limit(m) if m.contains("expression deeper than 50")), "{err:?}" diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 5945829d..18db6879 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -1,11 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // // krishiv.plan.v1: the stable serialized form of a physical plan (ADR 0003, -// P2-T173; design P2-T171 I1). Compatibility rules (ADR 0026): N reads N-1 -// and N-1 reads N. Fields are only ever added with new numbers; a removed -// field's number and name are `reserved`; an enum value is never reused. A -// writer that needs semantics an older reader would silently ignore raises -// `min_reader_version`, and an older reader rejects the plan. +// P2-T173; design P2-T171 I1). Compatibility (ADR 0040): a plan's header +// lists the features it uses, and a reader decodes every plan of its epoch +// whose features it knows. Fields are only ever added with new numbers; a +// removed field's number and name are `reserved`; an enum value is never +// reused. +// +// Every element added from here on carries a comment line saying which +// class it is (`tools/plan-protogen --check` enforces it; the elements that +// predate the rule are listed in `tools/plan-protogen/baseline.txt`): +// - `feature: `: a reader without the feature would misread it or fail +// with an unhelpful error (a new oneof variant, enum value or message, a +// field whose absence changes results). The id is in the codec's +// registry (`crates/krishiv-plan-codec/src/features.rs`); the fields of a +// new message inherit its class. +// - `ignorable`: an older reader may skip it and still compute the same +// result (a hint, an explain detail). // // Plans are DAG-flat: nodes, expressions and explain nodes live in tables in // post-order and refer to children by table index (always an earlier entry), @@ -19,9 +30,12 @@ syntax = "proto3"; package krishiv.plan.v1; message PhysicalPlan { - // The writer's format version (4). + // Frozen at 5 (the last sequential format version): readers that predate + // features check it with `min_reader_version`. uint32 format_version = 1; - // The oldest reader version that understands every field in use. + // For readers that predate features: 2 for a plan without features, else + // the highest legacy level of its features (3 limit_offset, 4 window, + // 6 for a feature added after the format version was frozen). uint32 min_reader_version = 2; // Plan nodes in post-order; a node's inputs are earlier entries. repeated Node nodes = 3; @@ -35,6 +49,20 @@ message PhysicalPlan { ExprEnv env = 9; // Expressions in post-order; an expression's inputs are earlier entries. repeated Expr exprs = 10; + // The ids of the features the plan uses, sorted and without duplicates. + repeated string features = 11; + // The wire epoch, written explicitly (1). Absent marks a plan written + // before the feature header existed: its features follow from its + // content. + uint32 epoch = 12; +} + +// The header fields of a `PhysicalPlan`, read without decoding its tables. +message PlanHeader { + uint32 format_version = 1; + uint32 min_reader_version = 2; + repeated string features = 11; + uint32 epoch = 12; } message ExprEnv { @@ -67,7 +95,7 @@ message Node { UnionNode union = 12; WriteFilesNode write_files = 13; ExchangeNode exchange = 14; - // v4: a plan with a window needs a v4 reader. + // feature: window WindowNode window = 15; } } @@ -128,7 +156,8 @@ enum SubqueryValueKind { message LimitNode { uint32 input = 1; uint64 limit = 2; - // v3: the number of leading rows to skip. v2 bytes leave this as zero. + // feature: limit_offset (when nonzero) + // The number of leading rows to skip. v2 bytes leave this as zero. uint64 offset = 3; } @@ -156,8 +185,9 @@ message TopNNode { uint64 limit = 3; } -// v4: window functions over an input sorted by the partitioning columns, -// then the ordering. +// feature: window +// Window functions over an input sorted by the partitioning columns, then +// the ordering. message WindowNode { uint32 input = 1; repeated uint64 partition_by = 2; diff --git a/tools/plan-protogen/Cargo.lock b/tools/plan-protogen/Cargo.lock index 4d637787..b8db3d15 100644 --- a/tools/plan-protogen/Cargo.lock +++ b/tools/plan-protogen/Cargo.lock @@ -273,6 +273,7 @@ name = "plan-protogen" version = "0.0.0" dependencies = [ "prost-build", + "prost-types", "protox", ] diff --git a/tools/plan-protogen/Cargo.toml b/tools/plan-protogen/Cargo.toml index 2b1f5073..4407042c 100644 --- a/tools/plan-protogen/Cargo.toml +++ b/tools/plan-protogen/Cargo.toml @@ -13,6 +13,7 @@ publish = false [dependencies] prost-build = "0.14.4" +prost-types = "0.14.4" protox = "0.9" [workspace] diff --git a/tools/plan-protogen/baseline.txt b/tools/plan-protogen/baseline.txt new file mode 100644 index 00000000..75366be8 --- /dev/null +++ b/tools/plan-protogen/baseline.txt @@ -0,0 +1,364 @@ +# The elements of proto/krishiv/plan/v1/plan.proto that predate the compatibility-class +# annotations (ADR 0040). Never add a line: a new element carries a `// feature: ` +# or `// ignorable` comment instead. A line is removed only with its element. +AggregateExpr +AggregateExpr.arg_type +AggregateExpr.args +AggregateExpr.function +AggregateExpr.ignore_nulls +AggregateExpr.input +AggregateExpr.name +AggregateFunction +AggregateFunction.AGGREGATE_FUNCTION_AVG +AggregateFunction.AGGREGATE_FUNCTION_COUNT +AggregateFunction.AGGREGATE_FUNCTION_FIRST +AggregateFunction.AGGREGATE_FUNCTION_LAST +AggregateFunction.AGGREGATE_FUNCTION_MAX +AggregateFunction.AGGREGATE_FUNCTION_MIN +AggregateFunction.AGGREGATE_FUNCTION_SUM +AggregateFunction.AGGREGATE_FUNCTION_UNSPECIFIED +AggregateInput +AggregateInput.AGGREGATE_INPUT_RAW +AggregateInput.AGGREGATE_INPUT_STATE +AggregateInput.AGGREGATE_INPUT_UNSPECIFIED +AggregateNode +AggregateNode.aggregates +AggregateNode.distinct +AggregateNode.group_by +AggregateNode.input +AggregateNode.output +AggregateOutput +AggregateOutput.AGGREGATE_OUTPUT_FINAL +AggregateOutput.AGGREGATE_OUTPUT_PARTIAL +AggregateOutput.AGGREGATE_OUTPUT_UNSPECIFIED +AliasExpr +AliasExpr.expr +AliasExpr.output +BoolList +BoolList.values +BoundColumn +BoundColumn.data_type +BoundColumn.id +BoundColumn.name +BoundColumn.nullable +Cast +Cast.expr +Cast.mode +Cast.to +CastMode +CastMode.CAST_MODE_ANSI +CastMode.CAST_MODE_LEGACY +CastMode.CAST_MODE_TRY +CastMode.CAST_MODE_UNSPECIFIED +ColumnPath +ColumnPath.parts +DataType +DataType.Kind +DataType.Kind.KIND_BINARY +DataType.Kind.KIND_BOOLEAN +DataType.Kind.KIND_BYTE +DataType.Kind.KIND_DATE +DataType.Kind.KIND_DECIMAL +DataType.Kind.KIND_DOUBLE +DataType.Kind.KIND_FLOAT +DataType.Kind.KIND_INTEGER +DataType.Kind.KIND_LONG +DataType.Kind.KIND_NULL +DataType.Kind.KIND_SHORT +DataType.Kind.KIND_STRING +DataType.Kind.KIND_TIMESTAMP +DataType.Kind.KIND_TIMESTAMP_NTZ +DataType.Kind.KIND_UNSPECIFIED +DataType.kind +DataType.precision +DataType.scale +DayTimeField +DayTimeField.DAY_TIME_FIELD_DAY +DayTimeField.DAY_TIME_FIELD_HOUR +DayTimeField.DAY_TIME_FIELD_MINUTE +DayTimeField.DAY_TIME_FIELD_SECOND +DayTimeField.DAY_TIME_FIELD_UNSPECIFIED +DayTimeLiteral +DayTimeLiteral.end +DayTimeLiteral.micros +DayTimeLiteral.start +DecimalLiteral +DecimalLiteral.data_type +DecimalLiteral.unscaled +Empty +ExchangeNode +ExchangeNode.id +ExchangeNode.input +ExchangeNode.partitioning +ExplainNode +ExplainNode.children +ExplainNode.detail +ExplainNode.name +Expr +Expr.alias +Expr.cast +Expr.column +Expr.function +Expr.literal +Expr.origin +Expr.outer_column +ExprEnv +ExprEnv.ansi +ExprEnv.query_start_micros +ExprEnv.user +ExprEnv.zone +FilesSource +FilesSource.OptionsEntry +FilesSource.OptionsEntry.key +FilesSource.OptionsEntry.value +FilesSource.format +FilesSource.options +FilesSource.paths +FilterNode +FilterNode.input +FilterNode.layout +FilterNode.predicate +FilterNode.subqueries +FrameBound +FrameBound.FRAME_BOUND_CURRENT_ROW +FrameBound.FRAME_BOUND_UNBOUNDED_FOLLOWING +FrameBound.FRAME_BOUND_UNBOUNDED_PRECEDING +FrameBound.FRAME_BOUND_UNSPECIFIED +FrameUnits +FrameUnits.FRAME_UNITS_RANGE +FrameUnits.FRAME_UNITS_ROWS +FrameUnits.FRAME_UNITS_UNSPECIFIED +Function +Function.args +Function.builtin +Function.data_type +Function.deterministic +Function.distinct +Function.filter +Function.kind +Function.nullable +Function.registry +FunctionKind +FunctionKind.FUNCTION_KIND_AGGREGATE +FunctionKind.FUNCTION_KIND_GENERATOR +FunctionKind.FUNCTION_KIND_SCALAR +FunctionKind.FUNCTION_KIND_UNSPECIFIED +FunctionKind.FUNCTION_KIND_WINDOW +HashJoinNode +HashJoinNode.build +HashJoinNode.build_keys +HashJoinNode.distribution +HashJoinNode.join_type +HashJoinNode.null_safe +HashJoinNode.probe +HashJoinNode.probe_keys +HashJoinNode.residual +HashJoinNode.runtime_filters +HashPartitioning +HashPartitioning.keys +JoinDistribution +JoinDistribution.JOIN_DISTRIBUTION_BROADCAST +JoinDistribution.JOIN_DISTRIBUTION_PARTITIONED +JoinDistribution.JOIN_DISTRIBUTION_UNSPECIFIED +JoinFilter +JoinFilter.id +JoinFilter.key +JoinFilter.options +JoinType +JoinType.JOIN_TYPE_INNER +JoinType.JOIN_TYPE_LEFT_ANTI +JoinType.JOIN_TYPE_LEFT_OUTER +JoinType.JOIN_TYPE_LEFT_SEMI +JoinType.JOIN_TYPE_NULL_AWARE_ANTI +JoinType.JOIN_TYPE_UNSPECIFIED +KeyFilterOptions +KeyFilterOptions.bloom +KeyFilterOptions.bloom_fpp +KeyFilterOptions.bloom_max_bytes +KeyFilterOptions.bloom_max_fpp +KeyFilterOptions.bloom_max_items +KeyFilterOptions.exact_max_values +KeyFilterOptions.max_values +LimitNode +LimitNode.input +LimitNode.limit +Literal +Literal.binary +Literal.boolean +Literal.byte +Literal.date +Literal.day_time +Literal.decimal +Literal.double +Literal.float +Literal.integer +Literal.long +Literal.null +Literal.short +Literal.string +Literal.timestamp +Literal.timestamp_ntz +Literal.year_month +NestedLoopJoinNode +NestedLoopJoinNode.build +NestedLoopJoinNode.join_type +NestedLoopJoinNode.probe +NestedLoopJoinNode.residual +Node +Node.aggregate +Node.exchange +Node.filter +Node.hash_join +Node.limit +Node.nested_loop_join +Node.one_row +Node.project +Node.replicate_rows +Node.scan +Node.sort +Node.top_n +Node.union +Node.write_files +NullLiteral +NullLiteral.data_type +OptionValue +OptionValue.plain +OptionValue.secret +Origin +Origin.plan_id +Origin.sql +Partitioning +Partitioning.broadcast +Partitioning.hash +Partitioning.range +Partitioning.round_robin +Partitioning.single +PhysicalPlan +PhysicalPlan.env +PhysicalPlan.epoch +PhysicalPlan.explain +PhysicalPlan.explain_root +PhysicalPlan.exprs +PhysicalPlan.features +PhysicalPlan.format_version +PhysicalPlan.min_reader_version +PhysicalPlan.nodes +PhysicalPlan.root +PhysicalPlan.schema +PhysicalPlan.warnings +PlanHeader +PlanHeader.epoch +PlanHeader.features +PlanHeader.format_version +PlanHeader.min_reader_version +ProjectColumn +ProjectColumn.index +ProjectColumn.name +ProjectExpr +ProjectExpr.expr +ProjectExpr.name +ProjectItem +ProjectItem.column +ProjectItem.expr +ProjectNode +ProjectNode.input +ProjectNode.items +ProjectNode.layout +ProjectNode.nullable +ProjectNode.subqueries +RangePartitioning +RangePartitioning.keys +ReadSource +ReadSource.files +ReadSource.table +RegistryFunction +RegistryFunction.id +RegistryFunction.name +ReplicateRowsNode +ReplicateRowsNode.count +ReplicateRowsNode.input +ReplicateRowsNode.layout +ReplicateRowsNode.subqueries +Residual +Residual.conjuncts +Residual.layout +RowStageSettings +RowStageSettings.max_pass_rate +RowStageSettings.min_rows +ScanFilter +ScanFilter.exact +ScanFilter.layout +ScanFilter.predicate +ScanNode +ScanNode.filter +ScanNode.id +ScanNode.output +ScanNode.projection +ScanNode.row_stage +ScanNode.runtime_filters +ScanNode.source +ScanNode.version +ScanRuntimeFilter +ScanRuntimeFilter.column +ScanRuntimeFilter.id +ScanRuntimeFilter.key_type +SecretRef +SecretRef.key +SecretRef.scope +SortKey +SortKey.column +SortKey.descending +SortKey.nulls_first +SortNode +SortNode.input +SortNode.keys +SqlSpan +SqlSpan.end +SqlSpan.source +SqlSpan.start +SubqueryPlan +SubqueryPlan.data_type +SubqueryPlan.id +SubqueryPlan.kind +SubqueryPlan.plan +SubqueryValueKind +SubqueryValueKind.SUBQUERY_VALUE_KIND_EXISTS +SubqueryValueKind.SUBQUERY_VALUE_KIND_SCALAR +SubqueryValueKind.SUBQUERY_VALUE_KIND_UNSPECIFIED +TableSource +TableSource.OptionsEntry +TableSource.OptionsEntry.key +TableSource.OptionsEntry.value +TableSource.name +TableSource.options +TopNNode +TopNNode.input +TopNNode.keys +TopNNode.limit +UnionNode +UnionNode.inputs +WindowFrame +WindowFrame.lower +WindowFrame.units +WindowFrame.upper +WindowFunctionPlan +WindowFunctionPlan.aggregate +WindowFunctionPlan.cume_dist +WindowFunctionPlan.dense_rank +WindowFunctionPlan.filter +WindowFunctionPlan.frame +WindowFunctionPlan.name +WindowFunctionPlan.ntile +WindowFunctionPlan.percent_rank +WindowFunctionPlan.rank +WindowFunctionPlan.row_number +WriteFilesNode +WriteFilesNode.input +YearMonthField +YearMonthField.YEAR_MONTH_FIELD_MONTH +YearMonthField.YEAR_MONTH_FIELD_UNSPECIFIED +YearMonthField.YEAR_MONTH_FIELD_YEAR +YearMonthLiteral +YearMonthLiteral.end +YearMonthLiteral.months +YearMonthLiteral.start diff --git a/tools/plan-protogen/src/classes.rs b/tools/plan-protogen/src/classes.rs new file mode 100644 index 00000000..1322549c --- /dev/null +++ b/tools/plan-protogen/src/classes.rs @@ -0,0 +1,307 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Every element of `krishiv.plan.v1` added after the feature header (ADR +//! 0040) says which compatibility class it is, in a comment line: +//! +//! - `feature: `: a reader without the feature would misread it; `` +//! must be in the codec's feature registry; +//! - `ignorable`: an older reader may skip it. +//! +//! The elements that predate the rule are listed in `baseline.txt`. A new +//! message or enum may carry the annotation once: its fields or values +//! inherit it. An element in the baseline that no longer exists is an error +//! too, so a removal (which must leave a `reserved` entry) is deliberate. + +use std::collections::BTreeSet; + +use prost_types::{DescriptorProto, EnumDescriptorProto, FileDescriptorProto, SourceCodeInfo}; + +// Field numbers of descriptor.proto, for `SourceCodeInfo` paths. +const FILE_MESSAGE: i32 = 4; +const FILE_ENUM: i32 = 5; +const MESSAGE_FIELD: i32 = 2; +const MESSAGE_NESTED: i32 = 3; +const MESSAGE_ENUM: i32 = 4; +const ENUM_VALUE: i32 = 2; + +/// The class an element's comment declares. +#[derive(Debug, Clone, PartialEq, Eq)] +enum Class { + Feature(String), + Ignorable, +} + +/// A named element of the schema and the class its comments declare. +#[derive(Debug)] +struct Element { + name: String, + class: Option, + /// The enclosing message or enum, by index into the element list. + parent: Option, +} + +fn class_of(info: Option<&SourceCodeInfo>, path: &[i32]) -> Option { + let location = info?.location.iter().find(|l| l.path == path)?; + let comments = [&location.leading_comments, &location.trailing_comments]; + comments + .into_iter() + .flatten() + .flat_map(|c| c.lines()) + .map(str::trim) + .find_map(|line| { + if let Some(id) = line.strip_prefix("feature:") { + let id = id.split_whitespace().next().unwrap_or_default(); + Some(Class::Feature(id.to_owned())) + } else if line == "ignorable" || line.starts_with("ignorable:") { + Some(Class::Ignorable) + } else { + None + } + }) +} + +struct Walker<'a> { + info: Option<&'a SourceCodeInfo>, + out: Vec, +} + +impl Walker<'_> { + fn push(&mut self, name: String, path: &[i32], parent: Option) -> usize { + let class = class_of(self.info, path); + self.out.push(Element { + name, + class, + parent, + }); + self.out.len() - 1 + } + + fn message(&mut self, m: &DescriptorProto, prefix: &str, path: &[i32], parent: Option) { + let name = format!("{prefix}{}", m.name()); + let at = self.push(name.clone(), path, parent); + for (i, f) in m.field.iter().enumerate() { + let field_path = [path, &[MESSAGE_FIELD, index(i)]].concat(); + self.push(format!("{name}.{}", f.name()), &field_path, Some(at)); + } + for (i, n) in m.nested_type.iter().enumerate() { + let nested = [path, &[MESSAGE_NESTED, index(i)]].concat(); + self.message(n, &format!("{name}."), &nested, Some(at)); + } + for (i, e) in m.enum_type.iter().enumerate() { + let nested = [path, &[MESSAGE_ENUM, index(i)]].concat(); + self.enumeration(e, &format!("{name}."), &nested, Some(at)); + } + } + + fn enumeration( + &mut self, + e: &EnumDescriptorProto, + prefix: &str, + path: &[i32], + parent: Option, + ) { + let name = format!("{prefix}{}", e.name()); + let at = self.push(name.clone(), path, parent); + for (i, v) in e.value.iter().enumerate() { + let value_path = [path, &[ENUM_VALUE, index(i)]].concat(); + self.push(format!("{name}.{}", v.name()), &value_path, Some(at)); + } + } +} + +fn index(i: usize) -> i32 { + i32::try_from(i).unwrap_or(i32::MAX) +} + +fn elements(file: &FileDescriptorProto) -> Vec { + let mut w = Walker { + info: file.source_code_info.as_ref(), + out: Vec::new(), + }; + for (i, m) in file.message_type.iter().enumerate() { + w.message(m, "", &[FILE_MESSAGE, index(i)], None); + } + for (i, e) in file.enum_type.iter().enumerate() { + w.enumeration(e, "", &[FILE_ENUM, index(i)], None); + } + w.out +} + +/// The names of every element of `file`, sorted (the baseline's form). +#[cfg(test)] +fn names(file: &FileDescriptorProto) -> Vec { + let mut names: Vec<_> = elements(file).into_iter().map(|e| e.name).collect(); + names.sort(); + names +} + +/// The feature ids in the registry source `features_rs`: the first item of +/// each `(id, Variant, ...)` registry line. +pub fn registered_features(features_rs: &str) -> BTreeSet { + features_rs + .lines() + .filter_map(|l| l.trim().strip_prefix('(')) + .filter_map(|l| l.split_once(',')) + .map(|(id, _)| id.trim().to_owned()) + .filter(|id| { + !id.is_empty() + && id + .chars() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_') + }) + .collect() +} + +/// The problems with `file`'s annotations against `baseline` and the +/// registered `features`, one message each. +pub fn check( + file: &FileDescriptorProto, + baseline: &BTreeSet, + features: &BTreeSet, +) -> Vec { + let elements = elements(file); + let mut problems = Vec::new(); + for e in &elements { + if let Some(Class::Feature(id)) = &e.class + && !features.contains(id) + { + problems.push(format!( + "{}: `feature: {id}` is not in the feature registry", + e.name + )); + } + if baseline.contains(&e.name) { + continue; + } + // The element itself or an enclosing new element declares a class. + let mut at = Some(e); + let mut declared = false; + while let Some(x) = at { + if x.class.is_some() { + declared = true; + break; + } + at = x + .parent + .map(|p| &elements[p]) + .filter(|p| !baseline.contains(&p.name)); + } + if !declared { + problems.push(format!( + "{}: a new element needs a `// feature: ` or `// ignorable` comment", + e.name + )); + } + } + let present: BTreeSet<&str> = elements.iter().map(|e| e.name.as_str()).collect(); + for name in baseline { + if !present.contains(name.as_str()) { + problems.push(format!( + "{name}: in baseline.txt but no longer in the schema (reserve its number and \ + remove the line)" + )); + } + } + problems +} + +#[cfg(test)] +mod tests { + use super::*; + + fn compile(source: &str) -> FileDescriptorProto { + let dir = std::env::temp_dir().join(format!( + "plan-protogen-test-{}-{:x}", + std::process::id(), + source.len() * 31 + source.bytes().map(usize::from).sum::() + )); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write(dir.join("t.proto"), source).unwrap(); + let fds = protox::compile(["t.proto"], [&dir]); + std::fs::remove_dir_all(&dir).unwrap(); + fds.unwrap().file.into_iter().next().unwrap() + } + + const BASE: &str = "syntax = \"proto3\";\npackage t;\n\ + message A {\n uint32 x = 1;\n}\nenum E {\n E_UNSPECIFIED = 0;\n}\n"; + + fn set(v: &[&str]) -> BTreeSet { + v.iter().map(|s| (*s).to_owned()).collect() + } + + fn baseline() -> BTreeSet { + set(&["A", "A.x", "E", "E.E_UNSPECIFIED"]) + } + + #[test] + fn the_baseline_needs_no_annotation() { + let file = compile(BASE); + assert_eq!(names(&file), Vec::from_iter(baseline())); + assert!(check(&file, &baseline(), &set(&[])).is_empty()); + } + + #[test] + fn a_new_field_or_value_needs_a_class() { + let file = compile( + &BASE + .replace("uint32 x = 1;", "uint32 x = 1;\n uint32 y = 2;") + .replace("E_UNSPECIFIED = 0;", "E_UNSPECIFIED = 0;\n E_NEW = 1;"), + ); + let problems = check(&file, &baseline(), &set(&[])); + assert_eq!(problems.len(), 2, "{problems:?}"); + assert!( + problems[0].starts_with("A.y: a new element needs"), + "{problems:?}" + ); + assert!( + problems[1].starts_with("E.E_NEW: a new element needs"), + "{problems:?}" + ); + } + + #[test] + fn annotated_elements_pass_and_new_messages_pass_it_on() { + let source = BASE.replace( + "uint32 x = 1;", + "uint32 x = 1;\n // feature: f\n B y = 2;\n uint32 z = 3; // ignorable: a hint\n", + ) + "// feature: f\n// What B holds.\nmessage B {\n uint32 w = 1;\n}\n"; + let file = compile(&source); + assert!(check(&file, &baseline(), &set(&["f"])).is_empty()); + // The id must be registered. + let problems = check(&file, &baseline(), &set(&["g"])); + assert_eq!(problems.len(), 2, "{problems:?}"); + assert!( + problems + .iter() + .all(|p| p.contains("`feature: f` is not in the feature registry")) + ); + } + + #[test] + fn a_field_of_a_baseline_message_does_not_inherit() { + // `A` predates the rule, so a field added to it declares its own class. + let source = BASE + .replace("message A {", "// feature: f\nmessage A {") + .replace("uint32 x = 1;", "uint32 x = 1;\n uint32 y = 2;"); + let problems = check(&compile(&source), &baseline(), &set(&["f"])); + assert_eq!(problems.len(), 1, "{problems:?}"); + assert!(problems[0].starts_with("A.y:"), "{problems:?}"); + } + + #[test] + fn a_removed_baseline_element_is_reported() { + let file = compile(&BASE.replace("uint32 x = 1;", "reserved 1;\n reserved \"x\";")); + let problems = check(&file, &baseline(), &set(&[])); + assert_eq!(problems.len(), 1, "{problems:?}"); + assert!( + problems[0].starts_with("A.x: in baseline.txt but no longer"), + "{problems:?}" + ); + } + + #[test] + fn registry_lines_are_parsed() { + let src = "registry! {\n (alpha_one, AlphaOne, 3, Semantic, \"a\"),\n \ + (beta, Beta, 6, Optimization, \"b\"),\n}\n (Not, an id)\n"; + assert_eq!(registered_features(src), set(&["alpha_one", "beta"])); + } +} diff --git a/tools/plan-protogen/src/main.rs b/tools/plan-protogen/src/main.rs index 1b6935f1..7f86c018 100644 --- a/tools/plan-protogen/src/main.rs +++ b/tools/plan-protogen/src/main.rs @@ -10,7 +10,14 @@ //! Run it from the repository root. `--check` regenerates in memory and //! fails when the file on disk differs, so CI notices a proto change without //! a regeneration. The compiler is `protox`: no `protoc` is needed. +//! +//! Both modes also check the compatibility-class annotations of the schema +//! (see [`classes`]) against `tools/plan-protogen/baseline.txt` and the +//! codec's feature registry. + +mod classes; +use std::collections::BTreeSet; use std::path::Path; use prost_build::Module; @@ -19,6 +26,8 @@ use std::process::ExitCode; const PROTO_DIR: &str = "proto"; const PROTO: &str = "krishiv/plan/v1/plan.proto"; const OUTPUT: &str = "crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs"; +const BASELINE: &str = "tools/plan-protogen/baseline.txt"; +const FEATURES: &str = "crates/krishiv-plan-codec/src/features.rs"; type Result = std::result::Result>; @@ -38,6 +47,7 @@ fn run() -> Result<()> { return Err("pass --write or --check".into()); } let fds = protox::compile([PROTO], [PROTO_DIR])?; + check_classes(&fds)?; let mut config = prost_build::Config::new(); // Deterministic iteration order for map fields (read options). config.btree_map(["."]); @@ -63,3 +73,24 @@ fn run() -> Result<()> { } Ok(()) } + +fn check_classes(fds: &prost_types::FileDescriptorSet) -> Result<()> { + let file = fds + .file + .iter() + .find(|f| f.name() == PROTO) + .ok_or("the plan schema was not compiled")?; + let baseline: BTreeSet = std::fs::read_to_string(BASELINE)? + .lines() + .map(str::trim) + .filter(|l| !l.is_empty() && !l.starts_with('#')) + .map(str::to_owned) + .collect(); + let features = classes::registered_features(&std::fs::read_to_string(FEATURES)?); + let problems = classes::check(file, &baseline, &features); + if problems.is_empty() { + Ok(()) + } else { + Err(format!("{PROTO}:\n {}", problems.join("\n ")).into()) + } +} From 1bfbfe9b1d97b41740f4b331e990516f2a7f3db7 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:35:58 +0200 Subject: [PATCH 8/9] plan-codec: bound the number of feature ids in a header Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-plan-codec/src/decode.rs | 6 ++++++ crates/krishiv-plan-codec/tests/features.rs | 3 +++ 2 files changed, 9 insertions(+) diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 391aa010..93bfee3c 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -133,6 +133,9 @@ enum Header { /// The longest feature id a reader accepts. const MAX_FEATURE_ID: usize = 64; +/// The most feature ids a header may list. +const MAX_FEATURES: usize = 256; + fn header(format_version: u32, min_reader: u32, epoch: u32, ids: &[String]) -> Result
{ if format_version == 0 { return Err(corrupt("no format version")); @@ -155,6 +158,9 @@ fn header(format_version: u32, min_reader: u32, epoch: u32, ids: &[String]) -> R supported: EPOCH, }); } + if ids.len() > MAX_FEATURES { + return Err(corrupt("more feature ids than any plan uses")); + } for id in ids { let well_formed = id.len() <= MAX_FEATURE_ID && id.starts_with(|c: char| c.is_ascii_lowercase()) diff --git a/crates/krishiv-plan-codec/tests/features.rs b/crates/krishiv-plan-codec/tests/features.rs index 4f74770a..650816bb 100644 --- a/crates/krishiv-plan-codec/tests/features.rs +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -177,6 +177,9 @@ fn malformed_headers_are_corrupt() { corrupt(&m, "a malformed feature id"); } let mut m = message("window_ranking"); + m.features = (0..257).map(|i| format!("f{i:04}")).collect(); + corrupt(&m, "more feature ids than any plan uses"); + let mut m = message("window_ranking"); m.features = vec!["window".to_owned(), "window".to_owned()]; corrupt(&m, "feature ids that are not sorted or repeat"); let mut m = message("window_ranking"); From 7f534f25d7c2ce395a64f391431db0bae9e0f512 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Sat, 24 Oct 2026 13:11:11 +0200 Subject: [PATCH 9/9] plan-protogen: classify the header fields; read ids only from the registry The header fields (features, epoch) and PlanHeader are ignorable, since readers that predate features check min_reader_version; they leave the pre-annotation baseline so --check requires their annotations. Feature ids come only from the registry! invocation, not from look-alike lines elsewhere in features.rs. Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .../src/proto/krishiv.plan.v1.pb.rs | 3 +++ proto/krishiv/plan/v1/plan.proto | 3 +++ tools/plan-protogen/baseline.txt | 7 ------- tools/plan-protogen/src/classes.rs | 17 ++++++++++++++++- 4 files changed, 22 insertions(+), 8 deletions(-) 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 c1216139..d8ac9c3c 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 @@ -33,15 +33,18 @@ pub struct PhysicalPlan { #[prost(message, repeated, tag = "10")] pub exprs: ::prost::alloc::vec::Vec, /// The ids of the features the plan uses, sorted and without duplicates. + /// ignorable: a reader that predates features checks min_reader_version. #[prost(string, repeated, tag = "11")] pub features: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, /// The wire epoch, written explicitly (1). Absent marks a plan written /// before the feature header existed: its features follow from its /// content. + /// ignorable: a reader that predates features checks min_reader_version. #[prost(uint32, tag = "12")] pub epoch: u32, } /// The header fields of a `PhysicalPlan`, read without decoding its tables. +/// ignorable: a reader-side view of `PhysicalPlan`, never written as such. #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct PlanHeader { #[prost(uint32, tag = "1")] diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 23e644f6..9ffa6f4d 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -50,14 +50,17 @@ message PhysicalPlan { // Expressions in post-order; an expression's inputs are earlier entries. repeated Expr exprs = 10; // The ids of the features the plan uses, sorted and without duplicates. + // ignorable: a reader that predates features checks min_reader_version. repeated string features = 11; // The wire epoch, written explicitly (1). Absent marks a plan written // before the feature header existed: its features follow from its // content. + // ignorable: a reader that predates features checks min_reader_version. uint32 epoch = 12; } // The header fields of a `PhysicalPlan`, read without decoding its tables. +// ignorable: a reader-side view of `PhysicalPlan`, never written as such. message PlanHeader { uint32 format_version = 1; uint32 min_reader_version = 2; diff --git a/tools/plan-protogen/baseline.txt b/tools/plan-protogen/baseline.txt index 75366be8..7ff88bb4 100644 --- a/tools/plan-protogen/baseline.txt +++ b/tools/plan-protogen/baseline.txt @@ -235,22 +235,15 @@ Partitioning.round_robin Partitioning.single PhysicalPlan PhysicalPlan.env -PhysicalPlan.epoch PhysicalPlan.explain PhysicalPlan.explain_root PhysicalPlan.exprs -PhysicalPlan.features PhysicalPlan.format_version PhysicalPlan.min_reader_version PhysicalPlan.nodes PhysicalPlan.root PhysicalPlan.schema PhysicalPlan.warnings -PlanHeader -PlanHeader.epoch -PlanHeader.features -PlanHeader.format_version -PlanHeader.min_reader_version ProjectColumn ProjectColumn.index ProjectColumn.name diff --git a/tools/plan-protogen/src/classes.rs b/tools/plan-protogen/src/classes.rs index 1322549c..0dc23534 100644 --- a/tools/plan-protogen/src/classes.rs +++ b/tools/plan-protogen/src/classes.rs @@ -135,10 +135,15 @@ fn names(file: &FileDescriptorProto) -> Vec { } /// The feature ids in the registry source `features_rs`: the first item of -/// each `(id, Variant, ...)` registry line. +/// each `(id, Variant, ...)` line inside the `registry! { ... }` invocation. +/// Lines elsewhere (the macro's definition, comments, other code) never +/// register an id. pub fn registered_features(features_rs: &str) -> BTreeSet { features_rs .lines() + .skip_while(|l| l.trim() != "registry! {") + .skip(1) + .take_while(|l| l.trim() != "}") .filter_map(|l| l.trim().strip_prefix('(')) .filter_map(|l| l.split_once(',')) .map(|(id, _)| id.trim().to_owned()) @@ -304,4 +309,14 @@ mod tests { (beta, Beta, 6, Optimization, \"b\"),\n}\n (Not, an id)\n"; assert_eq!(registered_features(src), set(&["alpha_one", "beta"])); } + + #[test] + fn registry_like_lines_outside_the_invocation_are_ignored() { + let src = "/*\n (phantom, Phantom, 6, Semantic, \"p\"),\n*/\n\ + macro_rules! registry {\n (ghost, Ghost, 6, Semantic, \"g\"),\n}\n\ + registry! {\n (real, Real, 6, Semantic, \"r\"),\n}\n\ + (after, After, 6, Semantic, \"a\"),\n"; + assert_eq!(registered_features(src), set(&["real"])); + assert!(registered_features("(only, Only, 6, Semantic, \"o\"),\n").is_empty()); + } }