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 01/16] wip: rollup/cube/grouping sets Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/lib.rs | 4 +- crates/krishiv-exec/src/operators/expand.rs | 144 ++++ crates/krishiv-exec/src/operators/mod.rs | 2 + .../krishiv-exec/src/physical/instantiate.rs | 28 +- crates/krishiv-exec/src/physical/mod.rs | 33 +- crates/krishiv-exec/src/planner/build.rs | 98 ++- crates/krishiv-optimizer/src/prune.rs | 46 ++ crates/krishiv-optimizer/src/stats.rs | 8 + crates/krishiv-plan/src/analyzer/analytics.rs | 692 ++++++++++++++++++ crates/krishiv-plan/src/analyzer/bind.rs | 4 + crates/krishiv-plan/src/analyzer/exprs.rs | 10 +- crates/krishiv-plan/src/analyzer/mod.rs | 1 + crates/krishiv-plan/src/analyzer/ops.rs | 127 +++- crates/krishiv-plan/src/analyzer/resolve.rs | 11 +- crates/krishiv-plan/src/analyzer/setops.rs | 2 +- crates/krishiv-plan/src/display.rs | 20 + crates/krishiv-plan/src/encode.rs | 11 + crates/krishiv-plan/src/lib.rs | 3 +- crates/krishiv-plan/src/plan.rs | 20 + crates/krishiv-plan/src/tests/analyzer.rs | 7 +- crates/krishiv-plan/src/tests/inventory.rs | 12 +- crates/krishiv-plan/src/tests/props.rs | 23 +- crates/krishiv-plan/src/walk.rs | 7 + crates/krishiv-session/src/dataframe.rs | 95 ++- .../krishiv-session/tests/scratch_rollup.rs | 84 +++ crates/krishiv-sql/src/lower.rs | 147 +++- 26 files changed, 1575 insertions(+), 64 deletions(-) create mode 100644 crates/krishiv-exec/src/operators/expand.rs create mode 100644 crates/krishiv-plan/src/analyzer/analytics.rs create mode 100644 crates/krishiv-session/tests/scratch_rollup.rs diff --git a/crates/krishiv-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index f6efa4c04..b9081da93 100644 --- a/crates/krishiv-exec/src/lib.rs +++ b/crates/krishiv-exec/src/lib.rs @@ -113,8 +113,8 @@ pub use operator::{ FlushOutput, GlobalState, LocalState, Operator, OperatorOutput, Sink, SinkStatus, Source, }; pub use operators::{ - DEFAULT_MAX_MERGE_FAN_IN, FileScan, FileScanMetrics, FileScanSource, FileSelections, FileWrite, - FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, MemoryScan, + DEFAULT_MAX_MERGE_FAN_IN, Expand, NullOf, FileScan, FileScanMetrics, FileScanSource, FileSelections, + FileWrite, FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, MemoryScan, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RowStageSettings, RuntimeFilters, ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, SortMetrics, TopN, UnionAlign, }; diff --git a/crates/krishiv-exec/src/operators/expand.rs b/crates/krishiv-exec/src/operators/expand.rs new file mode 100644 index 000000000..5f0a5c233 --- /dev/null +++ b/crates/krishiv-exec/src/operators/expand.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Each row once per projection: the first step of `ROLLUP`, `CUBE` and +//! `GROUPING SETS` (Spark's `Expand`). + +use std::sync::Arc; + +use arrow_array::{ArrayRef, RecordBatch, new_null_array}; +use arrow_schema::{DataType, SchemaRef}; +use krishiv_types::batch::Batch; + +use super::{Project, ScalarExpr, local_mut, no_state}; +use crate::context::TaskContext; +use crate::error::{Result, internal}; +use crate::operator::{LocalState, Operator, OperatorOutput}; + +/// Emits, for every input batch, one batch per projection: the batch +/// projected by it. The projections share one output schema. Rows keep their +/// weights in every copy. Output is produced one projection per +/// [`Operator::execute`] call (`HaveMoreOutput` until the last), so only +/// one copy of an input batch exists at a time, and the plain columns of +/// each copy share the input's buffers. +#[derive(Debug)] +pub struct Expand { + projects: Vec>, + schema: SchemaRef, +} + +impl Expand { + /// Expands batches by `projects`. + /// + /// # Errors + /// `INTERNAL_ERROR` if there is no projection, or two of them output + /// different columns. + pub fn try_new(projects: Vec>) -> Result> { + let first = projects + .first() + .ok_or_else(|| internal("an expansion without projections".to_owned()))?; + let schema = first.schema(); + for p in &projects { + let (a, b) = (p.schema(), &schema); + let same = a.fields().len() == b.fields().len() + && a.fields() + .iter() + .zip(b.fields()) + .all(|(x, y)| { + x.name() == y.name() + && x.data_type() == y.data_type() + && x.is_nullable() == y.is_nullable() + }); + if !same { + return Err(internal(format!( + "the projections of an expansion differ: {:?} and {:?}", + a.fields(), + b.fields() + ))); + } + } + Ok(Arc::new(Self { projects, schema })) + } +} + +/// The next projection to emit for the current input batch. +struct Cursor(usize); + +impl Operator for Expand { + fn name(&self) -> &'static str { + "Expand" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn local_state(&self, _ctx: &TaskContext) -> Result { + Ok(Box::new(Cursor(0))) + } + + fn execute(&self, input: &Batch, local: &mut LocalState) -> Result { + let cursor = local_mut::(local, "Expand")?; + if input.num_rows() == 0 { + cursor.0 = 0; + return Ok(OperatorOutput::NeedMoreInput(None)); + } + let project = self + .projects + .get(cursor.0) + .ok_or_else(|| internal("an expansion past its last projection".to_owned()))?; + let OperatorOutput::NeedMoreInput(Some(out)) = project.execute(input, &mut no_state())? + else { + return Err(internal("a projection did not output a batch".to_owned())); + }; + cursor.0 += 1; + if cursor.0 < self.projects.len() { + return Ok(OperatorOutput::HaveMoreOutput(out)); + } + cursor.0 = 0; + Ok(OperatorOutput::NeedMoreInput(Some(out))) + } +} + +/// A column of nulls, in the Arrow type of an input column. An expansion +/// nulls a grouping column in the carrier its other copies have: strings +/// reach an operator as `Utf8` from a registered table and as `Utf8View` +/// from a file. +#[derive(Debug)] +pub struct NullOf { + index: usize, + data_type: DataType, +} + +impl NullOf { + /// Nulls of input column `index` of batches of `input`. + /// + /// # Errors + /// `INTERNAL_ERROR` if the column is not in `input`. + pub fn new(input: &SchemaRef, index: usize) -> Result { + let field = input + .fields() + .get(index) + .ok_or_else(|| internal(format!("null-of index {index} out of range")))?; + Ok(Self { + index, + data_type: field.data_type().clone(), + }) + } +} + +impl ScalarExpr for NullOf { + fn data_type(&self) -> DataType { + self.data_type.clone() + } + + fn nullable(&self) -> bool { + true + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let column = batch + .columns() + .get(self.index) + .ok_or_else(|| internal(format!("null-of index {} out of range", self.index)))?; + Ok(new_null_array(column.data_type(), batch.num_rows())) + } +} diff --git a/crates/krishiv-exec/src/operators/mod.rs b/crates/krishiv-exec/src/operators/mod.rs index f8cf3b0ad..da85c64c6 100644 --- a/crates/krishiv-exec/src/operators/mod.rs +++ b/crates/krishiv-exec/src/operators/mod.rs @@ -25,6 +25,7 @@ //! - **`UNION`.** See [`union`] for how `UNION ALL` fan-in and the `UNION` //! (distinct) breaker are split between tasks. +mod expand; mod expr; mod file_scan; mod file_write; @@ -37,6 +38,7 @@ mod scan; mod sort; mod union; +pub use expand::{Expand, NullOf}; pub use expr::{FnPredicate, Predicate, ScalarExpr}; pub use file_scan::{FileScan, FileScanMetrics, FileScanSource, METADATA_COLUMN}; pub use file_write::{FileWrite, FileWriteSink}; diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index 266642cbd..728028a00 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -30,7 +30,7 @@ use crate::join::{ }; use crate::operator::{Operator, Sink, Source}; use crate::operators::{ - FileWrite, FileWriteSink, Filter, FilterPushdown, Predicate, Project, ProjectItem, + Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, }; use crate::pipeline::{ExecNode, PipelineBuilder}; @@ -278,6 +278,15 @@ impl Instantiation<'_> { } => self.with_subqueries(input, subqueries, |input, v| { replicate_op(input, layout, count, v) }), + PhysNode::Expand { + input, + layout, + projections, + nullable, + subqueries, + } => self.with_subqueries(input, subqueries, |input, v| { + expand_op(input, layout, projections, nullable, v) + }), PhysNode::Limit { input, limit } => { let input = self.node(input)?; let op = crate::operators::Limit::new(Arc::clone(&input.schema), *limit); @@ -547,6 +556,9 @@ pub(crate) fn project_op( name.clone(), Arc::new(Compiled(compile_one(expr, layout, schema, v)?)), ), + ProjectItemPlan::NullOf { index, name } => { + ProjectItem::expr(name.clone(), Arc::new(NullOf::new(schema, *index)?)) + } }) }) .collect::>>()?; @@ -556,6 +568,20 @@ pub(crate) fn project_op( } } +pub(crate) fn expand_op( + schema: &SchemaRef, + layout: &[AttrId], + projections: &[Vec], + nullable: &[bool], + v: &Values<'_>, +) -> Result> { + let projects = projections + .iter() + .map(|items| project_op(schema, layout, items, Some(nullable), v)) + .collect::>>()?; + Ok(Expand::try_new(projects)?) +} + pub(crate) fn replicate_op( schema: &SchemaRef, layout: &[AttrId], diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index f4f1ce667..2ca8958d9 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -51,7 +51,7 @@ pub(crate) use instantiate::node_ids; pub use instantiate::{SpillScopes, WorkerContext, instantiate}; pub(crate) use instantiate::{ Unbuilt, Values, aggregate_node, aggregate_spec, filter_op, hash_join, nested_loop_join, - project_op, replicate_op, scan_filter, sort_op, top_n_op, + expand_op, project_op, replicate_op, scan_filter, sort_op, top_n_op, }; use crate::aggregate::{AggregateExpr, AggregateOutput}; @@ -144,6 +144,20 @@ pub enum PhysNode { /// The subqueries `count` reads. subqueries: Vec, }, + /// Each row of `input` once per projection (`ROLLUP`, `CUBE`, + /// `GROUPING SETS`). + Expand { + /// The input. + input: Box, + /// The attribute of each input column. + layout: Vec, + /// The projections: the same output columns, each computed one way. + projections: Vec>, + /// Whether each output column is nullable. + nullable: Vec, + /// The subqueries the projections read. + subqueries: Vec, + }, /// `input` sorted by `keys` (columns of `input`). Sort { /// The input. @@ -252,6 +266,14 @@ pub enum ProjectItemPlan { /// The value. expr: Arc, }, + /// A column of nulls in the Arrow type of an input column (the cell an + /// expansion nulls). + NullOf { + /// The input column whose type it takes. + index: usize, + /// The output name. + name: String, + }, } /// Scalar or `EXISTS`. @@ -346,6 +368,7 @@ plan_data!(enum PhysNode { Project { input, layout, items, nullable, subqueries }, Limit { input, limit }, ReplicateRows { input, layout, count, subqueries }, + Expand { input, layout, projections, nullable, subqueries }, Sort { input, keys }, TopN { input, keys, limit }, Aggregate { input, aggregate }, @@ -358,7 +381,7 @@ plan_data!(struct ScanNode { id, scan, projection, filter, runtime_filters, row_ plan_data!(struct ScanFilter { layout, predicate, exact }); plan_data!(struct ScanRuntimeFilter { id, column, key_type }); plan_data!(struct JoinFilterPlan { id, key, options }); -plan_data!(enum ProjectItemPlan { Column { index, name }, Expr { name, expr } }); +plan_data!(enum ProjectItemPlan { Column { index, name }, Expr { name, expr }, NullOf { index, name } }); plan_data!( enum SubqueryValueKind { Scalar, @@ -386,6 +409,9 @@ impl PhysNode { | Self::Project { input, subqueries, .. } + | Self::Expand { + input, subqueries, .. + } | Self::ReplicateRows { input, subqueries, .. } => std::iter::once(input.as_ref()) @@ -412,6 +438,9 @@ impl PhysNode { | Self::Project { input, subqueries, .. } + | Self::Expand { + input, subqueries, .. + } | Self::ReplicateRows { input, subqueries, .. } => std::iter::once(input.as_mut()) diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index b7aadefb5..e9fcc721e 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -40,8 +40,8 @@ use crate::operator::{Operator, Source}; use crate::operators::{FilterPushdown, Limit, Project, ProjectItem, ScanPushdowns, SortKey, TopN}; use crate::physical::{ PhysNode, PhysicalPlan, ProjectItemPlan, ScanFilter, ScanId, ScanNode, ScanRuntimeFilter, - SubqueryId, SubqueryPlan, SubqueryValueKind, Values, filter_op, project_op, replicate_op, - scan_filter, sort_op, top_n_op, + SubqueryId, SubqueryPlan, SubqueryValueKind, Values, expand_op, filter_op, project_op, + replicate_op, scan_filter, sort_op, top_n_op, }; /// How far above the planner's first fresh attribute subquery ids start: @@ -348,6 +348,14 @@ impl<'a> Planner<'a> { let input = self.plan(input)?; self.replicate(input, count) } + PlanNode::Expand { + input, + projections, + output, + } => { + let input = self.plan(input)?; + self.expand(input, projections, output) + } PlanNode::OneRowRelation => Ok(one_row()), _ => { let text = p.to_string(); @@ -696,6 +704,92 @@ impl<'a> Planner<'a> { Ok(planned) } + /// Each row once per projection (`ROLLUP`, `CUBE`, `GROUPING SETS`): + /// plain columns pass through without a copy. + fn expand( + &mut self, + input: Planned, + projections: &[Vec>], + output: &[BoundColumn], + ) -> Result { + if projections.iter().any(|row| row.len() != output.len()) { + return Err(internal("an expansion row of the wrong width")); + } + // A column an expansion nulls keeps the Arrow type of the column + // another projection copies. + let copied: Vec> = (0..output.len()) + .map(|j| { + projections + .iter() + .find_map(|row| row.get(j).and_then(|e| column_of(e))) + .and_then(|c| input.column(c).ok()) + }) + .collect(); + let null_of = |j: usize, e: &Expr| { + let null = matches!(&e.node, ExprNode::Literal(Literal::Null(Some(_)))); + copied.get(j).copied().flatten().filter(|_| null) + }; + let computed: Vec<&Expr> = projections + .iter() + .flat_map(|row| row.iter().enumerate()) + .filter(|(j, e)| column_of(e).is_none() && null_of(*j, e).is_none()) + .map(|(_, e)| e.as_ref()) + .collect(); + let prepared = self.prepare(&computed)?; + let mut values = prepared.exprs.iter(); + let mut rows = Vec::with_capacity(projections.len()); + for row in projections { + let mut items = Vec::with_capacity(row.len()); + for (j, (e, out)) in row.iter().zip(output).enumerate() { + let name = out.name.clone(); + if let Some(c) = column_of(e) { + items.push(ProjectItemPlan::Column { + index: input.column(c)?, + name: Some(name), + }); + } else if let Some(index) = null_of(j, e) { + items.push(ProjectItemPlan::NullOf { index, name }); + } else { + let expr = values.next().ok_or_else(|| internal("expansion arity"))?; + items.push(ProjectItemPlan::Expr { + name, + expr: Arc::clone(expr), + }); + } + } + rows.push(items); + } + let nullable: Vec = output.iter().map(|c| c.nullable).collect(); + let layout: Vec = output.iter().map(|c| c.id).collect(); + let op = expand_op( + &input.schema, + &input.layout, + &rows, + &nullable, + &self.values(&prepared), + )?; + let names: Vec<&str> = output.iter().map(|c| c.name.as_str()).collect(); + let text = names.join(", "); + let mut planned = Self::attach( + input, + prepared, + op.schema(), + layout, + ("Expand", format!("{} projections [{text}]", projections.len())), + |input, layout, subqueries| PhysNode::Expand { + input, + layout, + projections: rows, + nullable, + subqueries, + }, + ); + // The rows are repeated with columns nulled: no runtime filter may + // reach a scan through the expansion. + planned.sites.clear(); + Ok(planned) + } + /// `SELECT exprs`: plain columns pass through without a copy. pub(crate) fn project(&mut self, input: Planned, exprs: &[Arc]) -> Result { let computed: Vec<&Expr> = exprs diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index cf56e0ef2..7fb5f23ab 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -76,6 +76,20 @@ impl Graph { grouping_refs(p, &mut need); child(input, Some(need)); } + PlanNode::Expand { + input, + projections, + output, + } => { + let mut need = Attrs::new(); + for j in expand_keep(output, req) { + projections + .iter() + .filter_map(|row| row.get(j)) + .for_each(|e| references_into(e, &mut need)); + } + child(input, Some(need)); + } PlanNode::SubqueryAlias { input, columns, .. } => { // A renaming alias keeps its input whole. child(input, columns.is_empty().then(|| req.clone())); @@ -197,6 +211,26 @@ impl Rewrite<'_> { replace(&rebuilt, PlanNode::Project { input, exprs: kept }) } } + PlanNode::Expand { + input, + projections, + output, + } => { + let kept = expand_keep(output, req); + if kept.len() == output.len() { + rebuilt + } else { + let pick = |row: &Vec>| -> Vec> { + kept.iter().filter_map(|&j| row.get(j).cloned()).collect() + }; + let node = PlanNode::Expand { + input: Arc::clone(input), + projections: projections.iter().map(pick).collect(), + output: kept.iter().filter_map(|&j| output.get(j).cloned()).collect(), + }; + replace(&rebuilt, node) + } + } PlanNode::Aggregate { input, grouping, @@ -304,6 +338,18 @@ fn keep(exprs: &[Arc], req: &Attrs) -> Vec> { kept } +/// The output columns of an expansion that are read (by position), keeping +/// at least one so the row count survives. +fn expand_keep(output: &[BoundColumn], req: &Attrs) -> Vec { + let mut kept: Vec = (0..output.len()) + .filter(|&j| req.contains(&output[j].id)) + .collect(); + if kept.is_empty() && !output.is_empty() { + kept.push(0); + } + kept +} + fn refs(exprs: &[Arc]) -> Attrs { let mut out = Attrs::new(); exprs.iter().for_each(|e| references_into(e, &mut out)); diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index fd00375b1..d96d5abf1 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -239,6 +239,14 @@ impl<'a> Estimator<'a> { // Each distinct row of an `INTERSECT ALL` / `EXCEPT ALL` count: // no count is known, the input's rows bound it from below. | PlanNode::ReplicateRows { input, .. } => Est::clone(&self.plan(input)), + // Every input row once per projection. + PlanNode::Expand { + input, projections, .. + } => { + let est = Est::clone(&self.plan(input)); + let rows = est.rows * to_f64(u64::try_from(projections.len()).unwrap_or(u64::MAX)); + est.with_rows(rows) + } PlanNode::Union { inputs, .. } => { let rows = inputs.iter().map(|i| self.plan(i).rows).sum(); Est::unknown(rows, ids()) diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs new file mode 100644 index 000000000..c6d34c79e --- /dev/null +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -0,0 +1,692 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Grouping analytics: `GROUP BY ROLLUP`, `CUBE` and `GROUPING SETS`, and the +//! `grouping` and `grouping_id` functions. +//! +//! The rewrite is Spark's (the `ResolveGroupingAnalytics` rule and the +//! `Expand` operator): every distinct grouping expression gets a nullable +//! copy `g`, `Expand` repeats each input row once per grouping set (`g` is +//! the expression for a set that has it and `NULL` for a set that has not) +//! with the set's `spark_grouping_id`, and an ordinary aggregate groups by +//! the copies and the id. Bit `n - 1 - j` of the id is 1 when the `j`-th +//! distinct grouping expression (in order of first appearance) is absent +//! from the set. Select items, `HAVING` and `ORDER BY` read the copies where +//! they name a grouping expression (aggregate calls keep reading the +//! input's own columns, which stay in the expansion), `grouping(e)` is bit +//! `e` of the id (a `TINYINT`) and `grouping_id(..)` is the id (a `BIGINT`). +//! +//! Sets that are equal as sets stay separate rows (`GROUPING SETS ((a), (a))` +//! has two rows per group): the expansion then also numbers its sets +//! (`_gen_grouping_pos`) and the aggregate groups by that too. A single +//! empty set over no grouping expression is a grand total, an aggregate +//! without `GROUP BY` (one row over empty input) in which `grouping_id()` is 0. +//! The virtual column `grouping__id` names the id. +//! +//! Exact behaviour, Spark 4.2: `tests/spark-golden/v4.2.0` (`group-analytics`, +//! `grouping_set`, `postgreSQL/groupingsets`). + +use std::collections::HashMap; +use std::sync::Arc; + +use krishiv_common::error::{ + GroupingColumnMismatch, GroupingIdColumnMismatch, GroupingSizeLimitExceeded, + UnsupportedGroupingExpression, WrongNumArgsWithoutSuggestion, +}; +use krishiv_types::DataType; + +use super::bind::unsupported; +use super::exprs::{Bound, Cx, pretty, semantic_key}; +use super::ops::col_expr; +use super::resolve::{Attr, Env, Rel, Resolver}; +use super::{AnalysisError, ArgType}; +use crate::{ + BoundColumn, CastMode, Expr, ExprNode, FunctionCall, FunctionKind, Grouping, Literal, + LogicalPlan, Origin, PlanNode, ResolvedExpr, +}; + +/// The most grouping expressions one `GROUP BY` takes: the grouping id is a +/// `BIGINT` with one bit each. +const MAX_GROUPING: usize = 63; + +/// `grouping` or `grouping_id`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum GroupingFn { + Grouping, + GroupingId, +} + +impl GroupingFn { + pub(super) fn from_name(name: &str) -> Option { + if name.eq_ignore_ascii_case("grouping") { + Some(Self::Grouping) + } else if name.eq_ignore_ascii_case("grouping_id") { + Some(Self::GroupingId) + } else { + None + } + } + + const fn name(self) -> &'static str { + match self { + Self::Grouping => "grouping", + Self::GroupingId => "grouping_id", + } + } + + const fn data_type(self) -> DataType { + match self { + Self::Grouping => DataType::Byte, + Self::GroupingId => DataType::Long, + } + } +} + +/// One distinct grouping expression. +pub(super) struct GroupExpr { + /// Its semantic key. + key: u128, + /// Its attribute in the expansion: a nullable copy. + pub(super) attr: Attr, +} + +/// The expansion an aggregate reads. +pub(super) struct ExpandPlan { + pub(super) plan: Arc, + /// The span of attribute ids the expansion defines. + pub(super) ids: Option<(u64, u64)>, +} + +/// An aggregate's grouping sets. +pub(super) struct Analytics { + /// The distinct grouping expressions, in order of first appearance. + pub(super) exprs: Vec, + /// The grouping id as the query names it (`grouping__id`); `None` for a + /// grand total, where it is 0. + gid: Option, + /// The expansion; `None` for a grand total. + pub(super) expand: Option, +} + +impl Analytics { + /// The attributes `HAVING` and `ORDER BY` may name: the grouping + /// expressions' copies and `grouping__id`. + pub(super) fn keys(&self) -> Vec { + self.exprs + .iter() + .map(|g| g.attr.clone()) + .chain(self.gid.iter().cloned()) + .collect() + } + + /// The position of `e` among the grouping expressions: a copy of one, or + /// an expression equal to one. + fn position(&self, e: &Expr, memo: &mut HashMap) -> Option { + if let ExprNode::Resolved(ResolvedExpr::Column(c)) = &e.node + && let Some(j) = self.exprs.iter().position(|g| g.attr.col.id == c.id) + { + return Some(j); + } + let key = semantic_key(e, memo); + self.exprs.iter().position(|g| g.key == key) + } +} + +fn literal(l: Literal) -> Arc { + Arc::new(Expr::new(ExprNode::Literal(l))) +} + +fn arg_type(data_type: DataType, nullable: bool) -> ArgType { + ArgType { + data_type: Some(data_type), + nullable, + } +} + +fn bound_of(col: &BoundColumn) -> Bound { + Bound { + expr: col_expr(col, Origin::UNKNOWN), + ty: arg_type(col.data_type.clone(), col.nullable), + agg: None, + outer: None, + lateral: Vec::new(), + shown: None, + } +} + +/// Spark's text for an expression in a grouping error: a column carries its +/// attribute id. +fn with_id(e: &Expr) -> String { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Column(c)) => format!("{}#{}", c.name, c.id.0), + _ => pretty(e), + } +} + +/// Whether a subquery is anywhere in `e`. +fn has_subquery(e: &Arc) -> bool { + let mut stack = vec![Arc::clone(e)]; + while let Some(x) = stack.pop() { + if matches!(x.node, ExprNode::Subquery { .. }) { + return true; + } + x.for_each_child(&mut |c| stack.push(Arc::clone(c))); + } + false +} + +/// The shape of a grouping: all its expressions in order, and each grouping +/// set as positions into them. +fn shape(grouping: &Grouping) -> (Vec>, Vec>) { + match grouping { + Grouping::GroupBy(v) => (v.clone(), vec![(0..v.len()).collect()]), + Grouping::Rollup(v) => ( + v.clone(), + (0..=v.len()).rev().map(|n| (0..n).collect()).collect(), + ), + Grouping::Cube(v) => { + // The full set first, then down to the empty one. + let n = v.len(); + let sets = (0..1_usize << n) + .rev() + .map(|m| (0..n).filter(|i| m >> (n - 1 - i) & 1 == 1).collect()) + .collect(); + (v.clone(), sets) + } + Grouping::GroupingSets { exprs, sets } => { + let mut flat = exprs.clone(); + let mut out = Vec::with_capacity(sets.len()); + for set in sets { + out.push((flat.len()..flat.len() + set.len()).collect()); + flat.extend(set.iter().cloned()); + } + (flat, out) + } + } +} + +impl Resolver<'_> { + /// `grouping(e)` or `grouping_id(..)` in a select list, `HAVING` or + /// `ORDER BY`: a call the aggregate resolves once it knows its grouping + /// sets ([`Self::resolve_grouping`]). Anywhere else it is + /// `UNSUPPORTED_GROUPING_EXPRESSION`. + #[inline(never)] + pub(super) fn grouping_call( + &mut self, + which: GroupingFn, + e: &Arc, + kids: &[Bound], + cx: &Cx<'_>, + ) -> Result { + // Spark has no source position for the error in `HAVING` and + // `ORDER BY`. + let origin = if matches!(cx.op, "Filter" | "Sort") { + Origin::UNKNOWN + } else { + e.origin + }; + if !cx.grouping { + return Err(AnalysisError::new(UnsupportedGroupingExpression {}, origin)); + } + if which == GroupingFn::Grouping && kids.len() != 1 { + return Err(AnalysisError::new( + WrongNumArgsWithoutSuggestion { + function_name: "`grouping`".to_owned(), + expected_num: "1".to_owned(), + actual_num: kids.len().to_string(), + docroot: super::exprs::SPARK_DOC_ROOT.to_owned(), + }, + e.origin, + )); + } + let call = FunctionCall { + name: which.name().to_owned(), + args: kids.iter().map(|k| Arc::clone(&k.expr)).collect(), + distinct: false, + filter: None, + }; + Ok(Bound { + expr: Arc::new(Expr::new(ExprNode::Function(call)).with_origin(e.origin)), + ty: arg_type(which.data_type(), false), + agg: kids.iter().find_map(|k| k.agg), + outer: kids.iter().filter_map(|k| k.outer).min(), + lateral: kids.iter().flat_map(|k| k.lateral.clone()).collect(), + shown: None, + }) + } + + /// A resolved `Expand` read again: its expressions bound over its + /// input, its output as it states it. + #[inline(never)] + pub(super) fn expand_node( + &mut self, + origin: Origin, + input: &Arc, + projections: &[Vec>], + output: &[BoundColumn], + env: &Arc, + ) -> Result { + let c = self.child(input, env)?; + let rels = [&c]; + let cx = Self::cx(&c.out, &rels, env, "Expand"); + let mut bound = Vec::with_capacity(projections.len()); + for row in projections { + let mut exprs = Vec::with_capacity(row.len()); + for e in row { + exprs.push(self.bind(e, &cx)?.expr); + } + bound.push(exprs); + } + let out = output + .iter() + .map(|col| Attr { + col: col.clone(), + qual: Arc::from(Vec::new()), + source: None, + }) + .collect(); + let node = PlanNode::Expand { + input: Arc::clone(&c.plan), + projections: bound, + output: output.to_vec(), + }; + Ok(self.make(origin, node, out, &[&c])) + } + + /// `e` with its `grouping` and `grouping_id` calls resolved against + /// `an`. + #[inline(never)] + pub(super) fn resolve_grouping( + &mut self, + e: &Arc, + an: &Analytics, + env: &Arc, + ) -> Result, AnalysisError> { + self.rewrite(e, an, false, env, &mut HashMap::new()) + } + + /// An aggregate over `ROLLUP`, `CUBE` or `GROUPING SETS`. + #[inline(never)] + pub(super) fn grouping_sets( + &mut self, + p: &LogicalPlan, + input: &Arc, + grouping: &Grouping, + aggregates: &[Arc], + env: &Arc, + ) -> Result { + let c = self.child(input, env)?; + let (flat, sets) = shape(grouping); + self.check_star_ordinals(&flat, aggregates)?; + // A single empty set over nothing is a grand total. + let grand = flat.is_empty() && sets.len() == 1; + let gid_col = (!grand).then(|| { + self.attr( + "spark_grouping_id", + &arg_type(DataType::Long, false), + ) + }); + let gid = gid_col.as_ref().map(|col| Attr { + col: BoundColumn { + name: "grouping__id".to_owned(), + ..col.clone() + }, + qual: Arc::from(Vec::new()), + source: None, + }); + let keys: Vec = gid.iter().cloned().collect(); + let items = self.select_items_in(&c, aggregates, env, "Aggregate", Some(&keys))?; + let bound = self.group_keys(&c, &flat, aggregates, &items, env)?; + let mut memo = HashMap::new(); + // The distinct expressions, and the sets as positions among them. + let mut distinct: Vec = Vec::new(); + let mut keys_of: Vec = Vec::new(); + let mut at = Vec::with_capacity(bound.len()); + for (i, b) in bound.iter().enumerate() { + let key = semantic_key(&b.expr, &mut memo); + let j = match keys_of.iter().position(|k| *k == key) { + Some(j) => j, + None => { + keys_of.push(key); + distinct.push(i); + keys_of.len() - 1 + } + }; + at.push(j); + } + let n = distinct.len(); + if n > MAX_GROUPING { + return Err(AnalysisError::new( + GroupingSizeLimitExceeded { + max_size: MAX_GROUPING.to_string(), + }, + p.origin, + )); + } + let sets: Vec> = sets + .into_iter() + .map(|s| { + let mut s: Vec = s.into_iter().map(|i| at[i]).collect(); + s.sort_unstable(); + s.dedup(); + s + }) + .collect(); + if grand { + let an = Analytics { + exprs: Vec::new(), + gid: None, + expand: None, + }; + let items = self.rewrite_items(items, &an, env)?; + return self.build_aggregate(p.origin, &c, Vec::new(), &items, true, Some(Arc::new(an))); + } + let Some(gid_col) = gid_col else { + return Err(unsupported("grouping sets without an id", p.origin)); + }; + let numbered = (0..sets.len()).any(|i| sets[..i].contains(&sets[i])); + // The expansion's input: the child's columns, and the grouping + // expressions that are not one of them. + let mut passed: Vec> = c.out.iter().map(|a| col_expr(&a.col, Origin::UNKNOWN)).collect(); + let mut passed_out: Vec = c.out.to_vec(); + let mut value: Vec = Vec::with_capacity(n); + let mut copies: Vec = Vec::with_capacity(n); + let mut computed = false; + for (j, &i) in distinct.iter().enumerate() { + let b = &bound[i]; + let column = match &b.expr.node { + ExprNode::Resolved(ResolvedExpr::Column(col)) => { + c.out.iter().find(|a| a.col.id == col.id).cloned() + } + _ => None, + }; + let name = match &column { + Some(a) => a.col.name.clone(), + None => pretty(&b.shown()), + }; + let v = match &column { + Some(a) => a.col.clone(), + None => { + let v = self.attr(&name, &b.ty); + passed.push(Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { + expr: Arc::clone(&b.expr), + output: v.clone(), + })))); + passed_out.push(Attr { + col: v.clone(), + qual: Arc::from(Vec::new()), + source: None, + }); + computed = true; + v + } + }; + let ty = arg_type(v.data_type.clone(), true); + let copy = self.attr(&name, &ty); + copies.push(GroupExpr { + key: keys_of[j], + attr: Attr { + col: copy, + qual: column.as_ref().map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), + source: column.and_then(|a| a.source), + }, + }); + value.push(v); + } + let below = if computed { + let out = passed_out.clone(); + self.projection(Origin::UNKNOWN, &c, passed, out) + } else { + c.clone() + }; + let pos_col = numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); + let mut output: Vec = passed_out.iter().map(|a| a.col.clone()).collect(); + output.extend(copies.iter().map(|g| g.attr.col.clone())); + output.push(gid_col.clone()); + output.extend(pos_col.iter().cloned()); + let mut projections = Vec::with_capacity(sets.len()); + for (si, set) in sets.iter().enumerate() { + let mut row: Vec> = passed_out + .iter() + .map(|a| col_expr(&a.col, Origin::UNKNOWN)) + .collect(); + let mut mask: i64 = 0; + for (j, v) in value.iter().enumerate() { + if set.contains(&j) { + row.push(col_expr(v, Origin::UNKNOWN)); + } else { + row.push(literal(Literal::Null(Some(v.data_type.clone())))); + mask |= 1 << (n - 1 - j); + } + } + row.push(literal(Literal::Long(mask))); + if pos_col.is_some() { + row.push(literal(Literal::Integer( + i32::try_from(si).map_err(|_| unsupported("that many grouping sets", p.origin))?, + ))); + } + projections.push(row); + } + let out_attrs: Vec = passed_out + .iter() + .cloned() + .chain(copies.iter().map(|g| g.attr.clone())) + .chain(std::iter::once(Attr { + col: gid_col.clone(), + qual: Arc::from(Vec::new()), + source: None, + })) + .chain(pos_col.iter().map(|col| Attr { + col: col.clone(), + qual: Arc::from(Vec::new()), + source: None, + })) + .collect(); + let node = PlanNode::Expand { + input: Arc::clone(&below.plan), + projections, + output, + }; + let expand = self.make(Origin::UNKNOWN, node, out_attrs, &[&below]); + let mut group: Vec = copies.iter().map(|g| bound_of(&g.attr.col)).collect(); + group.push(bound_of(&gid_col)); + group.extend(pos_col.iter().map(bound_of)); + let an = Analytics { + exprs: copies, + gid, + expand: Some(ExpandPlan { + plan: Arc::clone(&expand.plan), + ids: expand.ids, + }), + }; + let items = self.rewrite_items(items, &an, env)?; + self.build_aggregate(p.origin, &c, group, &items, false, Some(Arc::new(an))) + } + + /// The select `items` over the grouping sets: grouping expressions read + /// their copies, `grouping` and `grouping_id` are resolved. + #[inline(never)] + fn rewrite_items( + &mut self, + items: Vec, + an: &Analytics, + env: &Arc, + ) -> Result, AnalysisError> { + let mut memo = HashMap::new(); + let mut out = Vec::with_capacity(items.len()); + for b in items { + let expr = self.rewrite(&b.expr, an, true, env, &mut memo)?; + let ty = if Arc::ptr_eq(&expr, &b.expr) { + b.ty.clone() + } else { + self.type_of(&expr, env)? + }; + out.push(Bound { expr, ty, ..b }); + } + Ok(out) + } + + /// The type of the resolved `e`. + fn type_of(&mut self, e: &Arc, env: &Arc) -> Result { + let inner = match &e.node { + ExprNode::Resolved(ResolvedExpr::Alias { expr, .. }) => expr, + _ => e, + }; + Ok(self.bind(inner, &Self::cx(&[], &[], env, "Aggregate"))?.ty) + } + + /// `e` over the grouping sets: with `replace`, a subtree that is a + /// grouping expression reads its copy (aggregate calls and subqueries + /// are left as they are); `grouping` and `grouping_id` are resolved. + /// Nodes above a change are typed again, as the copies are nullable. + #[inline(never)] + fn rewrite( + &mut self, + e: &Arc, + an: &Analytics, + replace: bool, + env: &Arc, + memo: &mut HashMap, + ) -> Result, AnalysisError> { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) if f.kind == FunctionKind::Aggregate => { + return Ok(Arc::clone(e)); + } + ExprNode::Subquery { .. } => return Ok(Arc::clone(e)), + ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some() => { + return self.grouping_value(f, an, env, memo); + } + ExprNode::Resolved(ResolvedExpr::Alias { expr, output }) => { + let inner = self.rewrite(expr, an, replace, env, memo)?; + if Arc::ptr_eq(&inner, expr) { + return Ok(Arc::clone(e)); + } + let ty = self.type_of(&inner, env)?; + let output = BoundColumn { + nullable: ty.nullable, + ..output.clone() + }; + return Ok(Arc::new( + Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { + expr: inner, + output, + })) + .with_origin(e.origin), + )); + } + _ => {} + } + if replace && let Some(j) = an.position(e, memo) { + return Ok(col_expr(&an.exprs[j].attr.col, e.origin)); + } + let rebuilt = Arc::clone(e).map_children(|c| self.rewrite(&c, an, replace, env, memo))?; + if Arc::ptr_eq(&rebuilt, e) { + return Ok(rebuilt); + } + self.retype(rebuilt, env) + } + + /// `e`, whose children changed, typed again. + fn retype(&mut self, e: Arc, env: &Arc) -> Result, AnalysisError> { + let origin = e.origin; + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) if !has_subquery(&e) => { + self.call(f.id.name(), f.args.clone(), env, origin) + } + ExprNode::Resolved(ResolvedExpr::Function(f)) => { + // A call over a subquery keeps its type, nullable. + let mut f = f.clone(); + f.nullable = true; + Ok(Arc::new( + Expr::new(ExprNode::Resolved(ResolvedExpr::Function(f))).with_origin(origin), + )) + } + ExprNode::Cast { .. } if !has_subquery(&e) => { + Ok(self.bind(&e, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + } + _ => Ok(e), + } + } + + /// The value of a `grouping` or `grouping_id` call. + fn grouping_value( + &mut self, + f: &FunctionCall, + an: &Analytics, + env: &Arc, + memo: &mut HashMap, + ) -> Result, AnalysisError> { + let n = an.exprs.len(); + let id = match &an.gid { + Some(g) => col_expr(&g.col, Origin::UNKNOWN), + None => literal(Literal::Long(0)), + }; + let positions: Vec> = + f.args.iter().map(|a| an.position(a, memo)).collect(); + match GroupingFn::from_name(&f.name) { + Some(GroupingFn::GroupingId) => { + let exact = positions.len() == n + && positions.iter().enumerate().all(|(i, p)| *p == Some(i)); + if !f.args.is_empty() && !exact { + return Err(AnalysisError::new( + GroupingIdColumnMismatch { + grouping_id_column: f + .args + .iter() + .map(|a| with_id(a)) + .collect::>() + .join(","), + group_by_columns: an + .exprs + .iter() + .map(|g| format!("{}#{}", g.attr.col.name, g.attr.col.id.0)) + .collect::>() + .join(","), + }, + Origin::UNKNOWN, + )); + } + Ok(id) + } + _ => { + let (Some(arg), Some(Some(j))) = (f.args.first(), positions.first()) else { + return Err(AnalysisError::new( + GroupingColumnMismatch { + grouping: f.args.first().map(|a| with_id(a)).unwrap_or_default(), + grouping_columns: an + .exprs + .iter() + .map(|g| format!("{}#{}", g.attr.col.name, g.attr.col.id.0)) + .collect::>() + .join(","), + }, + Origin::UNKNOWN, + )); + }; + let _ = arg; + // Bit `n - 1 - j` of the id, as a `TINYINT`. + let shift = n - 1 - j; + let high = if shift == 0 { + id + } else { + self.call( + "div", + vec![id, literal(Literal::Long(1_i64 << shift))], + env, + Origin::UNKNOWN, + )? + }; + let bit = self.call( + "%", + vec![high, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )?; + let cast = Arc::new(Expr::new(ExprNode::Cast { + expr: bit, + to: DataType::Byte, + mode: CastMode::Default, + })); + Ok(self.bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + } + } + } +} diff --git a/crates/krishiv-plan/src/analyzer/bind.rs b/crates/krishiv-plan/src/analyzer/bind.rs index 3a8dcc69b..02f44f54f 100644 --- a/crates/krishiv-plan/src/analyzer/bind.rs +++ b/crates/krishiv-plan/src/analyzer/bind.rs @@ -227,6 +227,10 @@ fn check_function(ctx: &AnalysisContext, name: &str, origin: Origin) -> Result<( if crate::Builtin::from_name(name).is_some() { return Ok(()); } + // `grouping` and `grouping_id` are the analyzer's too (grouping sets). + if super::analytics::GroupingFn::from_name(name).is_some() { + return Ok(()); + } match ctx .functions .lookup(std::slice::from_ref(&name.to_owned()), ctx.profile) diff --git a/crates/krishiv-plan/src/analyzer/exprs.rs b/crates/krishiv-plan/src/analyzer/exprs.rs index 6eb5d5a27..b16b69299 100644 --- a/crates/krishiv-plan/src/analyzer/exprs.rs +++ b/crates/krishiv-plan/src/analyzer/exprs.rs @@ -21,6 +21,7 @@ use krishiv_common::error::{ use krishiv_types::DataType; use twox_hash::XxHash3_128; +use super::analytics::GroupingFn; use super::bind::{addr, same, unsupported}; use super::catalog::quoted; use super::coerce::{self, Arg, BadCast, Mismatch, ordinal, quoted_type}; @@ -50,6 +51,10 @@ pub(super) struct Cx<'s> { /// Over an aggregate: its input, where aggregate-function arguments /// resolve. pub(super) agg_scope: Option<&'s [Attr]>, + /// Whether `grouping(..)` and `grouping_id(..)` may appear: in the + /// select list, `HAVING` and `ORDER BY` of an aggregate over grouping + /// sets (see `analytics`). + pub(super) grouping: bool, /// A name `scope` holds several times is looked up further rather /// than ambiguous (Spark's `ORDER BY` over a select list). pub(super) lenient: bool, @@ -204,7 +209,7 @@ fn arg_of(b: &Bound) -> Arg<'_> { } /// Spark's documentation root, as its errors cite it. -const SPARK_DOC_ROOT: &str = "https://spark.apache.org/docs/latest"; +pub(super) const SPARK_DOC_ROOT: &str = "https://spark.apache.org/docs/latest"; fn column(c: &BoundColumn) -> ArgType { arg_type(c.data_type.clone(), c.nullable) @@ -444,6 +449,9 @@ impl Resolver<'_> { e.origin, )); } + if let Some(which) = GroupingFn::from_name(&f.name) { + return self.grouping_call(which, e, &kids[..n], cx); + } if let Some(op) = Builtin::from_name(&f.name) { return self.operator_call(op, e.origin, &kids[..n]); } diff --git a/crates/krishiv-plan/src/analyzer/mod.rs b/crates/krishiv-plan/src/analyzer/mod.rs index 44bca0668..560064ff2 100644 --- a/crates/krishiv-plan/src/analyzer/mod.rs +++ b/crates/krishiv-plan/src/analyzer/mod.rs @@ -62,6 +62,7 @@ //! provider's text. Catalog functions are not policy-checked yet (no //! catalog functions exist before P3). +mod analytics; mod bind; mod builtins; mod catalog; diff --git a/crates/krishiv-plan/src/analyzer/ops.rs b/crates/krishiv-plan/src/analyzer/ops.rs index 1e932bd85..8238aec0e 100644 --- a/crates/krishiv-plan/src/analyzer/ops.rs +++ b/crates/krishiv-plan/src/analyzer/ops.rs @@ -14,6 +14,7 @@ use krishiv_common::error::{ UnsupportedFeatureLateralColumnAliasInGroupBy, }; +use super::analytics::Analytics; use super::bind::{BoundTable, same, unsupported}; use super::catalog::quoted; use super::coerce::quoted_type; @@ -244,7 +245,7 @@ impl Resolver<'_> { ) -> Result { let items = self.select_items(c, exprs, env, "Aggregate")?; if items.iter().any(|b| b.agg.is_some()) { - return self.build_aggregate(p.origin, c, Vec::new(), &items, true); + return self.build_aggregate(p.origin, c, Vec::new(), &items, true, None); } Ok(self.lateral_project(p.origin, c, &items, &[])) } @@ -258,6 +259,20 @@ impl Resolver<'_> { exprs: &[Arc], env: &Arc, op: &'static str, + ) -> Result, AnalysisError> { + self.select_items_in(c, exprs, env, op, None) + } + + /// [`Self::select_items`] of an aggregate over grouping sets, which + /// may call `grouping` and `grouping_id` and name `grouping__id` + /// (`extra`, searched after the input). + pub(super) fn select_items_in( + &mut self, + c: &Rel, + exprs: &[Arc], + env: &Arc, + op: &'static str, + extra: Option<&[Attr]>, ) -> Result, AnalysisError> { let rels = [c]; let mut items = Vec::new(); @@ -274,6 +289,8 @@ impl Resolver<'_> { let b = { let mut cx = Self::cx(&c.out, &rels, env, op); cx.lateral = lca.then_some(&lateral[..]); + cx.extra = extra; + cx.grouping = extra.is_some(); self.bind(e, &cx)? }; let b = self.named(b)?; @@ -660,10 +677,16 @@ impl Resolver<'_> { let mut cx = Self::cx(&agg.out, &rels, env, op); cx.extra = Some(&info.keys); cx.agg_scope = Some(&info.input.out); + cx.grouping = info.analytics.is_some(); let mut bounds = Vec::new(); for e in exprs { bounds.push(self.bind(e, &cx)?); } + if let Some(an) = &info.analytics { + for b in &mut bounds { + b.expr = self.resolve_grouping(&b.expr, an, env)?; + } + } let mut memo = HashMap::new(); let mut known: Vec<(u128, Attr)> = info .aggregates @@ -754,8 +777,14 @@ impl Resolver<'_> { /// aliases if it has any; `pass` (attributes of the aggregate) are /// output after the select items. fn aggregate_rel(&self, origin: Origin, info: AggInfo, pass: &[Attr]) -> Rel { + // Over grouping sets, the aggregate reads their expansion; the + // attributes it defines are part of the aggregate's id span. + let (input, defined) = match info.analytics.as_deref().and_then(|a| a.expand.as_ref()) { + Some(e) => (Arc::clone(&e.plan), e.ids), + None => (Arc::clone(&info.input.plan), None), + }; let node = PlanNode::Aggregate { - input: Arc::clone(&info.input.plan), + input, grouping: Grouping::GroupBy(info.grouping.clone()), aggregates: info.aggregates.clone(), }; @@ -766,6 +795,7 @@ impl Resolver<'_> { self.lateral_project(origin, &a, top, pass) } }; + r.ids = super::resolve::hull(r.ids, defined); let info = Arc::new(info); r.widen = Some(Arc::new(Widen::Aggregate(Arc::clone(&info)))); r.agg = Some(info); @@ -781,38 +811,66 @@ impl Resolver<'_> { env: &Arc, ) -> Result { let Grouping::GroupBy(keys) = grouping else { - return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", p.origin)); + return self.grouping_sets(p, input, grouping, aggregates, env); }; let c = self.child(input, env)?; - let conf = self.ctx.conf; - let ordinal = |k: &Expr| match &k.node { - ExprNode::Literal(Literal::Integer(n)) if conf.group_by_ordinal => Some(*n), - _ => None, - }; + self.check_star_ordinals(keys, aggregates)?; + // Spark resolves the select list first, then ordinals, then the + // other keys. + let items = self.select_items(&c, aggregates, env, "Aggregate")?; + let group = self.group_keys(&c, keys, aggregates, &items, env)?; + self.build_aggregate(p.origin, &c, group, &items, false, None) + } + + /// `SELECT *` with `GROUP BY n` is an error (Spark's + /// `STAR_GROUP_BY_POS`). + pub(super) fn check_star_ordinals( + &self, + keys: &[Arc], + aggregates: &[Arc], + ) -> Result<(), AnalysisError> { let star = aggregates .iter() .any(|e| matches!(e.node, ExprNode::Star(_))); - if star && keys.iter().any(|k| ordinal(k).is_some()) { + if star && keys.iter().any(|k| self.ordinal(k).is_some()) { return Err(AnalysisError::new(StarGroupByPos {}, Origin::UNKNOWN)); } - // Spark resolves the select list first, then ordinals, then the - // other keys. - let items = self.select_items(&c, aggregates, env, "Aggregate")?; + Ok(()) + } + + /// `GROUP BY n` with ordinals on: `n`. + fn ordinal(&self, k: &Expr) -> Option { + match &k.node { + ExprNode::Literal(Literal::Integer(n)) if self.ctx.conf.group_by_ordinal => Some(*n), + _ => None, + } + } + + /// The grouping expressions `keys` over `c`: an ordinal names a select + /// item, other names may be select-list aliases. + pub(super) fn group_keys( + &mut self, + c: &Rel, + keys: &[Arc], + aggregates: &[Arc], + items: &[Bound], + env: &Arc, + ) -> Result, AnalysisError> { let mut group: Vec> = vec![None; keys.len()]; for (slot, k) in group.iter_mut().zip(keys) { - if let Some(n) = ordinal(k) { - *slot = Some(self.group_ordinal(n, k.origin, &items, &c)?); + if let Some(n) = self.ordinal(k) { + *slot = Some(self.group_ordinal(n, k.origin, items, c)?); } } let aliases: Vec<(String, Arc)> = aggregates .iter() - .filter(|_| conf.group_by_aliases) + .filter(|_| self.ctx.conf.group_by_aliases) .filter_map(|e| match &e.node { ExprNode::Alias { expr, name } => Some((name.clone(), Arc::clone(expr))), _ => None, }) .collect(); - let rels = [&c]; + let rels = [c]; for (slot, k) in group.iter_mut().zip(keys) { if slot.is_none() { let mut cx = Self::cx(&c.out, &rels, env, "Aggregate"); @@ -820,8 +878,7 @@ impl Resolver<'_> { *slot = Some(self.bind(k, &cx)?); } } - let group = group.into_iter().flatten().collect(); - self.build_aggregate(p.origin, &c, group, &items, false) + Ok(group.into_iter().flatten().collect()) } /// `GROUP BY n`: the `n`-th select item, unaliased. @@ -885,6 +942,7 @@ impl Resolver<'_> { group: Vec, items: &[Bound], global: bool, + analytics: Option>, ) -> Result { let mut memo = HashMap::new(); let gkeys: Vec<(u128, ArgType)> = group @@ -898,22 +956,33 @@ impl Resolver<'_> { .cloned() .collect(); self.check_aggregation(&base, &key_set, &mut memo, origin, global)?; - let keys: Vec = group - .iter() - .filter_map(|b| match &b.expr.node { - ExprNode::Resolved(ResolvedExpr::Column(col)) => { - c.out.iter().find(|a| a.col.id == col.id).cloned() - } - _ => None, - }) - .collect(); + // Over grouping sets, the grouping attributes are the expansion's. + let (keys, scope): (Vec, Vec) = match &analytics { + Some(a) => ( + a.keys(), + c.out.iter().chain(a.exprs.iter().map(|g| &g.attr)).cloned().collect(), + ), + None => ( + group + .iter() + .filter_map(|b| match &b.expr.node { + ExprNode::Resolved(ResolvedExpr::Column(col)) => { + c.out.iter().find(|a| a.col.id == col.id).cloned() + } + _ => None, + }) + .collect(), + c.out.to_vec(), + ), + }; let mut info = AggInfo { input: c.clone(), grouping: group.into_iter().map(|b| b.expr).collect(), aggregates: base.iter().map(|b| Arc::clone(&b.expr)).collect(), - out: base.iter().map(|b| out_attr(b, &c.out)).collect(), + out: base.iter().map(|b| out_attr(b, &scope)).collect(), keys, top: None, + analytics, }; if base.len() == items.len() { return Ok(self.aggregate_rel(origin, info, &[])); @@ -925,7 +994,7 @@ impl Resolver<'_> { let mut top = Vec::new(); for b in items { if b.lateral.is_empty() { - let a = out_attr(b, &c.out); + let a = out_attr(b, &scope); top.push(Bound { expr: col_expr(&a.col, b.expr.origin), ..b.clone() diff --git a/crates/krishiv-plan/src/analyzer/resolve.rs b/crates/krishiv-plan/src/analyzer/resolve.rs index da70c8c97..0c6f4f564 100644 --- a/crates/krishiv-plan/src/analyzer/resolve.rs +++ b/crates/krishiv-plan/src/analyzer/resolve.rs @@ -39,6 +39,9 @@ pub(super) struct AggInfo { /// above the aggregate: every select item, over the aggregate's output /// (`None` when no item does). pub(super) top: Option>, + /// The grouping sets, for an aggregate over `ROLLUP`, `CUBE` or + /// `GROUPING SETS`: its input is then their expansion. + pub(super) analytics: Option>, } /// How a relation gets back columns it dropped: Spark adds the attributes @@ -193,7 +196,7 @@ fn first_free_id(root: &LogicalPlan) -> Result { Ok(next) } -fn hull(a: Option<(u64, u64)>, b: Option<(u64, u64)>) -> Option<(u64, u64)> { +pub(super) fn hull(a: Option<(u64, u64)>, b: Option<(u64, u64)>) -> Option<(u64, u64)> { match (a, b) { (Some(x), Some(y)) => Some((x.0.min(y.0), x.1.max(y.1))), (x, None) | (None, x) => x, @@ -446,6 +449,7 @@ impl Resolver<'_> { aliases: &[], extra: None, agg_scope: None, + grouping: false, lenient: false, outer: true, env, @@ -624,6 +628,11 @@ impl Resolver<'_> { }; Ok(self.make(o, node, c.out.to_vec(), &[&c])) } + PlanNode::Expand { + input, + projections, + output, + } => self.expand_node(o, input, projections, output, env), PlanNode::Distinct { input } => self.distinct(p, input, env), PlanNode::Hint { input, diff --git a/crates/krishiv-plan/src/analyzer/setops.rs b/crates/krishiv-plan/src/analyzer/setops.rs index 9e25f448f..0ac2ae780 100644 --- a/crates/krishiv-plan/src/analyzer/setops.rs +++ b/crates/krishiv-plan/src/analyzer/setops.rs @@ -387,7 +387,7 @@ impl Resolver<'_> { exprs.push(alias(sum, name)); } let items = self.select_items(&u, &exprs, env, "Aggregate")?; - let agg = self.build_aggregate(Origin::UNKNOWN, &u, group, &items, false)?; + let agg = self.build_aggregate(Origin::UNKNOWN, &u, group, &items, false, None)?; let col = |i: usize| col_expr(&agg.out[i].col, Origin::UNKNOWN); let (lc, rc) = (col(n), col(n + 1)); let (predicate, count) = match op { diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index 7483abc98..092b6eeed 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -288,6 +288,26 @@ impl Display for NodeLine<'_> { PlanNode::ReplicateRows { count, .. } => { write!(f, "Generate replicaterows({})", Text(count)) } + // Spark's `Expand [[a, null, 1], [a, b, 0]], [a, b, spark_grouping_id]`. + PlanNode::Expand { + projections, + output, + .. + } => { + f.write_str("Expand [")?; + for (i, p) in projections.iter().enumerate() { + write!(f, "{}[", sep(i))?; + for (j, e) in p.iter().enumerate() { + write!(f, "{}{}", sep(j), Text(e))?; + } + f.write_str("]")?; + } + f.write_str("], [")?; + for (i, c) in output.iter().enumerate() { + write!(f, "{}{}#{}", sep(i), c.name, c.id.0)?; + } + f.write_str("]") + } PlanNode::WithCte { ctes, .. } => { f.write_str("WithCte [")?; for (i, c) in ctes.iter().enumerate() { diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index 2957c5fc4..976a4f499 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -791,6 +791,17 @@ impl<'k, 'a> Encoder<'k, 'a> { self.u8(21); self.expr(count); } + PlanNode::Expand { + projections, + output, + .. + } => { + self.u8(22); + self.len(projections.len()); + projections.iter().for_each(|p| self.exprs(p)); + self.len(output.len()); + output.iter().for_each(|c| self.bound(c)); + } } } diff --git a/crates/krishiv-plan/src/lib.rs b/crates/krishiv-plan/src/lib.rs index 35ddaf2e1..79607784f 100644 --- a/crates/krishiv-plan/src/lib.rs +++ b/crates/krishiv-plan/src/lib.rs @@ -67,7 +67,7 @@ //! | `PlanNode::Project` | U, R | Q2, Q4, Q6-Q9, Q11, Q14-Q17, Q19-Q22 | `select` | //! | `PlanNode::WithColumns` | U | - | `withColumn` | //! | `PlanNode::Filter` | U, R | Q1-Q12, Q14-Q22 | `filter`, `where` | -//! | `PlanNode::Aggregate` | U, R | Q1, Q3-Q5, Q7-Q13, Q15, Q16, Q18, Q21, Q22 | `groupBy.agg` | +//! | `PlanNode::Aggregate` | U, R | Q1, Q3-Q5, Q7-Q13, Q15, Q16, Q18, Q21, Q22 | `groupBy.agg`, `rollup.agg`, `cube.agg` | //! | `PlanNode::Join` | U, R | Q2, Q3, Q5, Q7-Q21 | `join` | //! | `PlanNode::Sort` | U, R | Q1-Q5, Q7-Q13, Q15, Q16, Q18, Q20-Q22 | `orderBy`, `sort` | //! | `PlanNode::Limit` | U, R | Q2, Q3, Q10, Q18, Q21 | `limit` | @@ -81,6 +81,7 @@ //! | `PlanNode::Distinct` | U | - | - | //! | `PlanNode::SetOp` | U | - | - | //! | `PlanNode::ReplicateRows` | R | - | - | +//! | `PlanNode::Expand` | R | - | - | //! | `ExprNode::Column` | U | all | `col`, `df["a"]` | //! | `ExprNode::Literal` | U, R | all | `lit` | //! | `ExprNode::Star` | U | Q1, Q4, Q13, Q21, Q22 | `select("*")` | diff --git a/crates/krishiv-plan/src/plan.rs b/crates/krishiv-plan/src/plan.rs index 59df152d5..d0e791be4 100644 --- a/crates/krishiv-plan/src/plan.rs +++ b/crates/krishiv-plan/src/plan.rs @@ -243,6 +243,19 @@ pub enum PlanNode { /// A `BIGINT` expression over the input: the copies of each row. count: Arc, }, + /// Each input row once per projection (Spark's `Expand`, the first step + /// of `ROLLUP`, `CUBE` and `GROUPING SETS`): row `i` of the output for + /// an input row is the row evaluated by `projections[i]`. Resolved + /// plans only: the analyzer builds it, with one expression per output + /// column in every projection. + Expand { + /// Input. + input: Arc, + /// The projections, in output order; at least one. + projections: Vec>>, + /// The output attributes, aligned with each projection. + output: Vec, + }, } /// The operation of a [`PlanNode::SetOp`]. @@ -554,6 +567,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } + | PlanNode::Expand { input, .. } | PlanNode::ReplicateRows { input, .. } => f(input), PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { f(left); @@ -585,6 +599,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } + | PlanNode::Expand { input, .. } | PlanNode::ReplicateRows { input, .. } => f(input), PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { f(left); @@ -644,6 +659,7 @@ impl LogicalPlan { JoinCondition::None | JoinCondition::Using(_) => {} }, PlanNode::Sort { order, .. } => order.iter().for_each(|o| f(&o.expr)), + PlanNode::Expand { projections, .. } => projections.iter().flatten().for_each(f), } } @@ -694,6 +710,9 @@ impl LogicalPlan { JoinCondition::None | JoinCondition::Using(_) => {} }, PlanNode::Sort { order, .. } => order.iter_mut().for_each(|o| f(&mut o.expr)), + PlanNode::Expand { projections, .. } => { + projections.iter_mut().flatten().for_each(f); + } } } @@ -748,6 +767,7 @@ impl LogicalPlan { ); } PlanNode::OneRowRelation => return Some(Vec::new()), + PlanNode::Expand { output, .. } => return Some(output.clone()), PlanNode::Scan(s) => return Some(s.output.clone()), PlanNode::CteRef(r) => return Some(r.output.clone()), PlanNode::Shift { input, offset } => { diff --git a/crates/krishiv-plan/src/tests/analyzer.rs b/crates/krishiv-plan/src/tests/analyzer.rs index 1e280cebe..c05b3ba6b 100644 --- a/crates/krishiv-plan/src/tests/analyzer.rs +++ b/crates/krishiv-plan/src/tests/analyzer.rs @@ -1423,13 +1423,8 @@ fn unsupported_inputs_fail_explicitly() { paths: vec![krishiv_common::secret::OptionValue::from_user("path", "/d")], options: std::collections::BTreeMap::new(), })); - let rollup = LogicalPlan::new(PlanNode::Aggregate { - input: Arc::new(LogicalPlan::table(["t"])), - grouping: Grouping::Rollup(vec![Arc::new(c("a"))]), - aggregates: vec![Arc::new(c("a"))], - }); let sha = LogicalPlan::table(["t"]).project(vec![Expr::function("factorial", vec![c("b")])]); - for plan in [files, rollup, sha] { + for plan in [files, sha] { assert_eq!(class(&run(plan, &ctx)), "KRISHIV_UNSUPPORTED_FEATURE"); } // An untyped NULL is the void type; an interval is not supported yet. diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index 011e3a136..fcd401999 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -12,7 +12,7 @@ use krishiv_types::DataType; use super::dsl::{c, int}; use super::tpch::QUERIES; use crate::{ - Builtin, ColumnRef, Expr, ExprNode, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, + Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, }; @@ -41,6 +41,7 @@ pub(super) fn kinds(plan: &LogicalPlan, out: &mut BTreeSet<&'static str>) { PlanNode::Distinct { .. } => "PlanNode::Distinct", PlanNode::SetOp { .. } => "PlanNode::SetOp", PlanNode::ReplicateRows { .. } => "PlanNode::ReplicateRows", + PlanNode::Expand { .. } => "PlanNode::Expand", }); plan.for_each_expr(&mut |e| expr_kinds(e, out)); plan.for_each_child(&mut |c| kinds(c, out)); @@ -153,6 +154,13 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { columns: vec![("b".to_owned(), Arc::new(c("a")))], }); let sum = Expr::builtin(Builtin::Add, vec![c("a"), int(1)]); + let grouped = |grouping| { + LogicalPlan::new(PlanNode::Aggregate { + input: Arc::new(df(1)), + grouping, + aggregates: vec![Arc::new(c("a"))], + }) + }; vec![ ("read.table", df(1)), ("read.parquet", parquet), @@ -162,6 +170,8 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { ("filter", df(1).filter(c("a"))), ("where", df(1).filter(c("a"))), ("groupBy.agg", df(1).aggregate(vec![c("a")], vec![c("a")])), + ("rollup.agg", grouped(Grouping::Rollup(vec![Arc::new(c("a"))]))), + ("cube.agg", grouped(Grouping::Cube(vec![Arc::new(c("a"))]))), ( "join", df(1).join(df(2), JoinKind::Inner, JoinCondition::None), diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index 715382a29..ec013bb5d 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -288,7 +288,7 @@ impl Gen<'_> { fn plan(&mut self, depth: u8) -> LogicalPlan { let origin = self.origin(); - let kinds = if depth == 0 { 4 } else { 21 }; + let kinds = if depth == 0 { 4 } else { 22 }; let d = depth.saturating_sub(1); let node = match self.pick(kinds) { 0 => PlanNode::Read(self.source()), @@ -424,6 +424,11 @@ impl Gen<'_> { input: self.input(d), count: Arc::new(self.expr(1)), }, + 21 => PlanNode::Expand { + input: self.input(d), + projections: self.many(1, 2, |g| g.exprs()), + output: self.many(0, 2, Gen::bound), + }, _ => PlanNode::WithCte { input: self.input(d), ctes: self.many(1, 2, |g| Cte { @@ -813,6 +818,22 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { } | PlanNode::Limit { input, limit: e } | PlanNode::ReplicateRows { input, count: e } => m_expr(e, w) || m_plan(input, w), + PlanNode::Expand { + input, + projections, + output, + } => { + site!(w, projections.push(Vec::new())); + site!(w, output.push(output.first().cloned().unwrap_or(BoundColumn { + id: AttrId(0), + name: "x".to_owned(), + data_type: DataType::Integer, + nullable: false, + }))); + projections.iter_mut().any(|p| m_exprs(p, w)) + || output.iter_mut().any(|c| m_bound(c, w)) + || m_plan(input, w) + } PlanNode::SetOp { left, right, diff --git a/crates/krishiv-plan/src/walk.rs b/crates/krishiv-plan/src/walk.rs index f79c41542..ee0d9ed3b 100644 --- a/crates/krishiv-plan/src/walk.rs +++ b/crates/krishiv-plan/src/walk.rs @@ -155,6 +155,7 @@ pub(crate) fn all_resolved(root: Node<'_>) -> bool { | PlanNode::SubqueryAlias { .. } | PlanNode::WithCte { .. } | PlanNode::JoinHint { .. } + | PlanNode::Expand { .. } | PlanNode::ReplicateRows { .. } => true, }, Node::Expr(e) => match &e.node { @@ -257,6 +258,12 @@ impl Teardown { self.plans.push(input); self.exprs.push(e); } + PlanNode::Expand { + input, projections, .. + } => { + self.plans.push(input); + self.exprs.extend(projections.into_iter().flatten()); + } PlanNode::Aggregate { input, grouping, diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index de0b9a94b..5c499c5bd 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -255,9 +255,73 @@ impl DataFrame { /// ``` #[must_use] pub fn group_by(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::GroupBy) + } + + /// Groups rows by `keys` and by each prefix of them, down to the grand + /// total, for [`GroupedData::agg`] (PySpark's `rollup`): a key a row is + /// rolled up over reads `NULL`. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s.table("t").rollup([col("a"), col("b")]).agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn rollup(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::Rollup) + } + + /// Groups rows by every subset of `keys`, for [`GroupedData::agg`] + /// (PySpark's `cube`): a key a row is rolled up over reads `NULL`. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s.table("t").cube([col("a"), col("b")]).agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn cube(&self, keys: impl IntoIterator>) -> GroupedData { + self.grouped(keys, Sets::Cube) + } + + /// Groups rows by each of `sets` (grouping expressions among `keys`), + /// for [`GroupedData::agg`] (PySpark's `groupingSets`): the output has + /// the columns of `keys`, `NULL` where a set does not group by them. + /// + /// ``` + /// # use krishiv_session::functions::{col, sum}; + /// # let s = krishiv_session::Session::new()?; + /// let df = s + /// .table("t") + /// .grouping_sets([vec![col("a")], vec![col("b")], vec![]], [col("a"), col("b")]) + /// .agg([sum(col("v"))]); + /// # Ok::<(), krishiv_common::error::KrishivError>(()) + /// ``` + #[must_use] + pub fn grouping_sets( + &self, + sets: impl IntoIterator>, + keys: impl IntoIterator>, + ) -> GroupedData { + let sets = sets + .into_iter() + .map(|set| set.into_iter().map(|c| c.expr).collect()) + .collect(); + self.grouped(keys, Sets::Explicit(sets)) + } + + fn grouped( + &self, + keys: impl IntoIterator>, + sets: Sets, + ) -> GroupedData { GroupedData { df: self.clone(), keys: keys.into_iter().map(|k| k.into().expr).collect(), + sets, } } @@ -749,11 +813,26 @@ impl DataFrame { } } -/// Rows grouped by keys ([`DataFrame::group_by`]). +/// Which groups of its keys a [`GroupedData`] aggregates. +#[derive(Debug, Clone)] +enum Sets { + /// All the keys together. + GroupBy, + /// Each prefix of the keys. + Rollup, + /// Each subset of the keys. + Cube, + /// These sets. + Explicit(Vec>), +} + +/// Rows grouped by keys ([`DataFrame::group_by`], [`DataFrame::rollup`], +/// [`DataFrame::cube`], [`DataFrame::grouping_sets`]). #[derive(Debug, Clone)] pub struct GroupedData { df: DataFrame, keys: Vec, + sets: Sets, } impl GroupedData { @@ -771,9 +850,21 @@ impl GroupedData { let keys: Vec> = self.keys.iter().cloned().map(Arc::new).collect(); let mut output = keys.clone(); output.extend(aggregates.into_iter().map(|c| Arc::new(c.expr))); + let grouping = match &self.sets { + Sets::GroupBy => Grouping::GroupBy(keys), + Sets::Rollup => Grouping::Rollup(keys), + Sets::Cube => Grouping::Cube(keys), + Sets::Explicit(sets) => Grouping::GroupingSets { + exprs: keys, + sets: sets + .iter() + .map(|set| set.iter().cloned().map(Arc::new).collect()) + .collect(), + }, + }; self.df.derive(PlanNode::Aggregate { input: Arc::clone(&self.df.plan), - grouping: Grouping::GroupBy(keys), + grouping, aggregates: output, }) } diff --git a/crates/krishiv-session/tests/scratch_rollup.rs b/crates/krishiv-session/tests/scratch_rollup.rs new file mode 100644 index 000000000..f12c60530 --- /dev/null +++ b/crates/krishiv-session/tests/scratch_rollup.rs @@ -0,0 +1,84 @@ +//! scratch +#![cfg(test)] +use std::sync::Arc; + +use arrow_array::{Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType as T, Field, Schema}; +use futures::executor::block_on; +use krishiv_session::Session; + +#[test] +fn scratch() { + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", T::Int32, false), + Field::new("b", T::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2, 3, 3])), + Arc::new(Int32Array::from(vec![1, 2, 1, 2, 1, 2])), + ], + ) + .unwrap(); + s.register_batches("td", schema, vec![batch]).unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("course", T::Utf8, false), + Field::new("year", T::Int32, false), + Field::new("earnings", T::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["dotNET", "Java", "dotNET", "dotNET", "Java"])), + Arc::new(Int32Array::from(vec![2012, 2012, 2012, 2013, 2013])), + Arc::new(Int32Array::from(vec![10000, 20000, 5000, 48000, 30000])), + ], + ) + .unwrap(); + s.register_batches("cs", schema, vec![batch]).unwrap(); + for q in std::fs::read_to_string("/home/krishiv/krishiv-data/tmp/s-rollup/q.sql") + .unwrap() + .split(";\n") + .filter(|q| !q.trim().is_empty()) + { + match block_on(s.sql(q)) { + Err(e) => println!("ERR {q}\n {e}"), + Ok(df) => match block_on(df.collect()) { + Err(e) => println!("RUNERR {q}\n {e}"), + Ok(b) => { + println!("OK {q}"); + for b in &b { + println!("{}", arrow_cast_fmt(b)); + } + } + }, + } + } +} + +fn arrow_cast_fmt(b: &RecordBatch) -> String { + let mut out = format!("{:?}\n", b.schema().fields().iter().map(|f| format!("{}:{}{}", f.name(), f.data_type(), if f.is_nullable() {"?"} else {""})).collect::>()); + for r in 0..b.num_rows() { + let cells: Vec = b + .columns() + .iter() + .map(|c| { + if c.is_null(r) { + "NULL".to_owned() + } else { + let d = arrow_array::cast::AsArray::as_primitive_opt::(c) + .map(|a| a.value(r).to_string()) + .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) + .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) + .or_else(|| arrow_array::cast::AsArray::as_string_opt::(c).map(|a| a.value(r).to_owned())); + d.unwrap_or_else(|| "?".to_owned()) + } + }) + .collect(); + out.push_str(&cells.join("\t")); + out.push('\n'); + } + out +} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index 6a4e5462c..8daa1db39 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -27,7 +27,8 @@ use krishiv_plan::{ use krishiv_types::DecimalType; use crate::ast::{ - Alias, BinaryOp, Expr, ExprKind, FromClause, GroupItemKind, JoinConstraint, JoinKind, Literal, + Alias, BinaryOp, Expr, ExprKind, FromClause, GroupBy, GroupItem, GroupItemKind, GroupWith, + GroupingSet, JoinConstraint, JoinKind, Literal, NullTreatment, NumKind, Query, Relation, RelationExt, Select, SetExpr, SetExprKind, SetOperator, SetQuantifier, SortItem, Span, TableFactor, TableFactorKind, Truth, UnaryOp, }; @@ -36,6 +37,9 @@ use crate::{ParseError, ParserOptions, parse_query}; type R = Result; +/// The most grouping sets one `GROUP BY` expands to. +const MAX_GROUPING_SETS: usize = 1 << 16; + /// Parses and lowers one query; spans index `source`. /// /// # Errors @@ -389,22 +393,14 @@ impl Lowerer { s.span, )); } - let mut keys = Vec::new(); - if let Some(g) = &s.group_by { - if g.with.is_some() { - return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", g.span)); - } - for item in &g.items { - match &item.kind { - GroupItemKind::Expr(e) => keys.push(arc(self.expr(e)?)), - _ => return Err(unsupported("ROLLUP, CUBE and GROUPING SETS", item.span)), - } - } - } + let grouping = match &s.group_by { + Some(g) => self.grouping(g)?, + None => Grouping::GroupBy(Vec::new()), + }; plan = self.node( PlanNode::Aggregate { input: arc(plan), - grouping: Grouping::GroupBy(keys), + grouping, aggregates: items, }, s.span, @@ -437,6 +433,129 @@ impl Lowerer { Ok(plan) } + /// `GROUP BY`, as Spark's `AstBuilder` shapes it: plain keys, `WITH + /// ROLLUP` / `WITH CUBE` over them, `GROUPING SETS (..)` after them (the + /// keys only fix the grouping expressions), or items that are + /// `ROLLUP`, `CUBE` and `GROUPING SETS` among keys, whose grouping sets + /// are the cross product of the items' (a key is a set of itself). + #[inline(never)] + fn grouping(&mut self, g: &GroupBy) -> R { + let mut plain = Vec::with_capacity(g.items.len()); + for item in &g.items { + if let GroupItemKind::Expr(e) = &item.kind { + plain.push(arc(self.expr(e)?)); + } + } + let all_plain = plain.len() == g.items.len(); + match (&g.with, all_plain) { + (None, true) => return Ok(Grouping::GroupBy(plain)), + (Some(GroupWith::Rollup), true) => return Ok(Grouping::Rollup(plain)), + (Some(GroupWith::Cube), true) => return Ok(Grouping::Cube(plain)), + (Some(GroupWith::GroupingSets(sets)), true) => { + let sets = sets + .iter() + .map(|s| self.set(s)) + .collect::>>()?; + return Ok(Grouping::GroupingSets { exprs: plain, sets }); + } + (Some(_), false) => { + return Err(unsupported( + "ROLLUP, CUBE and GROUPING SETS items with WITH ROLLUP, WITH CUBE or GROUPING SETS", + g.span, + )); + } + (None, false) => {} + } + let mut exprs = Vec::new(); + let mut product: Vec>> = vec![Vec::new()]; + for item in &g.items { + let sets = self.item_sets(item, &mut exprs)?; + if product.len().saturating_mul(sets.len()) > MAX_GROUPING_SETS { + return Err(unsupported( + &format!("more than {MAX_GROUPING_SETS} grouping sets"), + g.span, + )); + } + product = product + .iter() + .flat_map(|head| { + sets.iter().map(move |tail| { + head.iter().chain(tail).cloned().collect::>() + }) + }) + .collect(); + } + Ok(Grouping::GroupingSets { + exprs, + sets: product, + }) + } + + /// The grouping sets one `GROUP BY` item stands for; the expressions it + /// names are appended to `written` in order. + fn item_sets( + &mut self, + item: &GroupItem, + written: &mut Vec>, + ) -> R>>> { + Ok(match &item.kind { + GroupItemKind::Expr(e) => { + let e = arc(self.expr(e)?); + written.push(Arc::clone(&e)); + vec![vec![e]] + } + GroupItemKind::Set(s) => { + let set = self.set(s)?; + written.extend(set.iter().cloned()); + vec![set] + } + GroupItemKind::GroupingSets(items) => { + let mut sets = Vec::new(); + for item in items { + sets.extend(self.item_sets(item, written)?); + } + sets + } + GroupItemKind::Rollup(elements) | GroupItemKind::Cube(elements) => { + let rollup = matches!(item.kind, GroupItemKind::Rollup(_)); + let mut parts = Vec::with_capacity(elements.len()); + for element in elements { + let part = self.set(element)?; + written.extend(part.iter().cloned()); + parts.push(part); + } + // `ROLLUP (a, b)`: (a, b), (a), (); `CUBE (a, b)`: every + // subset of the elements. + let n = parts.len(); + if !rollup && n >= usize::BITS as usize - 1 + || !rollup && (1_usize << n) > MAX_GROUPING_SETS + { + return Err(unsupported( + &format!("more than {MAX_GROUPING_SETS} grouping sets"), + item.span, + )); + } + let chosen: Vec> = if rollup { + (0..=n).rev().map(|k| (0..k).collect()).collect() + } else { + (0..1_usize << n) + .rev() + .map(|m| (0..n).filter(|i| m >> (n - 1 - i) & 1 == 1).collect()) + .collect() + }; + chosen + .into_iter() + .map(|c| c.into_iter().flat_map(|i| parts[i].iter().cloned()).collect()) + .collect() + } + }) + } + + /// The expressions of one grouping set. + fn set(&mut self, s: &GroupingSet) -> R>> { + s.exprs.iter().map(|e| Ok(arc(self.expr(e)?))).collect() + } + fn from(&mut self, f: &FromClause) -> R { if f.pivot.is_some() || f.unpivot.is_some() { return Err(unsupported("PIVOT and UNPIVOT", f.span)); From 7bb422483cacecf444a1a884e5a4e967cae8dfc6 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:59:20 +0200 Subject: [PATCH 02/16] feat(sql): ROLLUP, CUBE, GROUPING SETS, grouping() and grouping_id() via Expand Spark's ResolveGroupingAnalytics shape: an Expand plan node (plan data, streamed one copy per projection) feeds an ordinary aggregate. Adds the DataFrame rollup/cube/grouping_sets, Spark golden tests, optimizer cases and commutative semantic keys. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/lib.rs | 9 +- crates/krishiv-exec/src/operators/expand.rs | 13 +- crates/krishiv-exec/src/operators/tests.rs | 124 ++++++ .../krishiv-exec/src/physical/instantiate.rs | 4 +- crates/krishiv-exec/src/physical/mod.rs | 4 +- crates/krishiv-exec/src/planner/build.rs | 5 +- crates/krishiv-optimizer/src/prune.rs | 5 +- crates/krishiv-optimizer/tests/it/rules.rs | 39 ++ .../it__rules__expand_cube_pruned.snap | 7 + .../it__rules__expand_duplicate_sets.snap | 7 + .../it__rules__expand_grouping_sets.snap | 8 + .../snapshots/it__rules__expand_rollup.snap | 7 + .../tests/it/support/eval.rs | 24 ++ crates/krishiv-plan/src/analyzer/analytics.rs | 127 ++++-- crates/krishiv-plan/src/analyzer/exprs.rs | 38 +- crates/krishiv-plan/src/analyzer/ops.rs | 12 +- crates/krishiv-plan/src/tests/analyzer.rs | 4 +- crates/krishiv-plan/src/tests/inventory.rs | 9 +- crates/krishiv-plan/src/tests/props.rs | 17 +- .../tests/grouping_analytics.rs | 326 +++++++++++++++ .../tests/grouping_analytics_api.rs | 381 ++++++++++++++++++ .../krishiv-session/tests/scratch_rollup.rs | 84 ---- crates/krishiv-sql/src/lower.rs | 22 +- ...cts_lower_or_are_named_as_unsupported.snap | 4 +- 24 files changed, 1126 insertions(+), 154 deletions(-) create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap create mode 100644 crates/krishiv-session/tests/grouping_analytics.rs create mode 100644 crates/krishiv-session/tests/grouping_analytics_api.rs delete mode 100644 crates/krishiv-session/tests/scratch_rollup.rs diff --git a/crates/krishiv-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index b9081da93..d189d9e1c 100644 --- a/crates/krishiv-exec/src/lib.rs +++ b/crates/krishiv-exec/src/lib.rs @@ -113,10 +113,11 @@ pub use operator::{ FlushOutput, GlobalState, LocalState, Operator, OperatorOutput, Sink, SinkStatus, Source, }; pub use operators::{ - DEFAULT_MAX_MERGE_FAN_IN, Expand, NullOf, FileScan, FileScanMetrics, FileScanSource, FileSelections, - FileWrite, FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, MemoryScan, - Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, RowStageSettings, RuntimeFilters, - ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, SortMetrics, TopN, UnionAlign, + DEFAULT_MAX_MERGE_FAN_IN, Expand, FileScan, FileScanMetrics, FileScanSource, FileSelections, + FileWrite, FileWriteSink, Filter, FilterPushdown, FnPredicate, Limit, METADATA_COLUMN, + MemoryScan, NullOf, Predicate, Project, ProjectItem, PushedFilter, ReplicateRows, + RowStageSettings, RuntimeFilters, ScalarExpr, ScanPushdown, ScanPushdowns, Sort, SortKey, + SortMetrics, TopN, UnionAlign, }; pub use pipeline::{ExecNode, PipelineBuilder}; pub use runner::{QueryRunner, Scheduler}; diff --git a/crates/krishiv-exec/src/operators/expand.rs b/crates/krishiv-exec/src/operators/expand.rs index 5f0a5c233..aedb7a718 100644 --- a/crates/krishiv-exec/src/operators/expand.rs +++ b/crates/krishiv-exec/src/operators/expand.rs @@ -39,14 +39,11 @@ impl Expand { for p in &projects { let (a, b) = (p.schema(), &schema); let same = a.fields().len() == b.fields().len() - && a.fields() - .iter() - .zip(b.fields()) - .all(|(x, y)| { - x.name() == y.name() - && x.data_type() == y.data_type() - && x.is_nullable() == y.is_nullable() - }); + && a.fields().iter().zip(b.fields()).all(|(x, y)| { + x.name() == y.name() + && x.data_type() == y.data_type() + && x.is_nullable() == y.is_nullable() + }); if !same { return Err(internal(format!( "the projections of an expansion differ: {:?} and {:?}", diff --git a/crates/krishiv-exec/src/operators/tests.rs b/crates/krishiv-exec/src/operators/tests.rs index 2f344affe..d628fb75a 100644 --- a/crates/krishiv-exec/src/operators/tests.rs +++ b/crates/krishiv-exec/src/operators/tests.rs @@ -841,3 +841,127 @@ fn replicate_rows_keeps_weights_and_rejects_non_bigint_counts() { .unwrap_or_else(|| unreachable!("an error")); assert_eq!(err.error_class(), ErrorClass::InternalError); } + +/// A projection of `schema()`: `a` (or nulls of it), then `s` (or nulls). +fn expand_projection(null_a: bool, null_s: bool) -> Arc { + let item = |index: usize, name: &str, null: bool| { + if null { + let nulls = NullOf::new(&schema(), index).unwrap_or_else(|e| unreachable!("{e}")); + ProjectItem::expr(name, Arc::new(nulls)) + } else { + ProjectItem::column(index) + } + }; + Project::try_new(&schema(), vec![item(0, "a", null_a), item(1, "s", null_s)]) + .unwrap_or_else(|e| unreachable!("{e}")) +} + +#[test] +fn expand_streams_one_copy_per_projection_and_keeps_weights() { + let (_root, ctx) = ctx(); + let op = Expand::try_new(vec![ + expand_projection(false, false), + expand_projection(false, true), + expand_projection(true, true), + ]) + .unwrap_or_else(|e| unreachable!("{e}")); + assert_eq!(op.name(), "Expand"); + assert_eq!(op.schema(), schema()); + let mut local = op.local_state(&ctx).unwrap_or_else(|e| unreachable!("{e}")); + let input = weighted(&[1, 2], &[3, -4]); + // One copy per call: `HaveMoreOutput` until the last projection. + let mut copies = Vec::new(); + for _ in 0..2 { + let OperatorOutput::HaveMoreOutput(b) = op + .execute(&input, &mut local) + .unwrap_or_else(|e| unreachable!("{e}")) + else { + unreachable!("a copy and more to come") + }; + copies.push(b); + } + copies.extend(out_batch( + op.execute(&input, &mut local) + .unwrap_or_else(|e| unreachable!("{e}")), + )); + assert_eq!(copies.len(), 3); + assert_eq!(col_a(&copies[0]), [1, 2].map(Some)); + assert_eq!(col_a(&copies[1]), [1, 2].map(Some)); + assert_eq!(col_a(&copies[2]), [None, None]); + let strings = |b: &Batch| -> Vec> { + b.data() + .column(1) + .as_string::() + .iter() + .map(|v| v.map(str::to_owned)) + .collect() + }; + assert_eq!( + strings(&copies[0]), + [Some("s1".to_owned()), Some("s2".to_owned())] + ); + assert_eq!(strings(&copies[1]), [None, None]); + assert!(copies.iter().all(|b| b.weights() == input.weights())); + // The cursor starts again for the next input; an empty batch is skipped. + assert!( + out_batch( + op.execute(&ints(&[]), &mut local) + .unwrap_or_else(|e| unreachable!("{e}")) + ) + .is_none() + ); + assert!(matches!( + op.execute(&ints(&[7]), &mut local), + Ok(OperatorOutput::HaveMoreOutput(_)) + )); +} + +#[test] +fn expand_rejects_bad_plans() { + let class = |r: Result, KrishivError>| { + r.err() + .unwrap_or_else(|| unreachable!("an error")) + .error_class() + }; + assert_eq!(class(Expand::try_new(vec![])), ErrorClass::InternalError); + // A projection that outputs other columns. + let narrow = Project::try_new(&schema(), vec![ProjectItem::column(0)]) + .unwrap_or_else(|e| unreachable!("{e}")); + assert_eq!( + class(Expand::try_new(vec![ + expand_projection(false, false), + narrow + ])), + ErrorClass::InternalError + ); + let renamed = Project::try_new( + &schema(), + vec![ProjectItem::aliased(0, "x"), ProjectItem::column(1)], + ) + .unwrap_or_else(|e| unreachable!("{e}")); + assert_eq!( + class(Expand::try_new(vec![ + expand_projection(false, false), + renamed + ])), + ErrorClass::InternalError + ); + // Nulls of a column the input has not; and of a batch that lost it. + let err = NullOf::new(&schema(), 2) + .err() + .unwrap_or_else(|| unreachable!("an error")); + assert_eq!(err.error_class(), ErrorClass::InternalError); + let nulls = NullOf::new(&schema(), 1).unwrap_or_else(|e| unreachable!("{e}")); + assert_eq!(nulls.data_type(), DataType::Utf8); + assert!(nulls.nullable()); + let one_column = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])), + vec![Arc::new(Int64Array::from(vec![1_i64]))], + ) + .unwrap_or_else(|e| unreachable!("{e}")); + let err = nulls + .evaluate(&one_column) + .err() + .unwrap_or_else(|| unreachable!("an error")); + assert_eq!(err.error_class(), ErrorClass::InternalError); +} diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index 728028a00..58d5971fe 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -30,8 +30,8 @@ use crate::join::{ }; use crate::operator::{Operator, Sink, Source}; use crate::operators::{ - Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project, ProjectItem, - PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, + Expand, FileWrite, FileWriteSink, Filter, FilterPushdown, NullOf, Predicate, Project, + ProjectItem, PushedFilter, ReplicateRows, RuntimeFilters, ScanPushdowns, Sort, SortKey, TopN, }; use crate::pipeline::{ExecNode, PipelineBuilder}; use crate::planner::TableProviderResolver; diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index 2ca8958d9..9d4833bf2 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -50,8 +50,8 @@ pub use data::{PlanData, assert_plan_data}; pub(crate) use instantiate::node_ids; pub use instantiate::{SpillScopes, WorkerContext, instantiate}; pub(crate) use instantiate::{ - Unbuilt, Values, aggregate_node, aggregate_spec, filter_op, hash_join, nested_loop_join, - expand_op, project_op, replicate_op, scan_filter, sort_op, top_n_op, + Unbuilt, Values, aggregate_node, aggregate_spec, expand_op, filter_op, hash_join, + nested_loop_join, project_op, replicate_op, scan_filter, sort_op, top_n_op, }; use crate::aggregate::{AggregateExpr, AggregateOutput}; diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index e9fcc721e..a004d4315 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -775,7 +775,10 @@ impl<'a> Planner<'a> { prepared, op.schema(), layout, - ("Expand", format!("{} projections [{text}]", projections.len())), + ( + "Expand", + format!("{} projections [{text}]", projections.len()), + ), |input, layout, subqueries| PhysNode::Expand { input, layout, diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index 7fb5f23ab..332c99b11 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -226,7 +226,10 @@ impl Rewrite<'_> { let node = PlanNode::Expand { input: Arc::clone(input), projections: projections.iter().map(pick).collect(), - output: kept.iter().filter_map(|&j| output.get(j).cloned()).collect(), + output: kept + .iter() + .filter_map(|&j| output.get(j).cloned()) + .collect(), }; replace(&rebuilt, node) } diff --git a/crates/krishiv-optimizer/tests/it/rules.rs b/crates/krishiv-optimizer/tests/it/rules.rs index b69d424bf..f797e7a46 100644 --- a/crates/krishiv-optimizer/tests/it/rules.rs +++ b/crates/krishiv-optimizer/tests/it/rules.rs @@ -398,6 +398,23 @@ const CASES: &[(&str, &str)] = &[ "prune_exists", "SELECT a FROM t WHERE EXISTS (SELECT * FROM u WHERE u.a = t.a)", ), + // Grouping sets: the expansion keeps what the aggregate reads. + ( + "expand_rollup", + "SELECT a, sum(c) FROM t GROUP BY ROLLUP(a, b)", + ), + ( + "expand_cube_pruned", + "SELECT count(*) FROM t GROUP BY CUBE(a, b)", + ), + ( + "expand_grouping_sets", + "SELECT a, grouping(a), grouping_id(a, b), sum(c) FROM t WHERE c > 1 GROUP BY GROUPING SETS ((a), (b), ())", + ), + ( + "expand_duplicate_sets", + "SELECT a, sum(c) FROM t GROUP BY GROUPING SETS ((a), (a))", + ), ]; #[test] @@ -443,6 +460,28 @@ fn rules() { } } +#[test] +fn an_expansion_has_a_copy_of_its_rows_per_projection() { + let ctx = support::ctx(); + let rows = |sql: &str| { + let plan = support::resolve(sql, &ctx); + let optimized = optimize(&plan).unwrap(); + let mut node = Arc::clone(optimized.root()); + while !matches!(node.node, PlanNode::Expand { .. }) { + let mut next = None; + node.for_each_child(&mut |c| next = Some(Arc::clone(c))); + node = next.expect("an expansion below"); + } + krishiv_optimizer::estimated_rows(&node, &krishiv_optimizer::NoStatistics) + }; + let one_set = rows("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))"); + let four_sets = rows("SELECT count(*) FROM t GROUP BY CUBE(a, b)"); + assert!( + (four_sets / one_set - 2.0).abs() < 1e-9, + "{four_sets} vs {one_set}" + ); +} + #[test] fn top_n_is_detected() { let ctx = support::ctx(); diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap new file mode 100644 index 000000000..24e2e4dc7 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_cube_pruned.snap @@ -0,0 +1,7 @@ +--- +source: crates/krishiv-optimizer/tests/it/rules.rs +expression: "SELECT count(*) FROM t GROUP BY CUBE(a, b)" +--- +Aggregate [a#5, b#6, spark_grouping_id#3], [count(1) AS count(1)#4] ++- Expand [[a#0, b#1, 0L], [a#0, CAST(NULL AS STRING), 1L], [CAST(NULL AS INT), b#1, 2L], [CAST(NULL AS INT), CAST(NULL AS STRING), 3L]], [a#5, b#6, spark_grouping_id#3] + +- Scan spark_catalog.default.t [a#0, b#1] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap new file mode 100644 index 000000000..d5016c878 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_duplicate_sets.snap @@ -0,0 +1,7 @@ +--- +source: crates/krishiv-optimizer/tests/it/rules.rs +expression: "SELECT a, sum(c) FROM t GROUP BY GROUPING SETS ((a), (a))" +--- +Aggregate [a#5, spark_grouping_id#3, _gen_grouping_pos#6], [a#5, sum(c#2) AS sum(c)#4] ++- Expand [[c#2, a#0, 0L, 0], [c#2, a#0, 0L, 1]], [c#2, a#5, spark_grouping_id#3, _gen_grouping_pos#6] + +- Scan spark_catalog.default.t [a#0, c#2] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap new file mode 100644 index 000000000..b3333bb86 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_grouping_sets.snap @@ -0,0 +1,8 @@ +--- +source: crates/krishiv-optimizer/tests/it/rules.rs +expression: "SELECT a, grouping(a), grouping_id(a, b), sum(c) FROM t WHERE c > 1 GROUP BY GROUPING SETS ((a), (b), ())" +--- +Aggregate [a#7, b#8, spark_grouping_id#3], [a#7, CAST(((grouping__id#3 div 2L) % 2L) AS TINYINT) AS grouping(a)#4, grouping__id#3 AS grouping_id(a, b)#5, sum(c#2) AS sum(c)#6] ++- Expand [[c#2, a#0, CAST(NULL AS STRING), 1L], [c#2, CAST(NULL AS INT), b#1, 2L], [c#2, CAST(NULL AS INT), CAST(NULL AS STRING), 3L]], [c#2, a#7, b#8, spark_grouping_id#3] + +- Filter (c#2 > 1L) + +- Scan spark_catalog.default.t [a#0, b#1, c#2] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap new file mode 100644 index 000000000..9c75b03d3 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__rules__expand_rollup.snap @@ -0,0 +1,7 @@ +--- +source: crates/krishiv-optimizer/tests/it/rules.rs +expression: "SELECT a, sum(c) FROM t GROUP BY ROLLUP(a, b)" +--- +Aggregate [a#5, b#6, spark_grouping_id#3], [a#5, sum(c#2) AS sum(c)#4] ++- Expand [[c#2, a#0, b#1, 0L], [c#2, a#0, CAST(NULL AS STRING), 1L], [c#2, CAST(NULL AS INT), CAST(NULL AS STRING), 3L]], [c#2, a#5, b#6, spark_grouping_id#3] + +- Scan spark_catalog.default.t [a#0, b#1, c#2] diff --git a/crates/krishiv-optimizer/tests/it/support/eval.rs b/crates/krishiv-optimizer/tests/it/support/eval.rs index bc73a19b3..e0d75ce2b 100644 --- a/crates/krishiv-optimizer/tests/it/support/eval.rs +++ b/crates/krishiv-optimizer/tests/it/support/eval.rs @@ -177,6 +177,28 @@ impl Interp<'_> { rows, } } + PlanNode::Expand { + input, + projections, + output, + } => { + let rel = self.plan(input, outer); + let ix = index(&rel.cols); + let rows = rel + .rows + .iter() + .flat_map(|r| { + projections + .iter() + .map(|row| row.iter().map(|e| self.eval(e, &ix, r, outer)).collect()) + .collect::>() + }) + .collect(); + Rel { + cols: output.iter().map(|c| c.id).collect(), + rows, + } + } PlanNode::Filter { input, predicate } => { let mut preds = Vec::new(); conjuncts(predicate, &mut preds); @@ -746,6 +768,8 @@ fn apply(name: &str, a: &[V]) -> V { "*" => num(|x, y| Some(x * y / SCALE)), "/" => num(|x, y| (y != 0).then(|| x * SCALE / y)), "%" => num(|x, y| (y != 0).then(|| x % y)), + // Integral division: the quotient truncated to a whole number. + "div" => num(|x, y| (y != 0).then(|| x / y * SCALE)), "negative" => match &a[0] { V::Num(x) => V::Num(-x), _ => V::Null, diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs index c6d34c79e..8173c6063 100644 --- a/crates/krishiv-plan/src/analyzer/analytics.rs +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -47,6 +47,10 @@ use crate::{ /// `BIGINT` with one bit each. const MAX_GROUPING: usize = 63; +/// The most expressions of a `CUBE`: 2^16 grouping sets (the SQL front end +/// bounds `CUBE (..)` items the same way). +const MAX_CUBE_EXPRS: usize = 16; + /// `grouping` or `grouping_id`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum GroupingFn { @@ -104,6 +108,8 @@ pub(super) struct Analytics { gid: Option, /// The expansion; `None` for a grand total. pub(super) expand: Option, + /// Whether the expansion numbers its sets (some two are equal). + numbered: bool, } impl Analytics { @@ -117,6 +123,17 @@ impl Analytics { .collect() } + /// `HAVING` and `ORDER BY` over numbered sets: Spark reads the last + /// grouping key as the grouping id there, and that key is the set's + /// number, so a `grouping` or `grouping_id` call is + /// `UNSUPPORTED_GROUPING_EXPRESSION`. + pub(super) fn check_late_call(&self, e: &Arc) -> Result<(), AnalysisError> { + if !self.numbered { + return Ok(()); + } + reject_grouping_calls(std::slice::from_ref(e)) + } + /// The position of `e` among the grouping expressions: a copy of one, or /// an expression equal to one. fn position(&self, e: &Expr, memo: &mut HashMap) -> Option { @@ -130,6 +147,38 @@ impl Analytics { } } +/// `UNSUPPORTED_GROUPING_EXPRESSION` if any of `exprs` calls `grouping` or +/// `grouping_id`. Spark has no source position for it. +pub(super) fn reject_grouping_calls(exprs: &[Arc]) -> Result<(), AnalysisError> { + let mut seen = std::collections::HashSet::new(); + let mut stack: Vec> = exprs.to_vec(); + while let Some(x) = stack.pop() { + if !seen.insert(Arc::as_ptr(&x) as usize) { + continue; + } + if matches!(&x.node, ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some()) { + return Err(AnalysisError::new( + UnsupportedGroupingExpression {}, + Origin::UNKNOWN, + )); + } + x.for_each_child(&mut |c| stack.push(Arc::clone(c))); + } + Ok(()) +} + +/// A resolved call `e`, stated not nullable. +fn not_null(e: &Arc) -> Arc { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) => { + let mut f = f.clone(); + f.nullable = false; + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function(f))).with_origin(e.origin)) + } + _ => Arc::clone(e), + } +} + fn literal(l: Literal) -> Arc { Arc::new(Expr::new(ExprNode::Literal(l))) } @@ -313,17 +362,25 @@ impl Resolver<'_> { aggregates: &[Arc], env: &Arc, ) -> Result { + // A cube has a set per subset of its expressions: bound it before + // building them (`WITH CUBE` and `DataFrame.cube` arrive here). + if let Grouping::Cube(v) = grouping + && v.len() > MAX_CUBE_EXPRS + { + return Err(AnalysisError::new( + GroupingSizeLimitExceeded { + max_size: (1_u64 << MAX_CUBE_EXPRS).to_string(), + }, + p.origin, + )); + } let c = self.child(input, env)?; let (flat, sets) = shape(grouping); self.check_star_ordinals(&flat, aggregates)?; // A single empty set over nothing is a grand total. let grand = flat.is_empty() && sets.len() == 1; - let gid_col = (!grand).then(|| { - self.attr( - "spark_grouping_id", - &arg_type(DataType::Long, false), - ) - }); + let gid_col = + (!grand).then(|| self.attr("spark_grouping_id", &arg_type(DataType::Long, false))); let gid = gid_col.as_ref().map(|col| Attr { col: BoundColumn { name: "grouping__id".to_owned(), @@ -375,9 +432,17 @@ impl Resolver<'_> { exprs: Vec::new(), gid: None, expand: None, + numbered: false, }; let items = self.rewrite_items(items, &an, env)?; - return self.build_aggregate(p.origin, &c, Vec::new(), &items, true, Some(Arc::new(an))); + return self.build_aggregate( + p.origin, + &c, + Vec::new(), + &items, + true, + Some(Arc::new(an)), + ); } let Some(gid_col) = gid_col else { return Err(unsupported("grouping sets without an id", p.origin)); @@ -385,7 +450,11 @@ impl Resolver<'_> { let numbered = (0..sets.len()).any(|i| sets[..i].contains(&sets[i])); // The expansion's input: the child's columns, and the grouping // expressions that are not one of them. - let mut passed: Vec> = c.out.iter().map(|a| col_expr(&a.col, Origin::UNKNOWN)).collect(); + let mut passed: Vec> = c + .out + .iter() + .map(|a| col_expr(&a.col, Origin::UNKNOWN)) + .collect(); let mut passed_out: Vec = c.out.to_vec(); let mut value: Vec = Vec::with_capacity(n); let mut copies: Vec = Vec::with_capacity(n); @@ -406,10 +475,12 @@ impl Resolver<'_> { Some(a) => a.col.clone(), None => { let v = self.attr(&name, &b.ty); - passed.push(Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { - expr: Arc::clone(&b.expr), - output: v.clone(), - })))); + passed.push(Arc::new(Expr::new(ExprNode::Resolved( + ResolvedExpr::Alias { + expr: Arc::clone(&b.expr), + output: v.clone(), + }, + )))); passed_out.push(Attr { col: v.clone(), qual: Arc::from(Vec::new()), @@ -425,7 +496,9 @@ impl Resolver<'_> { key: keys_of[j], attr: Attr { col: copy, - qual: column.as_ref().map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), + qual: column + .as_ref() + .map_or_else(|| Arc::from(Vec::new()), |a| Arc::clone(&a.qual)), source: column.and_then(|a| a.source), }, }); @@ -437,7 +510,8 @@ impl Resolver<'_> { } else { c.clone() }; - let pos_col = numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); + let pos_col = + numbered.then(|| self.attr("_gen_grouping_pos", &arg_type(DataType::Integer, false))); let mut output: Vec = passed_out.iter().map(|a| a.col.clone()).collect(); output.extend(copies.iter().map(|g| g.attr.col.clone())); output.push(gid_col.clone()); @@ -459,9 +533,9 @@ impl Resolver<'_> { } row.push(literal(Literal::Long(mask))); if pos_col.is_some() { - row.push(literal(Literal::Integer( - i32::try_from(si).map_err(|_| unsupported("that many grouping sets", p.origin))?, - ))); + row.push(literal(Literal::Integer(i32::try_from(si).map_err( + |_| unsupported("that many grouping sets", p.origin), + )?))); } projections.push(row); } @@ -492,6 +566,7 @@ impl Resolver<'_> { let an = Analytics { exprs: copies, gid, + numbered, expand: Some(ExpandPlan { plan: Arc::clone(&expand.plan), ids: expand.ids, @@ -619,8 +694,7 @@ impl Resolver<'_> { Some(g) => col_expr(&g.col, Origin::UNKNOWN), None => literal(Literal::Long(0)), }; - let positions: Vec> = - f.args.iter().map(|a| an.position(a, memo)).collect(); + let positions: Vec> = f.args.iter().map(|a| an.position(a, memo)).collect(); match GroupingFn::from_name(&f.name) { Some(GroupingFn::GroupingId) => { let exact = positions.len() == n @@ -647,7 +721,7 @@ impl Resolver<'_> { Ok(id) } _ => { - let (Some(arg), Some(Some(j))) = (f.args.first(), positions.first()) else { + let (Some(_), Some(Some(j))) = (f.args.first(), positions.first()) else { return Err(AnalysisError::new( GroupingColumnMismatch { grouping: f.args.first().map(|a| with_id(a)).unwrap_or_default(), @@ -661,18 +735,20 @@ impl Resolver<'_> { Origin::UNKNOWN, )); }; - let _ = arg; // Bit `n - 1 - j` of the id, as a `TINYINT`. let shift = n - 1 - j; + // The divisors are literals that are not 0, so the id's + // bit is never null: `grouping` is not nullable (Spark). let high = if shift == 0 { id } else { - self.call( + let high = self.call( "div", vec![id, literal(Literal::Long(1_i64 << shift))], env, Origin::UNKNOWN, - )? + )?; + not_null(&high) }; let bit = self.call( "%", @@ -680,12 +756,15 @@ impl Resolver<'_> { env, Origin::UNKNOWN, )?; + let bit = not_null(&bit); let cast = Arc::new(Expr::new(ExprNode::Cast { expr: bit, to: DataType::Byte, mode: CastMode::Default, })); - Ok(self.bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))?.expr) + Ok(self + .bind(&cast, &Self::cx(&[], &[], env, "Aggregate"))? + .expr) } } } diff --git a/crates/krishiv-plan/src/analyzer/exprs.rs b/crates/krishiv-plan/src/analyzer/exprs.rs index b16b69299..3bf207a15 100644 --- a/crates/krishiv-plan/src/analyzer/exprs.rs +++ b/crates/krishiv-plan/src/analyzer/exprs.rs @@ -272,6 +272,8 @@ impl Resolver<'_> { stack.push(Step::Enter(rewritten, in_agg)); continue; } + // In `ORDER BY` Spark rejects the call before it resolves + // its argument (`HAVING` reports the argument first). let inner = in_agg || self.is_aggregate(&x); stack.push(Step::Exit(Arc::clone(&x), in_agg)); let mut kids = Vec::new(); @@ -1525,6 +1527,29 @@ fn pretty_literal(v: &Literal) -> String { } } +/// A binary operator whose operands commute (Spark's `Add`, `Multiply`, +/// `And`, `Or`, `EqualTo` and `EqualNullSafe`). +fn is_commutative(node: &ExprNode) -> bool { + let ExprNode::Resolved(ResolvedExpr::Function(f)) = node else { + return false; + }; + let op = match &f.id { + FunctionId::Builtin(b) => Some(*b), + FunctionId::Registry { name, .. } => Builtin::from_name(name), + }; + matches!( + op, + Some( + Builtin::Add + | Builtin::Multiply + | Builtin::And + | Builtin::Or + | Builtin::Eq + | Builtin::EqNullSafe + ) + ) +} + /// A semantic key: equal for expressions that compute the same value, /// whatever their origins, column spellings or (top-level) alias. pub(super) fn semantic_key(e: &Expr, memo: &mut HashMap) -> u128 { @@ -1557,8 +1582,17 @@ pub(super) fn semantic_key(e: &Expr, memo: &mut HashMap) -> u128 { other => format!("u{other:?}"), } .into_bytes(); - for c in kids { - b.extend_from_slice(&memo.get(&addr(c)).copied().unwrap_or(0).to_le_bytes()); + let mut hashes: Vec = kids + .iter() + .map(|c| memo.get(&addr(*c)).copied().unwrap_or(0)) + .collect(); + // Spark canonicalizes the operands of a commutative operator into + // one order, so `a + b` and `b + a` are the same expression. + if hashes.len() == 2 && is_commutative(&x.node) { + hashes.sort_unstable(); + } + for h in hashes { + b.extend_from_slice(&h.to_le_bytes()); } memo.insert(addr(x), XxHash3_128::oneshot(&b)); } diff --git a/crates/krishiv-plan/src/analyzer/ops.rs b/crates/krishiv-plan/src/analyzer/ops.rs index 8238aec0e..aff4ece89 100644 --- a/crates/krishiv-plan/src/analyzer/ops.rs +++ b/crates/krishiv-plan/src/analyzer/ops.rs @@ -684,6 +684,7 @@ impl Resolver<'_> { } if let Some(an) = &info.analytics { for b in &mut bounds { + an.check_late_call(&b.expr)?; b.expr = self.resolve_grouping(&b.expr, an, env)?; } } @@ -960,7 +961,11 @@ impl Resolver<'_> { let (keys, scope): (Vec, Vec) = match &analytics { Some(a) => ( a.keys(), - c.out.iter().chain(a.exprs.iter().map(|g| &g.attr)).cloned().collect(), + c.out + .iter() + .chain(a.exprs.iter().map(|g| &g.attr)) + .cloned() + .collect(), ), None => ( group @@ -1400,6 +1405,11 @@ impl Resolver<'_> { .map(|(o, expr)| SortOrder { expr, ..o.clone() }) .collect() }; + // Over no grouping sets, `ORDER BY grouping(..)` is rejected before + // its argument resolves (`HAVING` reports the argument first). + if c.agg.as_ref().is_none_or(|a| a.analytics.is_none()) { + super::analytics::reject_grouping_calls(&exprs)?; + } let (bounds, below, extended) = if c.agg.is_some() { self.over_agg(&c, &exprs, env, "Sort")? } else { diff --git a/crates/krishiv-plan/src/tests/analyzer.rs b/crates/krishiv-plan/src/tests/analyzer.rs index c05b3ba6b..647bec823 100644 --- a/crates/krishiv-plan/src/tests/analyzer.rs +++ b/crates/krishiv-plan/src/tests/analyzer.rs @@ -26,8 +26,8 @@ use crate::analyzer::{ }; use crate::walk::{Node, visit_distinct}; use crate::{ - ColumnRef, Cte, Expr, ExprNode, Grouping, JoinCondition, JoinKind, Literal, LogicalPlan, - MAX_DEPTH, Origin, Plan, PlanId, PlanNode, ReadSource, SourceId, SqlArgs, Star, TableVersion, + ColumnRef, Cte, Expr, ExprNode, JoinCondition, JoinKind, Literal, LogicalPlan, MAX_DEPTH, + Origin, Plan, PlanId, PlanNode, ReadSource, SourceId, SqlArgs, Star, TableVersion, }; fn goldens() -> Value { diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index fcd401999..862da651f 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -12,8 +12,8 @@ use krishiv_types::DataType; use super::dsl::{c, int}; use super::tpch::QUERIES; use crate::{ - Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, PlanNode, - ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, + Builtin, ColumnRef, Expr, ExprNode, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanId, + PlanNode, ReadSource, ResolvedExpr, SortOrder, SqlArgs, Star, }; /// The node kinds a plan uses, named as in the inventory. No wildcard arm: @@ -170,7 +170,10 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { ("filter", df(1).filter(c("a"))), ("where", df(1).filter(c("a"))), ("groupBy.agg", df(1).aggregate(vec![c("a")], vec![c("a")])), - ("rollup.agg", grouped(Grouping::Rollup(vec![Arc::new(c("a"))]))), + ( + "rollup.agg", + grouped(Grouping::Rollup(vec![Arc::new(c("a"))])), + ), ("cube.agg", grouped(Grouping::Cube(vec![Arc::new(c("a"))]))), ( "join", diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index ec013bb5d..24bfa55ad 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -426,7 +426,7 @@ impl Gen<'_> { }, 21 => PlanNode::Expand { input: self.input(d), - projections: self.many(1, 2, |g| g.exprs()), + projections: self.many(1, 2, Gen::exprs), output: self.many(0, 2, Gen::bound), }, _ => PlanNode::WithCte { @@ -824,12 +824,15 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { output, } => { site!(w, projections.push(Vec::new())); - site!(w, output.push(output.first().cloned().unwrap_or(BoundColumn { - id: AttrId(0), - name: "x".to_owned(), - data_type: DataType::Integer, - nullable: false, - }))); + site!( + w, + output.push(output.first().cloned().unwrap_or(BoundColumn { + id: AttrId(0), + name: "x".to_owned(), + data_type: DataType::Integer, + nullable: false, + })) + ); projections.iter_mut().any(|p| m_exprs(p, w)) || output.iter_mut().any(|c| m_bound(c, w)) || m_plan(input, w) diff --git a/crates/krishiv-session/tests/grouping_analytics.rs b/crates/krishiv-session/tests/grouping_analytics.rs new file mode 100644 index 000000000..a544dce4a --- /dev/null +++ b/crates/krishiv-session/tests/grouping_analytics.rs @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Spark 4.2 `SQLQueryTestSuite` goldens for grouping analytics (`ROLLUP`, `CUBE`, +//! `GROUPING SETS`, `grouping`, `grouping_id`): each statement of the vendored +//! `group-analytics` and `grouping_set` result files runs through a session and +//! its schema, rows (compared as sets of rows) or error class must match. +//! +//! The session has no `CREATE VIEW` or inline `VALUES` yet, so the driver turns +//! each literal `VALUES` table into a registered table of the same name. +#![forbid(unsafe_code)] + +use std::path::Path; +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Float64Type, Int8Type, Int16Type, Int32Type, Int64Type}; +use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use futures::executor::block_on; +use krishiv_common::error::KrishivError; +use krishiv_session::Session; + +type R = Result>; + +/// One golden statement: the query and its expected schema and output. +struct Case { + query: String, + schema: String, + output: String, +} + +fn load(file: &str) -> R> { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../tests/spark-golden/v4.2.0/results") + .join(file); + let text = std::fs::read_to_string(&path)?; + let cases: Vec = text + .split("-- !query\n") + .skip(1) + .filter_map(|block| { + let (query, rest) = block.split_once("\n-- !query schema\n")?; + let (schema, output) = rest.split_once("\n-- !query output\n")?; + Some(Case { + query: query.trim().to_owned(), + schema: schema.trim().to_owned(), + output: output.trim_end().to_owned(), + }) + }) + .collect(); + if cases.is_empty() { + return Err(format!("no statements in {file}").into()); + } + Ok(cases) +} + +fn spark_type(t: &DataType) -> String { + match t { + DataType::Int8 => "tinyint".into(), + DataType::Int16 => "smallint".into(), + DataType::Int32 => "int".into(), + DataType::Int64 => "bigint".into(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => "string".into(), + DataType::Boolean => "boolean".into(), + DataType::Float64 => "double".into(), + other => format!("{other}"), + } +} + +fn cell(c: &dyn Array, row: usize) -> String { + if c.is_null(row) { + return "NULL".into(); + } + match c.data_type() { + DataType::Int8 => c.as_primitive::().value(row).to_string(), + DataType::Int16 => c.as_primitive::().value(row).to_string(), + DataType::Int32 => c.as_primitive::().value(row).to_string(), + DataType::Int64 => c.as_primitive::().value(row).to_string(), + DataType::Utf8 => c.as_string::().value(row).to_owned(), + DataType::Boolean => c.as_boolean().value(row).to_string(), + DataType::Float64 => format!("{:?}", c.as_primitive::().value(row)), + other => format!("<{other}>"), + } +} + +fn render(schema: &Schema, batches: &[RecordBatch]) -> (String, Vec) { + let schema = schema + .fields() + .iter() + .map(|f| format!("{}:{}", f.name(), spark_type(f.data_type()))) + .collect::>() + .join(","); + let rows = batches + .iter() + .flat_map(|b| { + (0..b.num_rows()).map(|r| { + b.columns() + .iter() + .map(|c| cell(c.as_ref(), r)) + .collect::>() + .join("\t") + }) + }) + .collect(); + (format!("struct<{schema}>"), rows) +} + +/// A parsed `VALUES` literal table. +struct Table { + name: String, + columns: Vec, + rows: Vec>>, +} + +fn is_word(b: u8) -> bool { + b.is_ascii_alphanumeric() || b == b'_' +} + +fn skip_ws(bytes: &[u8], i: &mut usize) { + while bytes.get(*i).is_some_and(u8::is_ascii_whitespace) { + *i += 1; + } +} + +/// Find an inline `VALUES (..), (..) [AS] alias[(cols)]` (optionally wrapped in +/// parentheses) in `sql`; return the text with it replaced by the alias, and the table. +fn extract_values(sql: &str) -> Option<(String, Table)> { + let upper = sql.to_ascii_uppercase(); + let bytes = sql.as_bytes(); + let at = upper.match_indices("VALUES").map(|(i, _)| i).find(|&i| { + (i == 0 || bytes.get(i - 1).is_none_or(|&b| !is_word(b))) + && bytes.get(i + 6).is_none_or(|&b| !is_word(b)) + })?; + let before = sql[..at].trim_end(); + let wrapped = before.ends_with('('); + let start = if wrapped { before.len() - 1 } else { at }; + let mut i = at + 6; + let mut rows = Vec::new(); + loop { + skip_ws(bytes, &mut i); + if bytes.get(i) != Some(&b'(') { + return None; + } + let close = i + sql[i..].find(')')?; + rows.push( + sql[i + 1..close] + .split(',') + .map(|v| { + let v = v.trim(); + (!v.eq_ignore_ascii_case("NULL")).then(|| v.to_owned()) + }) + .collect::>(), + ); + i = close + 1; + skip_ws(bytes, &mut i); + if bytes.get(i) == Some(&b',') { + i += 1; + } else { + break; + } + } + if wrapped { + if bytes.get(i) != Some(&b')') { + return None; + } + i += 1; + skip_ws(bytes, &mut i); + } + if upper[i..].starts_with("AS ") { + i += 3; + skip_ws(bytes, &mut i); + } + let name_start = i; + while bytes.get(i).is_some_and(|&b| is_word(b)) { + i += 1; + } + let name = sql[name_start..i].to_owned(); + skip_ws(bytes, &mut i); + let mut columns = Vec::new(); + if bytes.get(i) == Some(&b'(') { + let close = i + sql[i..].find(')')?; + columns = sql[i + 1..close] + .split(',') + .map(|c| c.trim().to_owned()) + .collect(); + i = close + 1; + } + let text = format!("{}{}{}", &sql[..start], name, &sql[i..]); + Some(( + text, + Table { + name, + columns, + rows, + }, + )) +} + +fn register(s: &Session, t: &Table, name: &str) -> R { + let mut fields = Vec::new(); + let mut arrays: Vec = Vec::new(); + for (c, column) in t.columns.iter().enumerate() { + let col: Vec> = t + .rows + .iter() + .map(|r| r.get(c).and_then(Option::as_ref)) + .collect(); + let numeric = col.iter().flatten().all(|v| v.parse::().is_ok()); + let nullable = col.iter().any(Option::is_none); + if numeric { + fields.push(Field::new(column, DataType::Int32, nullable)); + arrays.push(Arc::new(Int32Array::from( + col.iter() + .map(|v| v.and_then(|v| v.parse().ok())) + .collect::>(), + ))); + } else { + fields.push(Field::new(column, DataType::Utf8, nullable)); + arrays.push(Arc::new(StringArray::from( + col.iter() + .map(|v| v.map(|v| v.trim_matches(['"', '\'']))) + .collect::>(), + ))); + } + } + let schema = Arc::new(Schema::new(fields)); + let batch = RecordBatch::try_new(Arc::clone(&schema), arrays)?; + s.register_batches(name, schema, vec![batch])?; + Ok(()) +} + +/// The error class a golden expects, from its JSON error body. +fn expected_class(output: &str) -> Option { + output + .lines() + .find_map(|l| l.trim().strip_prefix("\"errorClass\" : \"")) + .map(|l| l.trim_end_matches(['"', ',']).to_owned()) +} + +/// Run every statement of `file`; return the statements that disagree with Spark. +fn run(file: &str) -> R> { + let mut views: Vec<(String, Table)> = Vec::new(); + let mut bad = Vec::new(); + for case in load(file)? { + let upper = case.query.to_ascii_uppercase(); + if upper.starts_with("CREATE") { + let name = case + .query + .split_whitespace() + .skip_while(|w| !w.eq_ignore_ascii_case("VIEW")) + .nth(1); + let (Some(name), Some((_, table))) = (name, extract_values(&case.query)) else { + bad.push(format!("{}\n unsupported setup", case.query)); + continue; + }; + views.retain(|(n, _)| n != name); + views.push((name.to_owned(), table)); + continue; + } + if let Some(name) = upper.strip_prefix("DROP VIEW ") { + views.retain(|(n, _)| !n.eq_ignore_ascii_case(name.trim())); + continue; + } + let s = Session::new()?; + for (name, t) in &views { + register(&s, t, name)?; + } + let mut query = case.query.clone(); + if let Some((text, table)) = extract_values(&query) { + register(&s, &table, &table.name)?; + query = text; + } + let result = block_on(async { + let df = s.sql(&query).await?; + let schema = df.schema().await?; + Ok::<_, KrishivError>((schema, df.collect().await?)) + }); + if case.schema == "struct<>" { + let class = expected_class(&case.output); + match (result, class) { + (Err(e), Some(c)) if e.error_class().name() == c => {} + (Err(e), c) => bad.push(format!( + "{}\n error class {} expected {c:?}", + case.query, + e.error_class().name() + )), + (Ok(_), c) => bad.push(format!("{}\n expected error {c:?}", case.query)), + } + continue; + } + match result { + Err(e) => bad.push(format!("{}\n error {e}", case.query)), + Ok((schema, batches)) => { + let (schema, mut rows) = render(&schema, &batches); + let mut want: Vec = case.output.lines().map(str::to_owned).collect(); + // Rows tied under `ORDER BY` come in no fixed order. + rows.sort(); + want.sort(); + if schema != case.schema || rows != want { + bad.push(format!( + "{}\n schema {schema} want {}\n rows {rows:?}\n want {want:?}", + case.query, case.schema + )); + } + } + } + } + Ok(bad) +} + +fn check(file: &str) -> R { + let bad = run(file)?; + if bad.is_empty() { + Ok(()) + } else { + Err(format!("{} mismatches:\n{}", bad.len(), bad.join("\n")).into()) + } +} + +#[test] +fn group_analytics_golden() -> R { + check("group-analytics.sql.out") +} + +#[test] +fn grouping_set_golden() -> R { + check("grouping_set.sql.out") +} diff --git a/crates/krishiv-session/tests/grouping_analytics_api.rs b/crates/krishiv-session/tests/grouping_analytics_api.rs new file mode 100644 index 000000000..9c6381437 --- /dev/null +++ b/crates/krishiv-session/tests/grouping_analytics_api.rs @@ -0,0 +1,381 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Grouping analytics through the session: the DataFrame forms (`rollup`, +//! `cube`, `grouping_sets`) against the SQL forms, rolled-up columns' types, +//! `NULL` data against rolled-up `NULL`, and the errors of invalid uses. +//! The Spark goldens are in `grouping_analytics.rs`. +#![forbid(unsafe_code)] + +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Int8Type, Int32Type, Int64Type}; +use arrow_array::{Array, Int32Array, RecordBatch}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use futures::executor::block_on; +use krishiv_common::error::{ErrorClass, KrishivError}; +use krishiv_session::functions::{col, lit, sum}; +use krishiv_session::{DataFrame, Session}; + +type R = Result>; + +/// Tables `t(a, b, v)` (no nulls) and `n(a, v)` (with a real NULL key). +fn session() -> R { + let s = Session::new()?; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2])), + Arc::new(Int32Array::from(vec![1, 2, 1])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + )?; + s.register_batches("t", schema, vec![batch])?; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("v", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(1), None, None])), + Arc::new(Int32Array::from(vec![5, 7, 11])), + ], + )?; + s.register_batches("n", schema, vec![batch])?; + Ok(s) +} + +fn cell(c: &dyn Array, row: usize) -> String { + if c.is_null(row) { + return "NULL".into(); + } + match c.data_type() { + DataType::Int8 => c.as_primitive::().value(row).to_string(), + DataType::Int32 => c.as_primitive::().value(row).to_string(), + DataType::Int64 => c.as_primitive::().value(row).to_string(), + other => format!("<{other}>"), + } +} + +fn row_strings(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|b| { + (0..b.num_rows()).map(|r| { + b.columns() + .iter() + .map(|c| cell(c.as_ref(), r)) + .collect::>() + .join("\t") + }) + }) + .collect() +} + +/// The schema and the rows (tab-separated, sorted) of a DataFrame. +fn run(df: &DataFrame) -> Result<(SchemaRef, Vec), KrishivError> { + block_on(async { + let schema = df.schema().await?; + let mut rows = row_strings(&df.collect().await?); + rows.sort(); + Ok((schema, rows)) + }) +} + +fn sql(s: &Session, query: &str) -> Result<(SchemaRef, Vec), KrishivError> { + run(&block_on(s.sql(query))?) +} + +fn rows(s: &Session, query: &str) -> R> { + Ok(sql(s, query)?.1) +} + +fn sorted(rows: &[&str]) -> Vec { + let mut rows: Vec = rows.iter().map(|r| (*r).to_owned()).collect(); + rows.sort(); + rows +} + +fn error_class(s: &Session, query: &str) -> R { + match sql(s, query) { + Err(e) => Ok(e.error_class()), + Ok((_, rows)) => { + Err(format!("{query} should fail, it returned {} rows", rows.len()).into()) + } + } +} + +#[test] +fn dataframe_rollup_cube_and_grouping_sets_match_their_sql() -> R { + let s = session()?; + let sum_v = || [sum(col("v"))]; + let (schema, rollup) = run(&s.table("t").rollup([col("a"), col("b")]).agg(sum_v()))?; + assert_eq!( + rollup, + sorted(&[ + "1\t1\t10", + "1\t2\t20", + "2\t1\t30", + "1\tNULL\t30", + "2\tNULL\t30", + "NULL\tNULL\t60", + ]) + ); + let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, ["a", "b", "sum(v)"]); + assert_eq!( + rollup, + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY ROLLUP(a, b)")? + ); + + let (_, cube) = run(&s.table("t").cube([col("a"), col("b")]).agg(sum_v()))?; + assert_eq!( + cube, + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY CUBE(a, b)")? + ); + assert_eq!(cube.len(), 8); + + let sets = s + .table("t") + .grouping_sets( + [vec![col("a")], vec![col("b")], vec![]], + [col("a"), col("b")], + ) + .agg(sum_v()); + let (_, sets) = run(&sets)?; + assert_eq!( + sets, + sorted(&[ + "1\tNULL\t30", + "2\tNULL\t30", + "NULL\t1\t40", + "NULL\t2\t20", + "NULL\tNULL\t60", + ]) + ); + assert_eq!( + sets, + rows( + &s, + "SELECT a, b, sum(v) FROM t GROUP BY GROUPING SETS ((a), (b), ())" + )? + ); + Ok(()) +} + +#[test] +fn rolled_up_columns_become_nullable_and_grouping_is_a_tinyint() -> R { + let s = session()?; + let (schema, _) = sql( + &s, + "SELECT a, b, sum(v), grouping(a), grouping_id(a, b) FROM t GROUP BY ROLLUP(a, b)", + )?; + let fields: Vec<(String, DataType, bool)> = schema + .fields() + .iter() + .map(|f| (f.name().clone(), f.data_type().clone(), f.is_nullable())) + .collect(); + assert_eq!( + fields, + [ + ("a".to_owned(), DataType::Int32, true), + ("b".to_owned(), DataType::Int32, true), + ("sum(v)".to_owned(), DataType::Int64, true), + ("grouping(a)".to_owned(), DataType::Int8, false), + ("grouping_id(a, b)".to_owned(), DataType::Int64, false), + ] + ); + Ok(()) +} + +#[test] +fn grouping_tells_a_null_key_from_a_rolled_up_one() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, grouping(a), sum(v) FROM n GROUP BY ROLLUP(a)" + )?, + sorted(&["1\t0\t5", "NULL\t0\t18", "NULL\t1\t23"]) + ); + Ok(()) +} + +#[test] +fn grouping_id_sets_bit_n_minus_1_minus_position() -> R { + let s = session()?; + // `grouping_id(a, b)`: `a` is the high bit. + assert_eq!( + rows( + &s, + "SELECT a, b, grouping_id(a, b) FROM t GROUP BY GROUPING SETS ((a), (b), (a, b))" + )?, + sorted(&[ + "1\tNULL\t1", + "2\tNULL\t1", + "NULL\t1\t2", + "NULL\t2\t2", + "1\t1\t0", + "1\t2\t0", + "2\t1\t0", + ]) + ); + Ok(()) +} + +#[test] +fn mixed_plain_keys_and_rollup() -> R { + let s = session()?; + assert_eq!( + rows(&s, "SELECT a, b, sum(v) FROM t GROUP BY a, ROLLUP(b)")?, + sorted(&[ + "1\t1\t10", + "1\t2\t20", + "2\t1\t30", + "1\tNULL\t30", + "2\tNULL\t30" + ]) + ); + Ok(()) +} + +#[test] +fn operands_of_a_commutative_operator_are_one_grouping_expression() -> R { + let s = session()?; + // `b + a` is `a + b`: the select item reads the grouping copy, and the + // two sets are equal, so each group comes twice. + assert_eq!( + rows( + &s, + "SELECT b + a, sum(v) FROM t GROUP BY GROUPING SETS ((a + b), (b + a))" + )?, + sorted(&["2\t10", "2\t10", "3\t50", "3\t50"]) + ); + Ok(()) +} + +#[test] +fn empty_set_is_a_grand_total_even_over_no_rows() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT count(*), grouping_id() FROM t WHERE v > 100 GROUP BY GROUPING SETS (())" + )?, + ["0\t0"] + ); + assert_eq!( + rows( + &s, + "SELECT a, count(*) FROM t WHERE v > 100 GROUP BY ROLLUP(a)" + )?, + Vec::::new(), + "a rollup over no rows has no grand-total row, as in Spark" + ); + Ok(()) +} + +#[test] +fn having_and_order_by_read_grouping() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, sum(v) FROM t GROUP BY ROLLUP(a) HAVING grouping(a) = 1" + )?, + ["NULL\t60"] + ); + // `ORDER BY` follows the id: the grand total (id 1) comes first, descending. + let df = block_on( + s.sql("SELECT a, sum(v) FROM t GROUP BY ROLLUP(a) ORDER BY grouping_id(a) DESC, a"), + )?; + assert_eq!( + row_strings(&block_on(df.collect())?), + ["NULL\t60", "1\t30", "2\t30"] + ); + Ok(()) +} + +#[test] +fn invalid_uses_have_spark_error_classes() -> R { + let s = session()?; + let unsupported = ErrorClass::UnsupportedGroupingExpression; + // Outside a grouping aggregate, in any clause. + assert_eq!(error_class(&s, "SELECT grouping(a) FROM t")?, unsupported); + assert_eq!( + error_class(&s, "SELECT a FROM t GROUP BY a HAVING grouping(a) = 0")?, + unsupported + ); + assert_eq!( + error_class(&s, "SELECT count(*) FROM t ORDER BY grouping(a)")?, + unsupported + ); + // An argument that is no grouping expression, and a mismatched id list. + assert_eq!( + error_class(&s, "SELECT grouping(v) FROM t GROUP BY ROLLUP(a)")?, + ErrorClass::GroupingColumnMismatch + ); + assert_eq!( + error_class(&s, "SELECT grouping_id(b, a) FROM t GROUP BY ROLLUP(a, b)")?, + ErrorClass::GroupingIdColumnMismatch + ); + assert_eq!( + error_class(&s, "SELECT grouping(a, b) FROM t GROUP BY ROLLUP(a, b)")?, + ErrorClass::WrongNumArgsWithoutSuggestion + ); + // 64 grouping expressions do not fit the 64-bit id. + let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); + let query = format!( + "SELECT count(*) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ); + assert_eq!( + error_class(&s, &query)?, + ErrorClass::GroupingSizeLimitExceeded + ); + // A cube has 2^n sets: 17 expressions are too many, 16 still are not. + let cube = |n: usize| { + let keys: Vec = (0..n).map(|i| format!("a + {i}")).collect(); + format!( + "SELECT count(*) FROM t GROUP BY {} WITH CUBE", + keys.join(", ") + ) + }; + assert_eq!( + error_class(&s, &cube(17))?, + ErrorClass::GroupingSizeLimitExceeded + ); + let keys: Vec<_> = (0..17).map(|i| col("a").gt(lit(i))).collect(); + let df = s.table("t").cube(keys).agg([sum(col("v"))]); + assert_eq!( + run(&df).err().map(|e| e.error_class()), + Some(ErrorClass::GroupingSizeLimitExceeded) + ); + Ok(()) +} + +#[test] +fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { + let s = session()?; + assert_eq!( + rows( + &s, + "SELECT a, sum(v) FROM t GROUP BY GROUPING SETS ((a), (a))" + )?, + sorted(&["1\t30", "1\t30", "2\t30", "2\t30"]) + ); + // Spark reads the set's number where `HAVING` expects the id. + assert_eq!( + error_class( + &s, + "SELECT a FROM t GROUP BY GROUPING SETS ((a), (a)) HAVING grouping(a) = 0" + )?, + ErrorClass::UnsupportedGroupingExpression + ); + Ok(()) +} diff --git a/crates/krishiv-session/tests/scratch_rollup.rs b/crates/krishiv-session/tests/scratch_rollup.rs deleted file mode 100644 index f12c60530..000000000 --- a/crates/krishiv-session/tests/scratch_rollup.rs +++ /dev/null @@ -1,84 +0,0 @@ -//! scratch -#![cfg(test)] -use std::sync::Arc; - -use arrow_array::{Int32Array, RecordBatch, StringArray}; -use arrow_schema::{DataType as T, Field, Schema}; -use futures::executor::block_on; -use krishiv_session::Session; - -#[test] -fn scratch() { - let s = Session::new().unwrap(); - let schema = Arc::new(Schema::new(vec![ - Field::new("a", T::Int32, false), - Field::new("b", T::Int32, false), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(vec![1, 1, 2, 2, 3, 3])), - Arc::new(Int32Array::from(vec![1, 2, 1, 2, 1, 2])), - ], - ) - .unwrap(); - s.register_batches("td", schema, vec![batch]).unwrap(); - let schema = Arc::new(Schema::new(vec![ - Field::new("course", T::Utf8, false), - Field::new("year", T::Int32, false), - Field::new("earnings", T::Int32, false), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(StringArray::from(vec!["dotNET", "Java", "dotNET", "dotNET", "Java"])), - Arc::new(Int32Array::from(vec![2012, 2012, 2012, 2013, 2013])), - Arc::new(Int32Array::from(vec![10000, 20000, 5000, 48000, 30000])), - ], - ) - .unwrap(); - s.register_batches("cs", schema, vec![batch]).unwrap(); - for q in std::fs::read_to_string("/home/krishiv/krishiv-data/tmp/s-rollup/q.sql") - .unwrap() - .split(";\n") - .filter(|q| !q.trim().is_empty()) - { - match block_on(s.sql(q)) { - Err(e) => println!("ERR {q}\n {e}"), - Ok(df) => match block_on(df.collect()) { - Err(e) => println!("RUNERR {q}\n {e}"), - Ok(b) => { - println!("OK {q}"); - for b in &b { - println!("{}", arrow_cast_fmt(b)); - } - } - }, - } - } -} - -fn arrow_cast_fmt(b: &RecordBatch) -> String { - let mut out = format!("{:?}\n", b.schema().fields().iter().map(|f| format!("{}:{}{}", f.name(), f.data_type(), if f.is_nullable() {"?"} else {""})).collect::>()); - for r in 0..b.num_rows() { - let cells: Vec = b - .columns() - .iter() - .map(|c| { - if c.is_null(r) { - "NULL".to_owned() - } else { - let d = arrow_array::cast::AsArray::as_primitive_opt::(c) - .map(|a| a.value(r).to_string()) - .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) - .or_else(|| arrow_array::cast::AsArray::as_primitive_opt::(c).map(|a| a.value(r).to_string())) - .or_else(|| arrow_array::cast::AsArray::as_string_opt::(c).map(|a| a.value(r).to_owned())); - d.unwrap_or_else(|| "?".to_owned()) - } - }) - .collect(); - out.push_str(&cells.join("\t")); - out.push('\n'); - } - out -} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index 8daa1db39..f14c2e530 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -28,9 +28,9 @@ use krishiv_types::DecimalType; use crate::ast::{ Alias, BinaryOp, Expr, ExprKind, FromClause, GroupBy, GroupItem, GroupItemKind, GroupWith, - GroupingSet, JoinConstraint, JoinKind, Literal, - NullTreatment, NumKind, Query, Relation, RelationExt, Select, SetExpr, SetExprKind, - SetOperator, SetQuantifier, SortItem, Span, TableFactor, TableFactorKind, Truth, UnaryOp, + GroupingSet, JoinConstraint, JoinKind, Literal, NullTreatment, NumKind, Query, Relation, + RelationExt, Select, SetExpr, SetExprKind, SetOperator, SetQuantifier, SortItem, Span, + TableFactor, TableFactorKind, Truth, UnaryOp, }; use crate::validate::literal::{self, Number, TypedKind}; use crate::{ParseError, ParserOptions, parse_query}; @@ -452,10 +452,7 @@ impl Lowerer { (Some(GroupWith::Rollup), true) => return Ok(Grouping::Rollup(plain)), (Some(GroupWith::Cube), true) => return Ok(Grouping::Cube(plain)), (Some(GroupWith::GroupingSets(sets)), true) => { - let sets = sets - .iter() - .map(|s| self.set(s)) - .collect::>>()?; + let sets = sets.iter().map(|s| self.set(s)).collect::>>()?; return Ok(Grouping::GroupingSets { exprs: plain, sets }); } (Some(_), false) => { @@ -479,9 +476,8 @@ impl Lowerer { product = product .iter() .flat_map(|head| { - sets.iter().map(move |tail| { - head.iter().chain(tail).cloned().collect::>() - }) + sets.iter() + .map(move |tail| head.iter().chain(tail).cloned().collect::>()) }) .collect(); } @@ -545,7 +541,11 @@ impl Lowerer { }; chosen .into_iter() - .map(|c| c.into_iter().flat_map(|i| parts[i].iter().cloned()).collect()) + .map(|c| { + c.into_iter() + .flat_map(|i| parts[i].iter().cloned()) + .collect() + }) .collect() } }) diff --git a/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap b/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap index 2e9c27b99..ba9dd242f 100644 --- a/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap +++ b/crates/krishiv-sql/tests/snapshots/lower__constructs_lower_or_are_named_as_unsupported.snap @@ -21,9 +21,9 @@ SELECT a FROM t LATERAL VIEW explode(b) x AS y SELECT row_number() OVER (ORDER BY a) FROM t => [KRISHIV_UNSUPPORTED_FEATURE] The feature window functions is not supported by Krishiv yet. SQLSTATE: 0A000 @ 7..37 SELECT a FROM t GROUP BY ROLLUP (a) - => [KRISHIV_UNSUPPORTED_FEATURE] The feature ROLLUP, CUBE and GROUPING SETS is not supported by Krishiv yet. SQLSTATE: 0A000 @ 25..35 + => Aggregate [a] GROUPING SETS ((a), ()), [a] | +- Read t | SELECT a FROM t GROUP BY a WITH CUBE - => [KRISHIV_UNSUPPORTED_FEATURE] The feature ROLLUP, CUBE and GROUPING SETS is not supported by Krishiv yet. SQLSTATE: 0A000 @ 16..36 + => Aggregate CUBE [a], [a] | +- Read t | SELECT * FROM t NATURAL JOIN u => [KRISHIV_UNSUPPORTED_FEATURE] The feature NATURAL JOIN is not supported by Krishiv yet. SQLSTATE: 0A000 @ 16..30 SELECT * FROM t, LATERAL (SELECT a) From a41e41eb2161828041829a9e11a3b15f7a68d771 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:57:22 +0200 Subject: [PATCH 03/16] Address review threads for ROLLUP, CUBE and GROUPING SETS Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-common/src/stack.rs | 2 +- crates/krishiv-optimizer/src/stats.rs | 100 +++++++++++++++++- crates/krishiv-optimizer/tests/it/rules.rs | 48 ++++++++- crates/krishiv-plan/src/analyzer/analytics.rs | 100 +++++++++++++++--- .../tests/grouping_analytics.rs | 34 +++++- .../tests/grouping_analytics_api.rs | 48 +++++++-- crates/krishiv-sql/src/lower.rs | 2 +- crates/krishiv-sql/tests/lower.rs | 22 ++++ 8 files changed, 321 insertions(+), 35 deletions(-) diff --git a/crates/krishiv-common/src/stack.rs b/crates/krishiv-common/src/stack.rs index e44a30970..2a0488ec4 100644 --- a/crates/krishiv-common/src/stack.rs +++ b/crates/krishiv-common/src/stack.rs @@ -7,7 +7,7 @@ //! never needs more than `STEP` frames of the caller's stack. /// Levels a recursive pass may take on one stack segment. -pub const STEP: usize = 64; +pub const STEP: usize = 32; /// Runs `f` on a new thread with a 64 MiB stack (16 MiB in release /// builds), waiting for it; `None` when no thread can be started. A panic diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index d96d5abf1..8e5c52108 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -241,11 +241,107 @@ impl<'a> Estimator<'a> { | PlanNode::ReplicateRows { input, .. } => Est::clone(&self.plan(input)), // Every input row once per projection. PlanNode::Expand { - input, projections, .. + input, + projections, + output, } => { let est = Est::clone(&self.plan(input)); let rows = est.rows * to_f64(u64::try_from(projections.len()).unwrap_or(u64::MAX)); - est.with_rows(rows) + let rows = rows.max(1.0); + if projections.is_empty() + || output.is_empty() + || projections.iter().any(|p| p.len() != output.len()) + { + return est.with_rows(rows); + } + let per = est.rows; + let mut cols = HashMap::with_capacity(output.len()); + for (j, out) in output.iter().enumerate() { + let mut null_rows = 0.0; + let mut col_id: Option = None; + let mut mixed = false; + let mut lits: Vec<&Literal> = Vec::new(); + let mut other = false; + for proj in projections.iter() { + let s = strip_alias(&proj[j]); + match &s.node { + ExprNode::Literal(Literal::Null(_)) => null_rows += per, + ExprNode::Literal(l) => { + if col_id.is_some() { + mixed = true; + } + if !lits.contains(&l) { + lits.push(l); + } + } + ExprNode::Resolved( + ResolvedExpr::Column(c) | ResolvedExpr::OuterColumn(c), + ) => match est.cols.get(&c.id) { + Some(col) => { + null_rows += col.nulls * per; + match col_id { + None => col_id = Some(c.id), + Some(id) if id == c.id => {} + Some(_) => mixed = true, + } + if !lits.is_empty() { + mixed = true; + } + } + None => other = true, + }, + _ => other = true, + } + } + let nulls = (null_rows / rows).clamp(0.0, 1.0); + let distinct_lits = to_f64(u64::try_from(lits.len()).unwrap_or(u64::MAX)); + let ndv = if other { + rows + } else if let Some(id) = col_id { + let base = if mixed { + // Different input columns (or columns plus + // literals): at most their sum. + let mut seen = Vec::new(); + let mut sum = 0.0; + for proj in projections.iter() { + let s = strip_alias(&proj[j]); + if let ExprNode::Resolved( + ResolvedExpr::Column(c) | ResolvedExpr::OuterColumn(c), + ) = &s.node + && !seen.contains(&c.id) + { + seen.push(c.id); + sum += est.cols.get(&c.id).map(|c| c.ndv).unwrap_or(per); + } + } + sum + distinct_lits + } else if lits.is_empty() { + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + } else { + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits + }; + base.clamp(1.0, rows) + } else if lits.is_empty() { + // Every projection is NULL. + 1.0 + } else { + distinct_lits.clamp(1.0, rows) + }; + let range = if !mixed && !other && lits.is_empty() { + col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) + } else { + None + }; + cols.insert( + out.id, + Col { + ndv: ndv.clamp(1.0, rows), + nulls, + range, + }, + ); + } + Est { rows, cols } } PlanNode::Union { inputs, .. } => { let rows = inputs.iter().map(|i| self.plan(i).rows).sum(); diff --git a/crates/krishiv-optimizer/tests/it/rules.rs b/crates/krishiv-optimizer/tests/it/rules.rs index f797e7a46..27189251c 100644 --- a/crates/krishiv-optimizer/tests/it/rules.rs +++ b/crates/krishiv-optimizer/tests/it/rules.rs @@ -3,13 +3,14 @@ //! (when the result is deterministic) the same rows as the unoptimised plan //! on generated data. +use std::collections::HashMap; use std::sync::Arc; use krishiv_common::error::ErrorClass; -use krishiv_optimizer::{TopN, optimize}; +use krishiv_optimizer::{ColumnStats, Statistics, TableStats, TopN, optimize}; use krishiv_plan::{ Builtin, Expr, ExprNode, FunctionId, FunctionKind, Literal, LogicalPlan, Plan, PlanNode, - ResolvedExpr, ResolvedFunction, + ResolvedExpr, ResolvedFunction, Scan, }; use krishiv_types::DataType; @@ -476,9 +477,46 @@ fn an_expansion_has_a_copy_of_its_rows_per_projection() { }; let one_set = rows("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))"); let four_sets = rows("SELECT count(*) FROM t GROUP BY CUBE(a, b)"); - assert!( - (four_sets / one_set - 2.0).abs() < 1e-9, - "{four_sets} vs {one_set}" + assert_eq!(one_set, 2_000_000.0, "duplicate sets double the rows"); + assert_eq!(four_sets, 4_000_000.0, "a 2-cube quadruples the rows"); +} + +#[test] +fn expansion_copies_keep_their_input_ndv() { + #[derive(Debug)] + struct Stats(HashMap); + impl Statistics for Stats { + fn table(&self, scan: &Scan) -> Option { + let name = match &scan.source { + krishiv_plan::ReadSource::Table { name, .. } => name.last()?.clone(), + _ => return None, + }; + self.0.get(&name).cloned() + } + } + let stats = Stats(HashMap::from([( + "t".to_owned(), + TableStats { + rows: 1_000, + columns: HashMap::from([( + "a".to_owned(), + ColumnStats { + distinct: Some(10), + nulls: Some(0), + min: None, + max: None, + }, + )]), + }, + )])); + let ctx = support::ctx(); + let plan = support::resolve("SELECT a FROM t GROUP BY GROUPING SETS ((a), (a))", &ctx); + let optimized = optimize(&plan).unwrap(); + // Two projections of 1,000 rows; the grouping copy has 10 distinct + // values, the id one and the position two: 20 groups, not 2,000 rows. + assert_eq!( + krishiv_optimizer::estimated_rows(optimized.root(), &stats), + 20.0 ); } diff --git a/crates/krishiv-plan/src/analyzer/analytics.rs b/crates/krishiv-plan/src/analyzer/analytics.rs index 8173c6063..9e17591c9 100644 --- a/crates/krishiv-plan/src/analyzer/analytics.rs +++ b/crates/krishiv-plan/src/analyzer/analytics.rs @@ -45,7 +45,7 @@ use crate::{ /// The most grouping expressions one `GROUP BY` takes: the grouping id is a /// `BIGINT` with one bit each. -const MAX_GROUPING: usize = 63; +const MAX_GROUPING: usize = 64; /// The most expressions of a `CUBE`: 2^16 grouping sets (the SQL front end /// bounds `CUBE (..)` items the same way). @@ -623,7 +623,14 @@ impl Resolver<'_> { ) -> Result, AnalysisError> { match &e.node { ExprNode::Resolved(ResolvedExpr::Function(f)) if f.kind == FunctionKind::Aggregate => { - return Ok(Arc::clone(e)); + // `grouping` calls inside aggregate arguments still resolve, + // but ordinary columns stay on the input (no copy). + let rebuilt = + Arc::clone(e).map_children(|c| self.rewrite(&c, an, false, env, memo))?; + if Arc::ptr_eq(&rebuilt, e) { + return Ok(Arc::clone(e)); + } + return self.retype(rebuilt, env); } ExprNode::Subquery { .. } => return Ok(Arc::clone(e)), ExprNode::Function(f) if GroupingFn::from_name(&f.name).is_some() => { @@ -736,26 +743,91 @@ impl Resolver<'_> { )); }; // Bit `n - 1 - j` of the id, as a `TINYINT`. + // With 64 expressions the id can be negative (bit 63 set), + // where signed `div`/`%` on `i64::MIN` reads -1 instead of 1: + // there the quotient is floored before `pmod(_, 2)` (bit 63 + // is the sign). Smaller ids stay non-negative and keep the + // plain `div`/`%` extraction. let shift = n - 1 - j; - // The divisors are literals that are not 0, so the id's - // bit is never null: `grouping` is not nullable (Spark). - let high = if shift == 0 { - id + let bit = if n < 64 { + // The divisors are literals that are not 0, so the id's + // bit is never null: `grouping` is not nullable (Spark). + let high = if shift == 0 { + id + } else { + let high = self.call( + "div", + vec![id, literal(Literal::Long(1_i64 << shift))], + env, + Origin::UNKNOWN, + )?; + not_null(&high) + }; + self.call( + "%", + vec![high, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )? + } else if shift >= 63 { + let cond = self.call( + "<", + vec![Arc::clone(&id), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + self.call( + "when", + vec![cond, literal(Literal::Long(1)), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )? } else { + let pow = 1_i64 << shift; let high = self.call( "div", - vec![id, literal(Literal::Long(1_i64 << shift))], + vec![Arc::clone(&id), literal(Literal::Long(pow))], + env, + Origin::UNKNOWN, + )?; + let high = not_null(&high); + let rem = self.call( + "%", + vec![Arc::clone(&id), literal(Literal::Long(pow))], env, Origin::UNKNOWN, )?; - not_null(&high) + let neg = self.call( + "<", + vec![Arc::clone(&id), literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + let zero = self.call( + "==", + vec![rem, literal(Literal::Long(0))], + env, + Origin::UNKNOWN, + )?; + let nonzero = self.call("not", vec![zero], env, Origin::UNKNOWN)?; + let adjust = self.call("and", vec![neg, nonzero], env, Origin::UNKNOWN)?; + let one = literal(Literal::Long(1)); + let decremented = + self.call("-", vec![Arc::clone(&high), one], env, Origin::UNKNOWN)?; + let floor = self.call( + "when", + vec![adjust, decremented, Arc::clone(&high)], + env, + Origin::UNKNOWN, + )?; + // `pmod` (not `%`) so a negative quotient reads 0/1. + self.call( + "pmod", + vec![floor, literal(Literal::Long(2))], + env, + Origin::UNKNOWN, + )? }; - let bit = self.call( - "%", - vec![high, literal(Literal::Long(2))], - env, - Origin::UNKNOWN, - )?; let bit = not_null(&bit); let cast = Arc::new(Expr::new(ExprNode::Cast { expr: bit, diff --git a/crates/krishiv-session/tests/grouping_analytics.rs b/crates/krishiv-session/tests/grouping_analytics.rs index a544dce4a..0920d7102 100644 --- a/crates/krishiv-session/tests/grouping_analytics.rs +++ b/crates/krishiv-session/tests/grouping_analytics.rs @@ -28,11 +28,8 @@ struct Case { output: String, } -fn load(file: &str) -> R> { - let path = Path::new(env!("CARGO_MANIFEST_DIR")) - .join("../../tests/spark-golden/v4.2.0/results") - .join(file); - let text = std::fs::read_to_string(&path)?; +fn parse_cases(text: &str, file: &str) -> R> { + let total = text.split("-- !query\n").skip(1).count(); let cases: Vec = text .split("-- !query\n") .skip(1) @@ -49,9 +46,36 @@ fn load(file: &str) -> R> { if cases.is_empty() { return Err(format!("no statements in {file}").into()); } + if cases.len() != total { + return Err(format!( + "dropped {} of {total} statements in {file}", + total - cases.len() + ) + .into()); + } Ok(cases) } +fn load(file: &str) -> R> { + let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../tests/spark-golden/v4.2.0/results") + .join(file); + let text = std::fs::read_to_string(&path)?; + parse_cases(&text, file) +} + +#[test] +fn malformed_blocks_are_not_dropped_silently() { + let good = "-- !query\nSELECT 1\n-- !query schema\nstruct\n-- !query output\n1\n"; + assert_eq!(parse_cases(good, "good").unwrap().len(), 1); + let bad = + format!("{good}-- !query\nSELECT 2\n-- !query schema\nstruct\n-- !output\n2\n"); + assert!( + parse_cases(&bad, "bad").is_err(), + "a block without the schema/output separators must fail, not vanish" + ); +} + fn spark_type(t: &DataType) -> String { match t { DataType::Int8 => "tinyint".into(), diff --git a/crates/krishiv-session/tests/grouping_analytics_api.rs b/crates/krishiv-session/tests/grouping_analytics_api.rs index 9c6381437..0259df485 100644 --- a/crates/krishiv-session/tests/grouping_analytics_api.rs +++ b/crates/krishiv-session/tests/grouping_analytics_api.rs @@ -328,14 +328,17 @@ fn invalid_uses_have_spark_error_classes() -> R { error_class(&s, "SELECT grouping(a, b) FROM t GROUP BY ROLLUP(a, b)")?, ErrorClass::WrongNumArgsWithoutSuggestion ); - // 64 grouping expressions do not fit the 64-bit id. - let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); - let query = format!( - "SELECT count(*) FROM t GROUP BY ROLLUP({})", - keys.join(", ") - ); + // 65 grouping expressions do not fit the 64-bit id; 64 still do. + let rollup = |n: usize| { + let keys: Vec = (0..n).map(|i| format!("a + {i}")).collect(); + format!( + "SELECT count(*) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ) + }; + block_on(s.sql(&rollup(64)))?; // Analyzer-only: 65 sets would still run. assert_eq!( - error_class(&s, &query)?, + error_class(&s, &rollup(65))?, ErrorClass::GroupingSizeLimitExceeded ); // A cube has 2^n sets: 17 expressions are too many, 16 still are not. @@ -346,6 +349,7 @@ fn invalid_uses_have_spark_error_classes() -> R { keys.join(", ") ) }; + block_on(s.sql(&cube(16)))?; // Analyzer-only: 65,536 sets would still run. assert_eq!( error_class(&s, &cube(17))?, ErrorClass::GroupingSizeLimitExceeded @@ -359,6 +363,26 @@ fn invalid_uses_have_spark_error_classes() -> R { Ok(()) } +#[test] +fn grouping_reads_the_top_bit_of_a_64_expression_id() -> R { + let s = session()?; + let keys: Vec = (0..64).map(|i| format!("a + {i}")).collect(); + let query = format!( + "SELECT grouping(a + 0), grouping(a + 63) FROM t GROUP BY ROLLUP({})", + keys.join(", ") + ); + let got = rows(&s, &query)?; + assert!( + got.iter().any(|r| r == "1\t1"), + "the grand total (id -1) reads both bits as 1, got {got:?}" + ); + assert!( + !got.iter().any(|r| r.contains('-')), + "no negative bits, got {got:?}" + ); + Ok(()) +} + #[test] fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { let s = session()?; @@ -379,3 +403,13 @@ fn duplicate_sets_keep_their_rows_but_not_late_grouping_calls() -> R { ); Ok(()) } + +#[test] +fn grouping_calls_inside_aggregate_arguments_resolve() -> R { + let s = session()?; + assert_eq!( + rows(&s, "SELECT sum(grouping(a)) FROM t GROUP BY ROLLUP(a)")?, + sorted(&["0", "0", "3"]) + ); + Ok(()) +} diff --git a/crates/krishiv-sql/src/lower.rs b/crates/krishiv-sql/src/lower.rs index f14c2e530..8f1fe3488 100644 --- a/crates/krishiv-sql/src/lower.rs +++ b/crates/krishiv-sql/src/lower.rs @@ -508,7 +508,7 @@ impl Lowerer { GroupItemKind::GroupingSets(items) => { let mut sets = Vec::new(); for item in items { - sets.extend(self.item_sets(item, written)?); + sets.extend(self.deeper(|l| l.item_sets(item, &mut *written))?); } sets } diff --git a/crates/krishiv-sql/tests/lower.rs b/crates/krishiv-sql/tests/lower.rs index 33b469bba..013f87160 100644 --- a/crates/krishiv-sql/tests/lower.rs +++ b/crates/krishiv-sql/tests/lower.rs @@ -191,3 +191,25 @@ fn deep_queries_lower_on_a_small_stack_and_long_joins_hit_the_cap() { assert!(lower(&dotted).is_ok()); }); } + +#[test] +fn nested_grouping_sets_lower_on_a_small_stack() { + on_small_stack(|| { + // 900 nested `GROUPING SETS (...)`: the parser accepts it (below + // `MAX_DEPTH`), and lowering must hop stacks through the recursion. + let sql = format!( + "SELECT a FROM t GROUP BY {}a{}", + "GROUPING SETS (".repeat(900), + ")".repeat(900) + ); + assert!(lower(&sql).is_ok(), "nested GROUPING SETS lowers"); + // A single level nests the same way and keeps its sets. + let plan = lower("SELECT a FROM t GROUP BY GROUPING SETS (GROUPING SETS ((a)))") + .expect("nested sets lower"); + assert!( + plan.root().to_string().contains("GROUPING SETS"), + "keeps grouping sets: {}", + plan.root() + ); + }); +} From c997748ea516a76f82f25eeb5f54fc302ad75853 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:59:15 +0200 Subject: [PATCH 04/16] optimizer: name the IN-list literals in full (typos lint) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-optimizer/src/stats.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index 8e5c52108..cbbd57f2c 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -260,7 +260,7 @@ impl<'a> Estimator<'a> { let mut null_rows = 0.0; let mut col_id: Option = None; let mut mixed = false; - let mut lits: Vec<&Literal> = Vec::new(); + let mut literals: Vec<&Literal> = Vec::new(); let mut other = false; for proj in projections.iter() { let s = strip_alias(&proj[j]); @@ -270,8 +270,8 @@ impl<'a> Estimator<'a> { if col_id.is_some() { mixed = true; } - if !lits.contains(&l) { - lits.push(l); + if !literals.contains(&l) { + literals.push(l); } } ExprNode::Resolved( @@ -284,7 +284,7 @@ impl<'a> Estimator<'a> { Some(id) if id == c.id => {} Some(_) => mixed = true, } - if !lits.is_empty() { + if !literals.is_empty() { mixed = true; } } @@ -294,7 +294,7 @@ impl<'a> Estimator<'a> { } } let nulls = (null_rows / rows).clamp(0.0, 1.0); - let distinct_lits = to_f64(u64::try_from(lits.len()).unwrap_or(u64::MAX)); + let distinct_lits = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); let ndv = if other { rows } else if let Some(id) = col_id { @@ -315,19 +315,19 @@ impl<'a> Estimator<'a> { } } sum + distinct_lits - } else if lits.is_empty() { + } else if literals.is_empty() { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) } else { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits }; base.clamp(1.0, rows) - } else if lits.is_empty() { + } else if literals.is_empty() { // Every projection is NULL. 1.0 } else { distinct_lits.clamp(1.0, rows) }; - let range = if !mixed && !other && lits.is_empty() { + let range = if !mixed && !other && literals.is_empty() { col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) } else { None From b3a96acfda20943ad6ee3f181d19d28e4bf9e2f2 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:00:24 +0200 Subject: [PATCH 05/16] optimizer: finish the literals rename (typos lint) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-optimizer/src/stats.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index cbbd57f2c..255ab67b3 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -294,7 +294,7 @@ impl<'a> Estimator<'a> { } } let nulls = (null_rows / rows).clamp(0.0, 1.0); - let distinct_lits = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); + let distinct_literals = to_f64(u64::try_from(literals.len()).unwrap_or(u64::MAX)); let ndv = if other { rows } else if let Some(id) = col_id { @@ -314,18 +314,18 @@ impl<'a> Estimator<'a> { sum += est.cols.get(&c.id).map(|c| c.ndv).unwrap_or(per); } } - sum + distinct_lits + sum + distinct_literals } else if literals.is_empty() { est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) } else { - est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_lits + est.cols.get(&id).map(|c| c.ndv).unwrap_or(per) + distinct_literals }; base.clamp(1.0, rows) } else if literals.is_empty() { // Every projection is NULL. 1.0 } else { - distinct_lits.clamp(1.0, rows) + distinct_literals.clamp(1.0, rows) }; let range = if !mixed && !other && literals.is_empty() { col_id.and_then(|id| est.cols.get(&id).and_then(|c| c.range)) From 873c73a9a45a0c0c7b1dcfe5f76b3218cece8b3c Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Fri, 23 Oct 2026 10:02:13 +0200 Subject: [PATCH 06/16] optimizer: default selectivity for literal equalities without a distinct count Without a distinct count the estimator took every value of a column as distinct, so `col = literal` and `col IN (list)` estimated 1/rows of a table: about one row. A low-cardinality string column with only footer statistics (row counts and bounds, no distinct counts) then made a large input look like the smallest one, and join ordering built it: TPC-H q10 built the 15M `l_returnflag = 'R'` rows of lineitem at SF10, and q21 built `lineitem x orders` through `o_orderstatus = 'F'`. `Col` now records whether its distinct count comes from statistics; a literal equality or IN list on a column without one takes DEFAULT_SELECTIVITY (0.2). Equalities between columns keep 1/ndv. Rule and property tests that need a selective filter on tables without statistics now plan under explicit all-distinct statistics, the estimates they were written against; the TPC-H plans under default estimates and under footer statistics change (q02, q05, q07, q08, q10, q17, q19, q21 and q08, q10, q21 respectively). Design: P2-T166 (PR 1a, part 1). Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .../planner/snapshots/planner__tpch__q08.snap | 64 +++++++-------- .../planner/snapshots/planner__tpch__q10.snap | 32 ++++---- .../planner/snapshots/planner__tpch__q21.snap | 34 ++++---- crates/krishiv-optimizer/src/stats.rs | 28 +++++-- .../krishiv-optimizer/tests/it/join_order.rs | 54 +++++++++++++ crates/krishiv-optimizer/tests/it/props.rs | 6 +- crates/krishiv-optimizer/tests/it/rules.rs | 6 +- .../tests/it/snapshots/it__tpch__q02.snap | 77 ++++++++++--------- .../tests/it/snapshots/it__tpch__q05.snap | 38 ++++----- .../tests/it/snapshots/it__tpch__q07.snap | 35 +++++---- .../tests/it/snapshots/it__tpch__q08.snap | 44 +++++------ .../tests/it/snapshots/it__tpch__q10.snap | 14 ++-- .../tests/it/snapshots/it__tpch__q17.snap | 32 ++++---- .../tests/it/snapshots/it__tpch__q19.snap | 8 +- .../tests/it/snapshots/it__tpch__q21.snap | 32 ++++---- .../krishiv-optimizer/tests/it/support/mod.rs | 30 ++++++++ 16 files changed, 321 insertions(+), 213 deletions(-) diff --git a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q08.snap b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q08.snap index f2d09d2ef..ea1bb47da 100644 --- a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q08.snap +++ b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q08.snap @@ -9,38 +9,38 @@ Sort [o_year#60 ASC NULLS FIRST] +- Exchange hashpartitioning(o_year#60) +- Project [o_year#60, CASE WHEN (nation#62 == 'BRAZIL') THEN volume#61 ELSE 0.0000 END, volume#61] +- Project [extract('year', o_orderdate#36) AS o_year#60, (l_extendedprice#21 * (1 - l_discount#22)) AS volume#61, n_name#54 AS nation#62] - +- HashJoin Inner, [n_nationkey#53] = [s_nationkey#12], RuntimeFilter #7: n_nationkey#53 -> spark_catalog.default.nation.n_nationkey - :- Project [n_nationkey#53, n_name#54] - : +- Scan spark_catalog.default.nation [n_nationkey#49, n_name#50, n_regionkey#51] (ids +4) + +- HashJoin Inner, [l_suppkey#18] = [s_suppkey#9], RuntimeFilter #7: l_suppkey#18 -> spark_catalog.default.lineitem.l_suppkey + :- Project [l_suppkey#18, l_extendedprice#21, l_discount#22, o_orderdate#36] + : +- HashJoin Inner, [l_partkey#17] = [p_partkey#0], RuntimeFilter #5: l_partkey#17 -> spark_catalog.default.lineitem.l_partkey + : :- Project [l_partkey#17, l_suppkey#18, l_extendedprice#21, l_discount#22, o_orderdate#36] + : : +- HashJoin Inner, [l_orderkey#16] = [o_orderkey#32], RuntimeFilter #4: l_orderkey#16 -> spark_catalog.default.lineitem.l_orderkey + : : :- Scan spark_catalog.default.lineitem [l_orderkey#16, l_partkey#17, l_suppkey#18, l_extendedprice#21, l_discount#22] + : : +- BroadcastExchange + : : +- Project [o_orderkey#32, o_orderdate#36] + : : +- HashJoin Inner, [o_custkey#33] = [c_custkey#41], RuntimeFilter #3: o_custkey#33 -> spark_catalog.default.orders.o_custkey + : : :- Filter ((o_orderdate#36 >= DATE '1995-01-01') and (o_orderdate#36 <= DATE '1996-12-31')) + : : : +- Scan spark_catalog.default.orders [o_orderkey#32, o_custkey#33, o_orderdate#36], PushedFilters: [((o_orderdate#36 >= DATE '1995-01-01') and (o_orderdate#36 <= DATE '1996-12-31'))] + : : +- BroadcastExchange + : : +- Project [c_custkey#41] + : : +- HashJoin Inner, [c_nationkey#44] = [n_nationkey#49], RuntimeFilter #2: c_nationkey#44 -> spark_catalog.default.customer.c_nationkey + : : :- Scan spark_catalog.default.customer [c_custkey#41, c_nationkey#44] + : : +- BroadcastExchange + : : +- Project [n_nationkey#49] + : : +- HashJoin Inner, [n_regionkey#51] = [r_regionkey#57], RuntimeFilter #1: n_regionkey#51 -> spark_catalog.default.nation.n_regionkey + : : :- Project [n_nationkey#49, n_regionkey#51] + : : : +- Scan spark_catalog.default.nation [n_nationkey#49, n_name#50, n_regionkey#51] + : : +- BroadcastExchange + : : +- Project [r_regionkey#57] + : : +- Filter (r_name#58 == 'AMERICA') + : : +- Scan spark_catalog.default.region [r_regionkey#57, r_name#58], PushedFilters: [(r_name#58 == 'AMERICA')] + : +- BroadcastExchange + : +- Project [p_partkey#0] + : +- Filter (p_type#4 == 'ECONOMY ANODIZED STEEL') + : +- Scan spark_catalog.default.part [p_partkey#0, p_type#4], PushedFilters: [(p_type#4 == 'ECONOMY ANODIZED STEEL')] +- BroadcastExchange - +- Project [s_nationkey#12, o_orderdate#36, l_extendedprice#21, l_discount#22] - +- HashJoin Inner, [s_suppkey#9] = [l_suppkey#18], RuntimeFilter #6: s_suppkey#9 -> spark_catalog.default.supplier.s_suppkey + +- Project [s_suppkey#9, n_name#54] + +- HashJoin Inner, [s_nationkey#12] = [n_nationkey#53], RuntimeFilter #6: s_nationkey#12 -> spark_catalog.default.supplier.s_nationkey :- Scan spark_catalog.default.supplier [s_suppkey#9, s_nationkey#12] +- BroadcastExchange - +- Project [o_orderdate#36, l_suppkey#18, l_extendedprice#21, l_discount#22] - +- HashJoin Inner, [n_regionkey#51] = [r_regionkey#57], RuntimeFilter #5: n_regionkey#51 -> spark_catalog.default.nation.n_regionkey - :- Project [n_regionkey#51, o_orderdate#36, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- HashJoin Inner, [n_nationkey#49] = [c_nationkey#44], RuntimeFilter #4: n_nationkey#49 -> spark_catalog.default.nation.n_nationkey - : :- Project [n_nationkey#49, n_regionkey#51] - : : +- Scan spark_catalog.default.nation [n_nationkey#49, n_name#50, n_regionkey#51] - : +- BroadcastExchange - : +- Project [c_nationkey#44, o_orderdate#36, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- HashJoin Inner, [c_custkey#41] = [o_custkey#33], RuntimeFilter #3: c_custkey#41 -> spark_catalog.default.customer.c_custkey - : :- Scan spark_catalog.default.customer [c_custkey#41, c_nationkey#44] - : +- BroadcastExchange - : +- Project [o_custkey#33, o_orderdate#36, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- HashJoin Inner, [o_orderkey#32] = [l_orderkey#16], RuntimeFilter #2: o_orderkey#32 -> spark_catalog.default.orders.o_orderkey - : :- Filter ((o_orderdate#36 >= DATE '1995-01-01') and (o_orderdate#36 <= DATE '1996-12-31')) - : : +- Scan spark_catalog.default.orders [o_orderkey#32, o_custkey#33, o_orderdate#36], PushedFilters: [((o_orderdate#36 >= DATE '1995-01-01') and (o_orderdate#36 <= DATE '1996-12-31'))] - : +- BroadcastExchange - : +- Project [l_orderkey#16, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- HashJoin Inner, [l_partkey#17] = [p_partkey#0], RuntimeFilter #1: l_partkey#17 -> spark_catalog.default.lineitem.l_partkey - : :- Scan spark_catalog.default.lineitem [l_orderkey#16, l_partkey#17, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- BroadcastExchange - : +- Project [p_partkey#0] - : +- Filter (p_type#4 == 'ECONOMY ANODIZED STEEL') - : +- Scan spark_catalog.default.part [p_partkey#0, p_type#4], PushedFilters: [(p_type#4 == 'ECONOMY ANODIZED STEEL')] - +- BroadcastExchange - +- Project [r_regionkey#57] - +- Filter (r_name#58 == 'AMERICA') - +- Scan spark_catalog.default.region [r_regionkey#57, r_name#58], PushedFilters: [(r_name#58 == 'AMERICA')] + +- Project [n_nationkey#53, n_name#54] + +- Scan spark_catalog.default.nation [n_nationkey#49, n_name#50, n_regionkey#51] (ids +4) diff --git a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q10.snap b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q10.snap index 620c03b42..8cb5322a4 100644 --- a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q10.snap +++ b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q10.snap @@ -9,20 +9,20 @@ TopN 20, [revenue#37 DESC NULLS LAST] +- HashAggregate keys=[c_custkey#0, c_name#1, c_acctbal#5, c_phone#4, n_name#34, c_address#2, c_comment#7], functions=[sum((l_extendedprice#22 * (1 - l_discount#23)))] +- Exchange hashpartitioning(c_custkey#0, c_name#1, c_acctbal#5, c_phone#4, n_name#34, c_address#2, c_comment#7) +- Project [c_custkey#0, c_name#1, c_acctbal#5, c_phone#4, n_name#34, c_address#2, c_comment#7, (l_extendedprice#22 * (1 - l_discount#23))] - +- Project [n_name#34, c_custkey#0, c_name#1, c_address#2, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23] - +- HashJoin Inner, [n_nationkey#33] = [c_nationkey#3], RuntimeFilter #3: n_nationkey#33 -> spark_catalog.default.nation.n_nationkey - :- Scan spark_catalog.default.nation [n_nationkey#33, n_name#34] + +- Project [c_custkey#0, c_name#1, c_address#2, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23, n_name#34] + +- HashJoin Inner, [c_nationkey#3] = [n_nationkey#33], RuntimeFilter #3: c_nationkey#3 -> spark_catalog.default.customer.c_nationkey + :- Project [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23] + : +- HashJoin Inner, [c_custkey#0] = [o_custkey#9], RuntimeFilter #2: c_custkey#0 -> spark_catalog.default.customer.c_custkey + : :- Scan spark_catalog.default.customer [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7] + : +- BroadcastExchange + : +- Project [l_extendedprice#22, l_discount#23, o_custkey#9] + : +- HashJoin Inner, [l_orderkey#17] = [o_orderkey#8], RuntimeFilter #1: l_orderkey#17 -> spark_catalog.default.lineitem.l_orderkey + : :- Project [l_orderkey#17, l_extendedprice#22, l_discount#23] + : : +- Filter (l_returnflag#25 == 'R') + : : +- Scan spark_catalog.default.lineitem [l_orderkey#17, l_extendedprice#22, l_discount#23, l_returnflag#25], PushedFilters: [(l_returnflag#25 == 'R')] + : +- BroadcastExchange + : +- Project [o_orderkey#8, o_custkey#9] + : +- Filter ((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01')) + : +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12], PushedFilters: [((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01'))] +- BroadcastExchange - +- Project [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23] - +- HashJoin Inner, [c_custkey#0] = [o_custkey#9], RuntimeFilter #2: c_custkey#0 -> spark_catalog.default.customer.c_custkey - :- Scan spark_catalog.default.customer [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7] - +- BroadcastExchange - +- Project [o_custkey#9, l_extendedprice#22, l_discount#23] - +- HashJoin Inner, [o_orderkey#8] = [l_orderkey#17], RuntimeFilter #1: o_orderkey#8 -> spark_catalog.default.orders.o_orderkey - :- Project [o_orderkey#8, o_custkey#9] - : +- Filter ((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01')) - : +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12], PushedFilters: [((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01'))] - +- BroadcastExchange - +- Project [l_orderkey#17, l_extendedprice#22, l_discount#23] - +- Filter (l_returnflag#25 == 'R') - +- Scan spark_catalog.default.lineitem [l_orderkey#17, l_extendedprice#22, l_discount#23, l_returnflag#25], PushedFilters: [(l_returnflag#25 == 'R')] + +- Scan spark_catalog.default.nation [n_nationkey#33, n_name#34] diff --git a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q21.snap b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q21.snap index da33645dd..fd337ca37 100644 --- a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q21.snap +++ b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q21.snap @@ -11,25 +11,25 @@ TopN 100, [numwait#68 DESC NULLS LAST, s_name#1 ASC NULLS FIRST] +- Project [s_name#1] +- HashJoin LeftAnti, [l_orderkey#7] = [l_orderkey#52], residual: (NOT (l_suppkey#54 == l_suppkey#9)) :- HashJoin LeftSemi, [l_orderkey#7] = [l_orderkey#36], residual: (NOT (l_suppkey#38 == l_suppkey#9)), RuntimeFilter #4: l_orderkey#7 -> spark_catalog.default.lineitem.l_orderkey - : :- Project [s_name#1, l_orderkey#7, l_suppkey#9] - : : +- HashJoin Inner, [s_nationkey#3] = [n_nationkey#32], RuntimeFilter #3: s_nationkey#3 -> spark_catalog.default.supplier.s_nationkey - : : :- Project [s_name#1, s_nationkey#3, l_orderkey#7, l_suppkey#9] - : : : +- HashJoin Inner, [s_suppkey#0] = [l_suppkey#9], RuntimeFilter #2: s_suppkey#0 -> spark_catalog.default.supplier.s_suppkey - : : : :- Scan spark_catalog.default.supplier [s_suppkey#0, s_name#1, s_nationkey#3] + : :- Project [l_orderkey#7, l_suppkey#9, s_name#1] + : : +- HashJoin Inner, [l_suppkey#9] = [s_suppkey#0], RuntimeFilter #3: l_suppkey#9 -> spark_catalog.default.lineitem.l_suppkey + : : :- Project [l_orderkey#7, l_suppkey#9] + : : : +- HashJoin Inner, [l_orderkey#7] = [o_orderkey#23], RuntimeFilter #1: l_orderkey#7 -> spark_catalog.default.lineitem.l_orderkey + : : : :- Project [l_orderkey#7, l_suppkey#9] + : : : : +- Filter (l_receiptdate#19 > l_commitdate#18) + : : : : +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_commitdate#18, l_receiptdate#19], PushedFilters: [(l_receiptdate#19 > l_commitdate#18)] : : : +- BroadcastExchange - : : : +- Project [l_orderkey#7, l_suppkey#9] - : : : +- HashJoin Inner, [l_orderkey#7] = [o_orderkey#23], RuntimeFilter #1: l_orderkey#7 -> spark_catalog.default.lineitem.l_orderkey - : : : :- Project [l_orderkey#7, l_suppkey#9] - : : : : +- Filter (l_receiptdate#19 > l_commitdate#18) - : : : : +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_commitdate#18, l_receiptdate#19], PushedFilters: [(l_receiptdate#19 > l_commitdate#18)] - : : : +- BroadcastExchange - : : : +- Project [o_orderkey#23] - : : : +- Filter (o_orderstatus#25 == 'F') - : : : +- Scan spark_catalog.default.orders [o_orderkey#23, o_orderstatus#25], PushedFilters: [(o_orderstatus#25 == 'F')] + : : : +- Project [o_orderkey#23] + : : : +- Filter (o_orderstatus#25 == 'F') + : : : +- Scan spark_catalog.default.orders [o_orderkey#23, o_orderstatus#25], PushedFilters: [(o_orderstatus#25 == 'F')] : : +- BroadcastExchange - : : +- Project [n_nationkey#32] - : : +- Filter (n_name#33 == 'SAUDI ARABIA') - : : +- Scan spark_catalog.default.nation [n_nationkey#32, n_name#33], PushedFilters: [(n_name#33 == 'SAUDI ARABIA')] + : : +- Project [s_suppkey#0, s_name#1] + : : +- HashJoin Inner, [s_nationkey#3] = [n_nationkey#32], RuntimeFilter #2: s_nationkey#3 -> spark_catalog.default.supplier.s_nationkey + : : :- Scan spark_catalog.default.supplier [s_suppkey#0, s_name#1, s_nationkey#3] + : : +- BroadcastExchange + : : +- Project [n_nationkey#32] + : : +- Filter (n_name#33 == 'SAUDI ARABIA') + : : +- Scan spark_catalog.default.nation [n_nationkey#32, n_name#33], PushedFilters: [(n_name#33 == 'SAUDI ARABIA')] : +- BroadcastExchange : +- Project [l_orderkey#36, l_suppkey#38] : +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_commitdate#18, l_receiptdate#19] (ids +29) diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index 78fad1991..06f208af8 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -2,7 +2,8 @@ //! Table statistics and cardinality estimates. //! //! Estimates follow the textbook independence assumptions: a predicate's -//! selectivity is `1 / NDV` for an equality with a literal, the covered +//! selectivity is `1 / NDV` for an equality with a literal (when the +//! column's distinct count is known, else [`DEFAULT_SELECTIVITY`]), the covered //! fraction of `[min, max]` for a range, `1 / max(NDV)` for an equi-join, //! and [`DEFAULT_SELECTIVITY`] when nothing better is known; conjunctions //! multiply and disjunctions add. Without statistics a table holds @@ -77,8 +78,11 @@ pub fn estimated_rows(plan: &Arc, stats: &dyn Statistics) -> f64 { /// What is known of one attribute. #[derive(Debug, Clone)] pub(crate) struct Col { - /// Distinct non-null values, at least 1. + /// Distinct non-null values, at least 1. Without statistics, every + /// value is taken as distinct (`ndv` = rows). pub(crate) ndv: f64, + /// Whether `ndv` comes from statistics rather than that assumption. + pub(crate) ndv_known: bool, /// The fraction of rows that are null. pub(crate) nulls: f64, /// Bounds, as numbers on one scale. @@ -114,6 +118,7 @@ impl Est { fn distinct(rows: f64) -> Col { Col { ndv: rows.max(1.0), + ndv_known: false, nulls: 0.0, range: None, } @@ -283,6 +288,7 @@ impl<'a> Estimator<'a> { let s = stats.columns.get(&c.name).cloned().unwrap_or_default(); let col = Col { ndv: s.distinct.map_or(rows, to_f64).clamp(1.0, rows), + ndv_known: s.distinct.is_some(), nulls: s.nulls.map_or(0.0, |n| (to_f64(n) / rows).min(1.0)), range: s .min @@ -359,7 +365,19 @@ pub(crate) fn selectivity(pred: &Expr, est: &Est) -> f64 { (Some((Builtin::Eq | Builtin::EqNullSafe, [a, b], _)), _) => { match (column(a, est), column(b, est)) { (Some(x), Some(y)) => 1.0 / x.ndv.max(y.ndv), - (Some(x), None) | (None, Some(x)) => 1.0 / x.ndv, + // `col = literal`: one value of the column's. Without a + // distinct count, `1 / rows` would claim every value is + // unique (one matching row); a low-cardinality column + // (`l_returnflag = 'R'`) would then look tiny and be + // built. + (Some(x), None) | (None, Some(x)) => { + let with_literal = literal(a).is_some() || literal(b).is_some(); + if with_literal && !x.ndv_known { + DEFAULT_SELECTIVITY + } else { + 1.0 / x.ndv + } + } (None, None) => DEFAULT_SELECTIVITY, } } @@ -374,8 +392,8 @@ pub(crate) fn selectivity(pred: &Expr, est: &Est) -> f64 { } } (Some((Builtin::In, [v, list @ ..], _)), _) => match column(v, est) { - Some(c) => to_f64(u64::try_from(list.len()).unwrap_or(u64::MAX)) / c.ndv, - None => DEFAULT_SELECTIVITY, + Some(c) if c.ndv_known => to_f64(u64::try_from(list.len()).unwrap_or(u64::MAX)) / c.ndv, + _ => DEFAULT_SELECTIVITY, }, (Some((Builtin::IsNull, [a], _)), _) => { column(a, est).map_or(DEFAULT_SELECTIVITY, |c| c.nulls) diff --git a/crates/krishiv-optimizer/tests/it/join_order.rs b/crates/krishiv-optimizer/tests/it/join_order.rs index a2aebd2a0..f13de51d4 100644 --- a/crates/krishiv-optimizer/tests/it/join_order.rs +++ b/crates/krishiv-optimizer/tests/it/join_order.rs @@ -424,3 +424,57 @@ fn implied_filters_of_on_conditions_keep_the_rows() { "SELECT t.a, u.d FROM t LEFT JOIN u ON (t.b = 'x' AND u.d = 'y') OR (t.b = 'z' AND u.a > 2)", ); } + +/// The fixture's statistics without distinct counts for the columns +/// `keep` does not name, as Parquet footers give them (footers carry no +/// distinct counts). +fn without_distinct(keep: &[&str]) -> Fixture { + let mut f = fixture(); + for t in f.0.values_mut() { + for (name, c) in &mut t.columns { + if !keep.contains(&name.as_str()) { + c.distinct = None; + } + } + } + f +} + +/// `col = literal` and `col IN (...)` read the column's distinct count +/// when one is known; without one the selectivity is the default, not +/// `1 / rows` (which would estimate a single matching row for +/// `l_returnflag = 'R'` and make 60M rows look like the smallest input). +#[test] +fn literal_equalities_without_a_distinct_count_use_the_default() { + let ctx = support::ctx(); + let rows = |sql: &str, stats: &Fixture| { + krishiv_optimizer::estimated_rows(support::resolve(sql, &ctx).root(), stats) + }; + let known = fixture(); + let footers = without_distinct(&[]); + let lineitem = 59_986_052.0; + let close = |a: f64, b: f64| (a - b).abs() <= 1e-6 * b; + let eq = "SELECT l_orderkey FROM lineitem WHERE l_returnflag = 'R'"; + assert!( + close(rows(eq, &known), lineitem / 3.0), + "{}", + rows(eq, &known) + ); + assert!( + close(rows(eq, &footers), lineitem * 0.2), + "{}", + rows(eq, &footers) + ); + // The literal on the left reads the same. + let flipped = "SELECT l_orderkey FROM lineitem WHERE 'R' = l_returnflag"; + assert!(close(rows(flipped, &footers), lineitem * 0.2)); + let within = "SELECT l_orderkey FROM lineitem WHERE l_shipmode IN ('MAIL', 'SHIP')"; + assert!(close(rows(within, &known), lineitem * 2.0 / 7.0)); + assert!(close(rows(within, &footers), lineitem * 0.2)); + // A distinct count from statistics still counts without the others. + let only = without_distinct(&["l_returnflag"]); + assert!(close(rows(eq, &only), lineitem / 3.0)); + // An equality between two columns of one input keeps `1 / ndv`. + let cols = "SELECT l_orderkey FROM lineitem WHERE l_commitdate = l_receiptdate"; + assert!(close(rows(cols, &footers), 1.0), "{}", rows(cols, &footers)); +} diff --git a/crates/krishiv-optimizer/tests/it/props.rs b/crates/krishiv-optimizer/tests/it/props.rs index 4e431c2e4..313771348 100644 --- a/crates/krishiv-optimizer/tests/it/props.rs +++ b/crates/krishiv-optimizer/tests/it/props.rs @@ -7,7 +7,7 @@ use std::collections::HashSet; use std::sync::Arc; -use krishiv_optimizer::optimize; +use krishiv_optimizer::{optimize, optimize_with}; use krishiv_plan::{ AttrId, Builtin, Expr, ExprNode, FunctionId, Grouping, JoinCondition, JoinKind, LogicalPlan, PlanNode, ResolvedExpr, @@ -487,8 +487,8 @@ fn restriction_after_join_ordering_reaches_a_fixed_point() { ORDER BY c0, c1 LIMIT 4"; let ctx = support::ctx(); let plan = support::resolve(sql, &ctx); - let optimized = optimize(&plan).unwrap(); - let again = optimize(&optimized).unwrap(); + let optimized = optimize_with(&plan, &support::AllDistinct).unwrap(); + let again = optimize_with(&optimized, &support::AllDistinct).unwrap(); assert!( again.structurally_eq(&optimized), "not a fixed point:\n{}\n{}", diff --git a/crates/krishiv-optimizer/tests/it/rules.rs b/crates/krishiv-optimizer/tests/it/rules.rs index b69d424bf..e64f6452a 100644 --- a/crates/krishiv-optimizer/tests/it/rules.rs +++ b/crates/krishiv-optimizer/tests/it/rules.rs @@ -6,7 +6,7 @@ use std::sync::Arc; use krishiv_common::error::ErrorClass; -use krishiv_optimizer::{TopN, optimize}; +use krishiv_optimizer::{TopN, optimize, optimize_with}; use krishiv_plan::{ Builtin, Expr, ExprNode, FunctionId, FunctionKind, Literal, LogicalPlan, Plan, PlanNode, ResolvedExpr, ResolvedFunction, @@ -405,13 +405,13 @@ fn rules() { let ctx = support::ctx(); for (name, sql) in CASES { let plan = support::resolve(sql, &ctx); - let optimized = optimize(&plan).unwrap(); + let optimized = optimize_with(&plan, &support::AllDistinct).unwrap(); assert_eq!( optimized.root().resolved_output(), plan.root().resolved_output(), "{name} keeps its output" ); - let again = optimize(&optimized).unwrap(); + let again = optimize_with(&optimized, &support::AllDistinct).unwrap(); assert!( again.structurally_eq(&optimized), "{name} is a fixed point:\n{}", diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q02.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q02.snap index 9913a93db..49005db6b 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q02.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q02.snap @@ -6,39 +6,44 @@ Limit 100 +- Sort [s_acctbal#14 DESC NULLS LAST, n_name#22 ASC NULLS FIRST, s_name#10 ASC NULLS FIRST, p_partkey#0 ASC NULLS FIRST], true +- Project [s_acctbal#14, s_name#10, n_name#22, p_partkey#0, p_mfgr#2, s_address#11, s_phone#13, s_comment#15] +- Join Inner, (n_regionkey#23 == r_regionkey#25) - :- Project [n_name#22, n_regionkey#23, s_name#10, s_address#11, s_phone#13, s_acctbal#14, s_comment#15, p_partkey#0, p_mfgr#2] - : +- Join Inner, (s_nationkey#12 == n_nationkey#21) - : :- Scan spark_catalog.default.nation [n_nationkey#21, n_name#22, n_regionkey#23] (shared #1) - : +- Project [s_name#10, s_address#11, s_nationkey#12, s_phone#13, s_acctbal#14, s_comment#15, p_partkey#0, p_mfgr#2] - : +- Join Inner, (s_suppkey#9 == ps_suppkey#17) - : :- Scan spark_catalog.default.supplier [s_suppkey#9, s_name#10, s_address#11, s_nationkey#12, s_phone#13, s_acctbal#14, s_comment#15] (shared #2) - : +- Project [ps_suppkey#17, p_partkey#0, p_mfgr#2] - : +- Join Inner, ((ps_supplycost#19 == min(ps_supplycost)#47) and (p_partkey#0 == ps_partkey#16)) - : :- Scan spark_catalog.default.partsupp [ps_partkey#16, ps_suppkey#17, ps_supplycost#19] (shared #3) - : +- Project [p_partkey#0, p_mfgr#2, min(ps_supplycost)#47] - : +- Join Inner, (p_partkey#0 == ps_partkey#28) - : :- Project [p_partkey#0, p_mfgr#2] - : : +- Filter ((p_size#5 == 15) and (p_type#4 like '%BRASS')) - : : +- Scan spark_catalog.default.part [p_partkey#0, p_mfgr#2, p_type#4, p_size#5] - : +- Aggregate [ps_partkey#28], [min(ps_supplycost#31) AS min(ps_supplycost)#47, ps_partkey#28] - : +- Project [ps_partkey#28, ps_supplycost#31] - : +- Join Inner, (s_suppkey#33 == ps_suppkey#29) - : :- Shift +12 - : : +- Shared #3, printed above - : +- Project [s_suppkey#33] - : +- Join Inner, (s_nationkey#36 == n_nationkey#40) - : :- Project [s_suppkey#33, s_nationkey#36] - : : +- Shift +24 - : : +- Shared #2, printed above - : +- Project [n_nationkey#40] - : +- Join Inner, (n_regionkey#42 == r_regionkey#44) - : :- Project [n_nationkey#40, n_regionkey#42] - : : +- Shift +19 - : : +- Shared #1, printed above - : +- Project [r_regionkey#44] - : +- Shift +19 - : +- Filter (r_name#26 == 'EUROPE') - : +- Scan spark_catalog.default.region [r_regionkey#25, r_name#26] (shared #4) - +- Project [r_regionkey#25] - +- Filter (r_name#26 == 'EUROPE') - +- Shared #4, printed above + :- Project [r_regionkey#25] + : +- Filter (r_name#26 == 'EUROPE') + : +- Scan spark_catalog.default.region [r_regionkey#25, r_name#26] (shared #1) + +- Project [n_name#22, n_regionkey#23, s_name#10, s_address#11, s_phone#13, s_acctbal#14, s_comment#15, p_partkey#0, p_mfgr#2] + +- Join Inner, (s_nationkey#12 == n_nationkey#21) + :- Scan spark_catalog.default.nation [n_nationkey#21, n_name#22, n_regionkey#23] (shared #2) + +- Project [s_name#10, s_address#11, s_nationkey#12, s_phone#13, s_acctbal#14, s_comment#15, p_partkey#0, p_mfgr#2] + +- Join Inner, (s_suppkey#9 == ps_suppkey#17) + :- Scan spark_catalog.default.supplier [s_suppkey#9, s_name#10, s_address#11, s_nationkey#12, s_phone#13, s_acctbal#14, s_comment#15] (shared #3) + +- Project [ps_suppkey#17, p_partkey#0, p_mfgr#2] + +- Join Inner, ((ps_supplycost#19 == min(ps_supplycost)#47) and (p_partkey#0 == ps_partkey#16)) + :- Scan spark_catalog.default.partsupp [ps_partkey#16, ps_suppkey#17, ps_supplycost#19] (shared #4) + +- Project [min(ps_supplycost)#47, p_partkey#0, p_mfgr#2] + +- Join Inner, (p_partkey#0 == ps_partkey#28) + :- Aggregate [ps_partkey#28], [min(ps_supplycost#31) AS min(ps_supplycost)#47, ps_partkey#28] + : +- Join LeftSemi, (ps_partkey#28 == p_partkey#48) + : :- Project [ps_partkey#28, ps_supplycost#31] + : : +- Join Inner, (s_suppkey#33 == ps_suppkey#29) + : : :- Shift +12 + : : : +- Shared #4, printed above + : : +- Project [s_suppkey#33] + : : +- Join Inner, (s_nationkey#36 == n_nationkey#40) + : : :- Project [s_suppkey#33, s_nationkey#36] + : : : +- Shift +24 + : : : +- Shared #3, printed above + : : +- Project [n_nationkey#40] + : : +- Join Inner, (n_regionkey#42 == r_regionkey#44) + : : :- Project [n_nationkey#40, n_regionkey#42] + : : : +- Shift +19 + : : : +- Shared #2, printed above + : : +- Project [r_regionkey#44] + : : +- Shift +19 + : : +- Filter (r_name#26 == 'EUROPE') + : : +- Shared #1, printed above + : +- Shift +48 + : +- Project [p_partkey#0] + : +- Filter ((p_size#5 == 15) and (p_type#4 like '%BRASS')) + : +- Scan spark_catalog.default.part [p_partkey#0, p_mfgr#2, p_type#4, p_size#5] (shared #5) + +- Project [p_partkey#0, p_mfgr#2] + +- Filter ((p_size#5 == 15) and (p_type#4 like '%BRASS')) + +- Shared #5, printed above diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q05.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q05.snap index f3fc44e1e..66f133a33 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q05.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q05.snap @@ -4,23 +4,23 @@ expression: optimized.root().to_string() --- Sort [revenue#47 DESC NULLS LAST], true +- Aggregate [n_name#41], [n_name#41, sum((l_extendedprice#22 * (1 - l_discount#23))) AS revenue#47] - +- Project [l_extendedprice#22, l_discount#23, n_name#41] - +- Join Inner, ((l_orderkey#17 == o_orderkey#8) and (c_custkey#0 == o_custkey#9)) - :- Project [o_orderkey#8, o_custkey#9] - : +- Filter ((o_orderdate#12 >= DATE '1994-01-01') and (o_orderdate#12 < DATE '1995-01-01')) - : +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12] - +- Project [l_orderkey#17, l_extendedprice#22, l_discount#23, c_custkey#0, n_name#41] - +- Join Inner, (l_suppkey#19 == s_suppkey#33) - :- Scan spark_catalog.default.lineitem [l_orderkey#17, l_suppkey#19, l_extendedprice#22, l_discount#23] - +- Project [c_custkey#0, s_suppkey#33, n_name#41] - +- Join Inner, (c_nationkey#3 == s_nationkey#36) - :- Scan spark_catalog.default.customer [c_custkey#0, c_nationkey#3] - +- Project [s_suppkey#33, s_nationkey#36, n_name#41] - +- Join Inner, (s_nationkey#36 == n_nationkey#40) + +- Project [n_name#41, l_extendedprice#22, l_discount#23] + +- Join Inner, (n_regionkey#42 == r_regionkey#44) + :- Project [r_regionkey#44] + : +- Filter (r_name#45 == 'ASIA') + : +- Scan spark_catalog.default.region [r_regionkey#44, r_name#45] + +- Project [n_name#41, n_regionkey#42, l_extendedprice#22, l_discount#23] + +- Join Inner, (s_nationkey#36 == n_nationkey#40) + :- Scan spark_catalog.default.nation [n_nationkey#40, n_name#41, n_regionkey#42] + +- Project [l_extendedprice#22, l_discount#23, s_nationkey#36] + +- Join Inner, ((l_suppkey#19 == s_suppkey#33) and (l_orderkey#17 == o_orderkey#8)) + :- Scan spark_catalog.default.lineitem [l_orderkey#17, l_suppkey#19, l_extendedprice#22, l_discount#23] + +- Project [s_suppkey#33, s_nationkey#36, o_orderkey#8] + +- Join Inner, (c_nationkey#3 == s_nationkey#36) :- Scan spark_catalog.default.supplier [s_suppkey#33, s_nationkey#36] - +- Project [n_nationkey#40, n_name#41] - +- Join Inner, (n_regionkey#42 == r_regionkey#44) - :- Scan spark_catalog.default.nation [n_nationkey#40, n_name#41, n_regionkey#42] - +- Project [r_regionkey#44] - +- Filter (r_name#45 == 'ASIA') - +- Scan spark_catalog.default.region [r_regionkey#44, r_name#45] + +- Project [c_nationkey#3, o_orderkey#8] + +- Join Inner, (c_custkey#0 == o_custkey#9) + :- Scan spark_catalog.default.customer [c_custkey#0, c_nationkey#3] + +- Project [o_orderkey#8, o_custkey#9] + +- Filter ((o_orderdate#12 >= DATE '1994-01-01') and (o_orderdate#12 < DATE '1995-01-01')) + +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q07.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q07.snap index 0297bd5ec..67bd79f94 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q07.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q07.snap @@ -5,21 +5,22 @@ expression: optimized.root().to_string() Sort [supp_nation#48 ASC NULLS FIRST, cust_nation#49 ASC NULLS FIRST, l_year#50 ASC NULLS FIRST], true +- Aggregate [supp_nation#48, cust_nation#49, l_year#50], [supp_nation#48, cust_nation#49, l_year#50, sum(volume#51) AS revenue#52] +- Project [n_name#41 AS supp_nation#48, n_name#45 AS cust_nation#49, extract('year', l_shipdate#17) AS l_year#50, (l_extendedprice#12 * (1 - l_discount#13)) AS volume#51] - +- Join Inner, ((c_custkey#32 == o_custkey#24) and (o_orderkey#23 == l_orderkey#7)) - :- Scan spark_catalog.default.orders [o_orderkey#23, o_custkey#24] - +- Project [c_custkey#32, n_name#45, l_orderkey#7, l_extendedprice#12, l_discount#13, l_shipdate#17, n_name#41] - +- Join Inner, (c_nationkey#35 == n_nationkey#44) + +- Join Inner, ((c_nationkey#35 == n_nationkey#44) and (((n_name#41 == 'FRANCE') and (n_name#45 == 'GERMANY')) or ((n_name#41 == 'GERMANY') and (n_name#45 == 'FRANCE')))) + :- Shift +4 + : +- Filter ((n_name#41 == 'GERMANY') or (n_name#41 == 'FRANCE')) + : +- Scan spark_catalog.default.nation [n_nationkey#40, n_name#41] (shared #1) + +- Project [c_nationkey#35, n_name#41, l_extendedprice#12, l_discount#13, l_shipdate#17] + +- Join Inner, (c_custkey#32 == o_custkey#24) :- Scan spark_catalog.default.customer [c_custkey#32, c_nationkey#35] - +- Join Inner, (((n_name#41 == 'FRANCE') and (n_name#45 == 'GERMANY')) or ((n_name#41 == 'GERMANY') and (n_name#45 == 'FRANCE'))) - :- Shift +4 - : +- Filter ((n_name#41 == 'GERMANY') or (n_name#41 == 'FRANCE')) - : +- Scan spark_catalog.default.nation [n_nationkey#40, n_name#41] (shared #1) - +- Project [l_orderkey#7, l_extendedprice#12, l_discount#13, l_shipdate#17, n_name#41] - +- Join Inner, (s_suppkey#0 == l_suppkey#9) - :- Filter ((l_shipdate#17 >= DATE '1995-01-01') and (l_shipdate#17 <= DATE '1996-12-31')) - : +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_extendedprice#12, l_discount#13, l_shipdate#17] - +- Project [s_suppkey#0, n_name#41] - +- Join Inner, (s_nationkey#3 == n_nationkey#40) - :- Scan spark_catalog.default.supplier [s_suppkey#0, s_nationkey#3] - +- Filter ((n_name#41 == 'FRANCE') or (n_name#41 == 'GERMANY')) - +- Shared #1, printed above + +- Project [o_custkey#24, n_name#41, l_extendedprice#12, l_discount#13, l_shipdate#17] + +- Join Inner, (o_orderkey#23 == l_orderkey#7) + :- Scan spark_catalog.default.orders [o_orderkey#23, o_custkey#24] + +- Project [n_name#41, l_orderkey#7, l_extendedprice#12, l_discount#13, l_shipdate#17] + +- Join Inner, (s_nationkey#3 == n_nationkey#40) + :- Filter ((n_name#41 == 'FRANCE') or (n_name#41 == 'GERMANY')) + : +- Shared #1, printed above + +- Project [s_nationkey#3, l_orderkey#7, l_extendedprice#12, l_discount#13, l_shipdate#17] + +- Join Inner, (s_suppkey#0 == l_suppkey#9) + :- Scan spark_catalog.default.supplier [s_suppkey#0, s_nationkey#3] + +- Filter ((l_shipdate#17 >= DATE '1995-01-01') and (l_shipdate#17 <= DATE '1996-12-31')) + +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_extendedprice#12, l_discount#13, l_shipdate#17] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q08.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q08.snap index cbad225ca..768420300 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q08.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q08.snap @@ -9,28 +9,28 @@ Sort [o_year#60 ASC NULLS FIRST], true :- Project [n_nationkey#53, n_name#54] : +- Shift +4 : +- Scan spark_catalog.default.nation [n_nationkey#49, n_name#50, n_regionkey#51] (shared #1) - +- Project [s_nationkey#12, o_orderdate#36, l_extendedprice#21, l_discount#22] + +- Project [s_nationkey#12, l_extendedprice#21, l_discount#22, o_orderdate#36] +- Join Inner, (s_suppkey#9 == l_suppkey#18) :- Scan spark_catalog.default.supplier [s_suppkey#9, s_nationkey#12] - +- Project [o_orderdate#36, l_suppkey#18, l_extendedprice#21, l_discount#22] + +- Project [l_suppkey#18, l_extendedprice#21, l_discount#22, o_orderdate#36] +- Join Inner, (n_regionkey#51 == r_regionkey#57) - :- Project [n_regionkey#51, o_orderdate#36, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- Join Inner, (c_nationkey#44 == n_nationkey#49) - : :- Project [n_nationkey#49, n_regionkey#51] - : : +- Shared #1, printed above - : +- Project [c_nationkey#44, o_orderdate#36, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- Join Inner, (o_custkey#33 == c_custkey#41) - : :- Scan spark_catalog.default.customer [c_custkey#41, c_nationkey#44] - : +- Project [o_custkey#33, o_orderdate#36, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- Join Inner, (l_orderkey#16 == o_orderkey#32) - : :- Filter ((o_orderdate#36 >= DATE '1995-01-01') and (o_orderdate#36 <= DATE '1996-12-31')) - : : +- Scan spark_catalog.default.orders [o_orderkey#32, o_custkey#33, o_orderdate#36] - : +- Project [l_orderkey#16, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- Join Inner, (p_partkey#0 == l_partkey#17) - : :- Scan spark_catalog.default.lineitem [l_orderkey#16, l_partkey#17, l_suppkey#18, l_extendedprice#21, l_discount#22] - : +- Project [p_partkey#0] - : +- Filter (p_type#4 == 'ECONOMY ANODIZED STEEL') - : +- Scan spark_catalog.default.part [p_partkey#0, p_type#4] - +- Project [r_regionkey#57] - +- Filter (r_name#58 == 'AMERICA') - +- Scan spark_catalog.default.region [r_regionkey#57, r_name#58] + :- Project [r_regionkey#57] + : +- Filter (r_name#58 == 'AMERICA') + : +- Scan spark_catalog.default.region [r_regionkey#57, r_name#58] + +- Project [n_regionkey#51, l_suppkey#18, l_extendedprice#21, l_discount#22, o_orderdate#36] + +- Join Inner, (c_nationkey#44 == n_nationkey#49) + :- Project [n_nationkey#49, n_regionkey#51] + : +- Shared #1, printed above + +- Project [c_nationkey#44, l_suppkey#18, l_extendedprice#21, l_discount#22, o_orderdate#36] + +- Join Inner, (o_custkey#33 == c_custkey#41) + :- Scan spark_catalog.default.customer [c_custkey#41, c_nationkey#44] + +- Project [l_suppkey#18, l_extendedprice#21, l_discount#22, o_custkey#33, o_orderdate#36] + +- Join Inner, (p_partkey#0 == l_partkey#17) + :- Project [p_partkey#0] + : +- Filter (p_type#4 == 'ECONOMY ANODIZED STEEL') + : +- Scan spark_catalog.default.part [p_partkey#0, p_type#4] + +- Project [l_partkey#17, l_suppkey#18, l_extendedprice#21, l_discount#22, o_custkey#33, o_orderdate#36] + +- Join Inner, (l_orderkey#16 == o_orderkey#32) + :- Scan spark_catalog.default.lineitem [l_orderkey#16, l_partkey#17, l_suppkey#18, l_extendedprice#21, l_discount#22] + +- Filter ((o_orderdate#36 >= DATE '1995-01-01') and (o_orderdate#36 <= DATE '1996-12-31')) + +- Scan spark_catalog.default.orders [o_orderkey#32, o_custkey#33, o_orderdate#36] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q10.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q10.snap index fb9a56719..8b8b6f795 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q10.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q10.snap @@ -11,11 +11,11 @@ Limit 20 +- Project [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23] +- Join Inner, (c_custkey#0 == o_custkey#9) :- Scan spark_catalog.default.customer [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7] - +- Project [o_custkey#9, l_extendedprice#22, l_discount#23] + +- Project [l_extendedprice#22, l_discount#23, o_custkey#9] +- Join Inner, (l_orderkey#17 == o_orderkey#8) - :- Project [o_orderkey#8, o_custkey#9] - : +- Filter ((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01')) - : +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12] - +- Project [l_orderkey#17, l_extendedprice#22, l_discount#23] - +- Filter (l_returnflag#25 == 'R') - +- Scan spark_catalog.default.lineitem [l_orderkey#17, l_extendedprice#22, l_discount#23, l_returnflag#25] + :- Project [l_orderkey#17, l_extendedprice#22, l_discount#23] + : +- Filter (l_returnflag#25 == 'R') + : +- Scan spark_catalog.default.lineitem [l_orderkey#17, l_extendedprice#22, l_discount#23, l_returnflag#25] + +- Project [o_orderkey#8, o_custkey#9] + +- Filter ((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01')) + +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q17.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q17.snap index d9c9de33d..ad5e14a57 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q17.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q17.snap @@ -4,19 +4,19 @@ expression: optimized.root().to_string() --- Aggregate [], [(sum(l_extendedprice#5) / 7.0) AS avg_yearly#42] +- Project [l_extendedprice#5] - +- Join Inner, ((p_partkey#16 == l_partkey#26) and (CAST(l_quantity#4 AS DECIMAL(21,7)) < (0.2 * avg(l_quantity))#41)) - :- Project [l_quantity#4, l_extendedprice#5, p_partkey#16] - : +- Join Inner, (p_partkey#16 == l_partkey#1) - : :- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5] (shared #1) - : +- Project [p_partkey#16] - : +- Filter ((p_brand#19 == 'Brand#23') and (p_container#22 == 'MED BOX')) - : +- Scan spark_catalog.default.part [p_partkey#16, p_brand#19, p_container#22] (shared #2) - +- Aggregate [l_partkey#26], [(0.2 * avg(l_quantity#29)) AS (0.2 * avg(l_quantity))#41, l_partkey#26] - +- Join LeftSemi, (l_partkey#26 == p_partkey#59) - :- Project [l_partkey#26, l_quantity#29] - : +- Shift +25 - : +- Shared #1, printed above - +- Shift +43 - +- Project [p_partkey#16] - +- Filter ((p_brand#19 == 'Brand#23') and (p_container#22 == 'MED BOX')) - +- Shared #2, printed above + +- Join Inner, ((CAST(l_quantity#4 AS DECIMAL(21,7)) < (0.2 * avg(l_quantity))#41) and (p_partkey#16 == l_partkey#1)) + :- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5] (shared #1) + +- Project [(0.2 * avg(l_quantity))#41, p_partkey#16] + +- Join Inner, (p_partkey#16 == l_partkey#26) + :- Aggregate [l_partkey#26], [(0.2 * avg(l_quantity#29)) AS (0.2 * avg(l_quantity))#41, l_partkey#26] + : +- Join LeftSemi, (l_partkey#26 == p_partkey#59) + : :- Project [l_partkey#26, l_quantity#29] + : : +- Shift +25 + : : +- Shared #1, printed above + : +- Shift +43 + : +- Project [p_partkey#16] + : +- Filter ((p_brand#19 == 'Brand#23') and (p_container#22 == 'MED BOX')) + : +- Scan spark_catalog.default.part [p_partkey#16, p_brand#19, p_container#22] (shared #2) + +- Project [p_partkey#16] + +- Filter ((p_brand#19 == 'Brand#23') and (p_container#22 == 'MED BOX')) + +- Shared #2, printed above diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap index 2f3755c0c..b22844ab2 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap @@ -5,7 +5,7 @@ expression: optimized.root().to_string() Aggregate [], [sum((l_extendedprice#5 * (1 - l_discount#6))) AS revenue#25] +- Project [l_extendedprice#5, l_discount#6] +- Join Inner, ((((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#12')) and ((p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')) and (l_quantity#4 >= 1.00))) and (((l_quantity#4 <= 11.00) and ((p_size#21 >= 1) and (p_size#21 <= 5))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#23')) and ((p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')) and (l_quantity#4 >= 10.00))) and (((l_quantity#4 <= 20.00) and ((p_size#21 >= 1) and (p_size#21 <= 10))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#34')) and ((p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')) and (l_quantity#4 >= 20.00))) and (((l_quantity#4 <= 30.00) and ((p_size#21 >= 1) and (p_size#21 <= 15))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) - :- Filter ((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) - : +- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5, l_discount#6, l_shipinstruct#13, l_shipmode#14] - +- Filter ((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 5)) or ((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (p_size#21 >= 1)) and (p_size#21 <= 10))) or ((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 15))) - +- Scan spark_catalog.default.part [p_partkey#16, p_brand#19, p_size#21, p_container#22] + :- Filter ((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 5)) or ((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (p_size#21 >= 1)) and (p_size#21 <= 10))) or ((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 15))) + : +- Scan spark_catalog.default.part [p_partkey#16, p_brand#19, p_size#21, p_container#22] + +- Filter ((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) + +- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5, l_discount#6, l_shipinstruct#13, l_shipmode#14] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q21.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q21.snap index 92007a4a0..d581af153 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q21.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q21.snap @@ -8,22 +8,22 @@ Limit 100 +- Project [s_name#1] +- Join LeftAnti, ((l_orderkey#52 == l_orderkey#7) and (NOT (l_suppkey#54 == l_suppkey#9))) :- Join LeftSemi, ((l_orderkey#36 == l_orderkey#7) and (NOT (l_suppkey#38 == l_suppkey#9))) - : :- Project [l_orderkey#7, l_suppkey#9, s_name#1] - : : +- Join Inner, (o_orderkey#23 == l_orderkey#7) - : : :- Project [l_orderkey#7, l_suppkey#9, s_name#1] - : : : +- Join Inner, (s_suppkey#0 == l_suppkey#9) - : : : :- Project [l_orderkey#7, l_suppkey#9] - : : : : +- Filter (l_receiptdate#19 > l_commitdate#18) - : : : : +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_commitdate#18, l_receiptdate#19] (shared #1) - : : : +- Project [s_suppkey#0, s_name#1] - : : : +- Join Inner, (s_nationkey#3 == n_nationkey#32) - : : : :- Scan spark_catalog.default.supplier [s_suppkey#0, s_name#1, s_nationkey#3] - : : : +- Project [n_nationkey#32] - : : : +- Filter (n_name#33 == 'SAUDI ARABIA') - : : : +- Scan spark_catalog.default.nation [n_nationkey#32, n_name#33] - : : +- Project [o_orderkey#23] - : : +- Filter (o_orderstatus#25 == 'F') - : : +- Scan spark_catalog.default.orders [o_orderkey#23, o_orderstatus#25] + : :- Project [s_name#1, l_orderkey#7, l_suppkey#9] + : : +- Join Inner, (s_nationkey#3 == n_nationkey#32) + : : :- Project [n_nationkey#32] + : : : +- Filter (n_name#33 == 'SAUDI ARABIA') + : : : +- Scan spark_catalog.default.nation [n_nationkey#32, n_name#33] + : : +- Project [s_name#1, s_nationkey#3, l_orderkey#7, l_suppkey#9] + : : +- Join Inner, (s_suppkey#0 == l_suppkey#9) + : : :- Scan spark_catalog.default.supplier [s_suppkey#0, s_name#1, s_nationkey#3] + : : +- Project [l_orderkey#7, l_suppkey#9] + : : +- Join Inner, (o_orderkey#23 == l_orderkey#7) + : : :- Project [l_orderkey#7, l_suppkey#9] + : : : +- Filter (l_receiptdate#19 > l_commitdate#18) + : : : +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_commitdate#18, l_receiptdate#19] (shared #1) + : : +- Project [o_orderkey#23] + : : +- Filter (o_orderstatus#25 == 'F') + : : +- Scan spark_catalog.default.orders [o_orderkey#23, o_orderstatus#25] : +- Project [l_orderkey#36, l_suppkey#38] : +- Shift +29 : +- Shared #1, printed above diff --git a/crates/krishiv-optimizer/tests/it/support/mod.rs b/crates/krishiv-optimizer/tests/it/support/mod.rs index f1f5eb87a..2ffb179f4 100644 --- a/crates/krishiv-optimizer/tests/it/support/mod.rs +++ b/crates/krishiv-optimizer/tests/it/support/mod.rs @@ -137,3 +137,33 @@ impl Rng { items[usize::try_from(self.range(0, i64::try_from(items.len()).unwrap() - 1)).unwrap()] } } + +/// Default row counts with every column's distinct count known, all values +/// distinct: the estimates `optimize` made before a literal equality on a +/// column without a distinct count took the default selectivity. Rule tests +/// that need a selective filter on a table without statistics plan under +/// them. +#[derive(Debug)] +pub(crate) struct AllDistinct; + +impl krishiv_optimizer::Statistics for AllDistinct { + fn table(&self, scan: &krishiv_plan::Scan) -> Option { + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "a constant" + )] + let rows = krishiv_optimizer::DEFAULT_ROWS as u64; + let columns = scan.output.iter().map(|c| { + let stats = krishiv_optimizer::ColumnStats { + distinct: Some(rows), + ..krishiv_optimizer::ColumnStats::default() + }; + (c.name.clone(), stats) + }); + Some(krishiv_optimizer::TableStats { + rows, + columns: columns.collect(), + }) + } +} From aa2f23733fed7866f6196c1c11858dc8f550c1bb Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Fri, 23 Oct 2026 10:11:07 +0200 Subject: [PATCH 07/16] exec, plan-codec: join types that build their preserved side (format v4) Adds `RightSemi`, `RightAnti` and `RightOuter` to the physical join type: a left semi, left anti or left outer join that builds its left input and returns that side's rows by their matches. They are plan data only in this change: the hash and nested-loop join constructors refuse them with KRISHIV_UNSUPPORTED_FEATURE and the planner does not produce them yet. krishiv.plan.v1 gains the three join types as format version 4. A plan that uses one raises `min_reader_version` to 4, so an older reader refuses it as too new instead of as corrupt; other plans keep `min_reader_version` 2. The join-type decoder now matches every value explicitly (an unknown or unspecified value is corrupt) instead of mapping unlisted values to the null-aware anti join. New golden files under tests/golden/v4. Design: P2-T166 (PR 1a, part 2). Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/src/join/mod.rs | 63 +++++++++++++++++- crates/krishiv-exec/src/join/nested.rs | 8 ++- crates/krishiv-exec/src/join/probe.rs | 7 +- crates/krishiv-exec/tests/nested_loop.rs | 18 +++++ crates/krishiv-plan-codec/src/decode.rs | 16 +++-- crates/krishiv-plan-codec/src/encode.rs | 20 +++++- crates/krishiv-plan-codec/src/lib.rs | 2 +- .../src/proto/krishiv.plan.v1.pb.rs | 11 +++ .../tests/common/inventory.rs | 21 ++++++ .../tests/golden/v4/aggregate.bin | Bin 0 -> 1730 bytes .../tests/golden/v4/aggregate_final.bin | Bin 0 -> 848 bytes .../tests/golden/v4/exchange_broadcast.bin | Bin 0 -> 754 bytes .../tests/golden/v4/exchange_hash.bin | Bin 0 -> 758 bytes .../tests/golden/v4/exchange_range.bin | Bin 0 -> 764 bytes .../tests/golden/v4/exchange_round_robin.bin | Bin 0 -> 754 bytes .../tests/golden/v4/exchange_single.bin | Bin 0 -> 754 bytes .../tests/golden/v4/expressions.bin | Bin 0 -> 1470 bytes .../tests/golden/v4/file_scan.bin | Bin 0 -> 495 bytes .../tests/golden/v4/filter.bin | Bin 0 -> 967 bytes .../tests/golden/v4/hash_join.bin | Bin 0 -> 1030 bytes .../tests/golden/v4/hash_join_anti.bin | Bin 0 -> 879 bytes .../tests/golden/v4/hash_join_partitioned.bin | Bin 0 -> 1030 bytes .../tests/golden/v4/hash_join_right_anti.bin | Bin 0 -> 1030 bytes .../tests/golden/v4/hash_join_right_outer.bin | Bin 0 -> 1030 bytes .../tests/golden/v4/hash_join_right_semi.bin | Bin 0 -> 1030 bytes .../tests/golden/v4/limit.bin | Bin 0 -> 759 bytes .../tests/golden/v4/limit_offset.bin | Bin 0 -> 752 bytes .../tests/golden/v4/nested_loop_join.bin | Bin 0 -> 866 bytes .../tests/golden/v4/one_row.bin | Bin 0 -> 386 bytes .../tests/golden/v4/project.bin | Bin 0 -> 953 bytes .../golden/v4/project_without_nullable.bin | Bin 0 -> 392 bytes .../tests/golden/v4/replicate_rows.bin | Bin 0 -> 879 bytes .../tests/golden/v4/scan.bin | Bin 0 -> 742 bytes .../tests/golden/v4/sort.bin | Bin 0 -> 758 bytes .../tests/golden/v4/top_n.bin | Bin 0 -> 760 bytes .../tests/golden/v4/union.bin | Bin 0 -> 870 bytes .../tests/golden/v4/write_files.bin | Bin 0 -> 748 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 31 +++++++++ proto/krishiv/plan/v1/plan.proto | 5 ++ 39 files changed, 189 insertions(+), 13 deletions(-) create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/aggregate.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/aggregate_final.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/exchange_broadcast.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/exchange_hash.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/exchange_range.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/exchange_round_robin.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/exchange_single.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/expressions.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/file_scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/filter.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/hash_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/hash_join_anti.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/hash_join_partitioned.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/hash_join_right_anti.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/hash_join_right_outer.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/hash_join_right_semi.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/limit.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/limit_offset.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/nested_loop_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/one_row.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/project.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/project_without_nullable.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/replicate_rows.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/sort.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/top_n.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/union.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v4/write_files.bin diff --git a/crates/krishiv-exec/src/join/mod.rs b/crates/krishiv-exec/src/join/mod.rs index adddb6c88..2704a133c 100644 --- a/crates/krishiv-exec/src/join/mod.rs +++ b/crates/krishiv-exec/src/join/mod.rs @@ -150,8 +150,10 @@ use crate::operators::Predicate; use crate::pipeline::{ExecNode, PipelineBuilder}; use crate::rows; -/// The join type. The probe (left) side is preserved by the outer and anti -/// joins; the build (right) side never is. +/// The join type. The probe (left) side is preserved by the left outer and +/// anti joins; the build side is preserved by the right semi, right anti +/// and right outer joins, which build the input whose rows they return +/// (design P2-T166). #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum JoinType { /// Matching pairs. @@ -164,12 +166,30 @@ pub enum JoinType { LeftAnti, /// `NOT IN`: see the module docs. NullAwareAnti, + /// Build rows with at least one match, each once: a left semi join + /// that builds its left input. Planned data only for now: no operator + /// runs it yet. + RightSemi, + /// Build rows with no match: a left anti join that builds its left + /// input. Planned data only for now. + RightAnti, + /// Matching pairs, and each unmatched build row with NULL probe + /// columns: a left outer join that builds its preserved input. Planned + /// data only for now. + RightOuter, } impl JoinType { fn emits_build(self) -> bool { matches!(self, Self::Inner | Self::LeftOuter) } + + /// Whether the join returns rows of its build side by their matches + /// (right semi, right anti, right outer). + #[must_use] + pub fn builds_preserved(self) -> bool { + matches!(self, Self::RightSemi | Self::RightAnti | Self::RightOuter) + } } /// What a hash join joins. @@ -401,6 +421,11 @@ impl Shared { "a hash join needs matching, non-empty key lists".to_owned(), )); } + if join_type.builds_preserved() { + return Err(unsupported( + "hash joins that return rows of their build side by their matches", + )); + } if join_type == JoinType::NullAwareAnti && (n != 1 || filter.is_some() || null_safe.contains(&true)) { @@ -757,6 +782,40 @@ mod tests { } } + /// The joins that build their preserved side are plan data only: no + /// hash join runs them yet, and none is planned. + #[test] + fn joins_that_build_their_preserved_side_are_refused() { + for t in [ + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::RightOuter, + ] { + assert!(t.builds_preserved()); + let e = HashJoin::try_new(spec(t, 1, false)) + .err() + .unwrap_or_else(|| panic!("{t:?} was accepted")); + assert_eq!( + e.error_class(), + krishiv_common::error::ErrorClass::KrishivUnsupportedFeature, + "{t:?}" + ); + assert!( + e.to_string().contains("return rows of their build side"), + "{e}" + ); + } + for t in [ + JoinType::Inner, + JoinType::LeftOuter, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::NullAwareAnti, + ] { + assert!(!t.builds_preserved(), "{t:?}"); + } + } + #[test] fn specs_are_validated() { assert!(HashJoin::try_new(spec(JoinType::Inner, 0, false)).is_err()); diff --git a/crates/krishiv-exec/src/join/nested.rs b/crates/krishiv-exec/src/join/nested.rs index dfb58efb1..3f2052b7e 100644 --- a/crates/krishiv-exec/src/join/nested.rs +++ b/crates/krishiv-exec/src/join/nested.rs @@ -67,11 +67,17 @@ impl NestedLoopJoin { /// /// # Errors /// `KRISHIV_UNSUPPORTED_FEATURE` for the null-aware anti join (it needs - /// an equality key; the planner uses the hash join). + /// an equality key; the planner uses the hash join) and for the joins + /// that return rows of their build side by their matches. pub fn try_new(spec: NestedLoopSpec) -> Result { if spec.join_type == JoinType::NullAwareAnti { return Err(unsupported("a null-aware anti join without a key")); } + if spec.join_type.builds_preserved() { + return Err(unsupported( + "nested-loop joins that return rows of their build side by their matches", + )); + } let fields = |nullable: bool| -> Vec { spec.probe_schema .fields() diff --git a/crates/krishiv-exec/src/join/probe.rs b/crates/krishiv-exec/src/join/probe.rs index 282003af3..8dca8d045 100644 --- a/crates/krishiv-exec/src/join/probe.rs +++ b/crates/krishiv-exec/src/join/probe.rs @@ -266,7 +266,12 @@ impl Cursor { /// The rows the join type emits once every match is known. fn tail(&self, shared: &Shared) -> Result> { let want = match shared.join_type { - JoinType::Inner => return Ok(None), + // No probe row waits for the end: inner joins emit pairs as + // they match, and the joins that build their preserved side + // return build rows, never unmatched probe rows. + JoinType::Inner | JoinType::RightSemi | JoinType::RightAnti | JoinType::RightOuter => { + return Ok(None); + } JoinType::LeftSemi => MATCHED, JoinType::LeftOuter | JoinType::LeftAnti | JoinType::NullAwareAnti => UNMATCHED, }; diff --git a/crates/krishiv-exec/tests/nested_loop.rs b/crates/krishiv-exec/tests/nested_loop.rs index 67788034c..43d7de1fe 100644 --- a/crates/krishiv-exec/tests/nested_loop.rs +++ b/crates/krishiv-exec/tests/nested_loop.rs @@ -198,6 +198,24 @@ fn null_aware_and_weighted_inputs_are_refused() { }) .unwrap_err(); assert_eq!(e.error_class(), ErrorClass::KrishivUnsupportedFeature); + for join_type in [ + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::RightOuter, + ] { + let e = NestedLoopJoin::try_new(NestedLoopSpec { + join_type, + probe_schema: schema("p"), + build_schema: schema("b"), + filter: None, + }) + .unwrap_err(); + assert_eq!(e.error_class(), ErrorClass::KrishivUnsupportedFeature); + assert!( + e.to_string().contains("return rows of their build side"), + "{e}" + ); + } let weighted = |name: &str| { let b = &batches(name, &[Some(1), Some(2)], 2)[0]; vec![Batch::with_weights(b.data().clone(), Weights::Const(2)).unwrap()] diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 48d3d1124..5b3752440 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -568,12 +568,16 @@ fn aggregate_expr(a: &pb::AggregateExpr) -> Result { } 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, + Ok(match pb::JoinType::try_from(t) { + Ok(pb::JoinType::Inner) => JoinType::Inner, + Ok(pb::JoinType::LeftOuter) => JoinType::LeftOuter, + Ok(pb::JoinType::LeftSemi) => JoinType::LeftSemi, + Ok(pb::JoinType::LeftAnti) => JoinType::LeftAnti, + Ok(pb::JoinType::NullAwareAnti) => JoinType::NullAwareAnti, + Ok(pb::JoinType::RightSemi) => JoinType::RightSemi, + Ok(pb::JoinType::RightAnti) => JoinType::RightAnti, + Ok(pb::JoinType::RightOuter) => JoinType::RightOuter, + Ok(pb::JoinType::Unspecified) | Err(_) => return Err(corrupt("an unknown join type")), }) } diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index ab67f47fa..858c773db 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -28,6 +28,11 @@ use crate::error::{CodecError, Result}; use crate::proto as pb; use crate::{FORMAT_VERSION, MIN_READER_VERSION}; +/// The format version that added the joins that build their preserved +/// side (`RightSemi`, `RightAnti`, `RightOuter`): the reader a plan using +/// one needs. +const BUILD_PRESERVED_JOINS: u32 = 4; + /// Encodes `plan` as `krishiv.plan.v1` bytes. Deterministic: the same plan /// always gives the same bytes. /// @@ -297,7 +302,7 @@ impl Encoder { Ok(pb::HashJoinNode { probe: self.node_index(&join.probe)?, build: self.node_index(&join.build)?, - join_type: join_type(join.join_type) as i32, + join_type: self.join_type(join.join_type) as i32, probe_keys: indexes(&join.probe_keys), build_keys: indexes(&join.build_keys), null_safe: join.null_safe.clone(), @@ -311,11 +316,19 @@ impl Encoder { Ok(pb::NestedLoopJoinNode { probe: self.node_index(&join.probe)?, build: self.node_index(&join.build)?, - join_type: join_type(join.join_type) as i32, + join_type: self.join_type(join.join_type) as i32, residual: self.residual(join.residual.as_ref())?, }) } + fn join_type(&mut self, t: JoinType) -> pb::JoinType { + if t.builds_preserved() { + // An older reader does not know these join types. + self.min_reader = self.min_reader.max(BUILD_PRESERVED_JOINS); + } + join_type(t) + } + fn exchange(&self, exchange: &ExchangePlan) -> Result { Ok(pb::ExchangeNode { id: exchange.id.0, @@ -488,6 +501,9 @@ fn join_type(t: JoinType) -> pb::JoinType { JoinType::LeftSemi => pb::JoinType::LeftSemi, JoinType::LeftAnti => pb::JoinType::LeftAnti, JoinType::NullAwareAnti => pb::JoinType::NullAwareAnti, + JoinType::RightSemi => pb::JoinType::RightSemi, + JoinType::RightAnti => pb::JoinType::RightAnti, + JoinType::RightOuter => pb::JoinType::RightOuter, } } diff --git a/crates/krishiv-plan-codec/src/lib.rs b/crates/krishiv-plan-codec/src/lib.rs index 8ff657885..a451f1fff 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 = 3; +pub const FORMAT_VERSION: u32 = 4; /// 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 ccea08b3b..de4a9ed6a 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 @@ -886,6 +886,11 @@ pub enum JoinType { LeftSemi = 3, LeftAnti = 4, NullAwareAnti = 5, + /// v4: joins that build their preserved side and return its rows by + /// their matches. A plan that uses one needs a v4 reader. + RightSemi = 6, + RightAnti = 7, + RightOuter = 8, } impl JoinType { /// String value of the enum field names used in the ProtoBuf definition. @@ -900,6 +905,9 @@ impl JoinType { Self::LeftSemi => "JOIN_TYPE_LEFT_SEMI", Self::LeftAnti => "JOIN_TYPE_LEFT_ANTI", Self::NullAwareAnti => "JOIN_TYPE_NULL_AWARE_ANTI", + Self::RightSemi => "JOIN_TYPE_RIGHT_SEMI", + Self::RightAnti => "JOIN_TYPE_RIGHT_ANTI", + Self::RightOuter => "JOIN_TYPE_RIGHT_OUTER", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -911,6 +919,9 @@ impl JoinType { "JOIN_TYPE_LEFT_SEMI" => Some(Self::LeftSemi), "JOIN_TYPE_LEFT_ANTI" => Some(Self::LeftAnti), "JOIN_TYPE_NULL_AWARE_ANTI" => Some(Self::NullAwareAnti), + "JOIN_TYPE_RIGHT_SEMI" => Some(Self::RightSemi), + "JOIN_TYPE_RIGHT_ANTI" => Some(Self::RightAnti), + "JOIN_TYPE_RIGHT_OUTER" => Some(Self::RightOuter), _ => None, } } diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index cf1f5a041..6448dccc4 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -370,6 +370,27 @@ pub(crate) fn nodes() -> Vec<(&'static str, PhysNode)> { ..join(JoinType::NullAwareAnti) }), ), + ( + "hash_join_right_semi", + PhysNode::HashJoin(HashJoinPlan { + distribution: JoinDistribution::Partitioned, + ..join(JoinType::RightSemi) + }), + ), + ( + "hash_join_right_anti", + PhysNode::HashJoin(HashJoinPlan { + distribution: JoinDistribution::Partitioned, + ..join(JoinType::RightAnti) + }), + ), + ( + "hash_join_right_outer", + PhysNode::HashJoin(HashJoinPlan { + distribution: JoinDistribution::Partitioned, + ..join(JoinType::RightOuter) + }), + ), ( "hash_join_partitioned", PhysNode::HashJoin(HashJoinPlan { diff --git a/crates/krishiv-plan-codec/tests/golden/v4/aggregate.bin b/crates/krishiv-plan-codec/tests/golden/v4/aggregate.bin new file mode 100644 index 0000000000000000000000000000000000000000..9af9fc98cd4081e2d1b7b24e088483754f020978 GIT binary patch literal 1730 zcmcgt&2G~`5Z+lk@n)M+cT>Pv9ITwGNL+G4JwVkHs<>c!ETXI0NvyhdWpC7QGw#)r9F5htBsOR#TTrE#@ zarE*~8=dAu=|u5Bk5ePTC|BCYc@!%lo24m|L!-K~Ve+FmQ}@(eX`d8%q}EP#5s&h5 z89%68{)zwv+UKGXfZq~aL5Yi2u-CIB=$@2N%0`l9anTo&19xo?wi00}a{C8u#+bmt zn0+xb2DsWlb#ZptEw+X}+U4Txp$4e9i{Hmtur_*Es@Q>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_}?Wgg;C`ju^Jy0;|OW6pM;mG&^{f@N5JU0q|1V$$*_HM=> zT02Syqa-uIsV;+LnB_$r?Ro;->1WG({dW5O)5j0ZuA?)ZzA8<)I`guUU0_>W@NIYkp{@AlM%d}COL&AHM7wg`?scT4vxv+VT$+CXg-nCg literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/aggregate_final.bin b/crates/krishiv-plan-codec/tests/golden/v4/aggregate_final.bin new file mode 100644 index 0000000000000000000000000000000000000000..bd35ef78de337f67f83f68f0ac870769b847e600 GIT binary patch literal 848 zcmbVKv2N5r5Zzh(&U=#!arS~R7KL?HWQoQVN&!(S6wzQ(f?yTgjhpSMX-?iB8*9pX ziC&g3%(}D)R{Q9Ftg~Dz9JjV9j;)^K$kr!$r621@=$p2f!w*v-grXyadu@dP@^(ml zvc7mAwkXH1e6j&)L`tvZ6S7K7DZA0vMjRf>dkjm?!Bky&+kISIy$BbFS@=#t1o6>* zn7M`sfWL&r?ax5TYzb6zfL+G_D(@oTMt+lClJ1kT|F-KoZsi&rBf1YVrnbg(SmY2Q z%j03PZ@23_#$e%*C_?6;auwrpO-4G5RQxy6fm;N3HP+tL-Tzk@!bm3QzRDL?LpfZ{ zLc^y7ag-Lzl??%s8M8S%5Z+n)&U=$P;_L-sRutA%$WofCPzs1rAw+|rM1oc9yR%~K$lg<+qd`i# z6g&Vgz#~u~@h~ZP2Qy;_qzF28^?mmIexCIzKa}P|;wPyB{Ssf`tUa^M>6;VlY+0`{ z$nurlmJY$`5dC*`mTQIM)-}b6(+eEA`ZTZfeSHu8v!>4U?zwI9O}%aNgFa+-6lkzU zs!++CLJ2bTBRDLSr#WoiPK6MPwh;ae3jt*96GGxp z+!4DJFn06Nlg;r vq-d+!Rok+RU;*M#DtQD?A@#L8S%5Z+nuo%iN0#MukNSQOS($WofCPzs1rAw+|rM1oc9yR%~K$lg<+qd`i# z6g&Vgz#~u~@h~ZP2Qy;_qzF28^nLdH@9eHtfs8IBev;ZZFYyKLv}d+CeRE=+E$cP< zS-!H{(jhn#pn6wlxzRXoT~nMmv%rz7PxH#$H}_DTHFaio&ux=$>TR1J^sZ>C!GMjV z9>ow%v^3C?GKpU$9)e8}p~X>ARe7`2C_x4)gu_C6hTYNInGiyg2%&CR2q0~r5D*9A zj@Tt1M+L+oAtq3id_b(ljJ(r6H)ivY-q&!X9MAN%zt`Kv^^<7Upq@r!R44_+kX*@s zZGhYTAY2fR2&}*Ex{e!iJoUi1|1yoK`Iv;4?3sy&;0=7fo#%56Uc7OAGc;#$UCzk} z!{Fkt!H3==8S%5Z+nuo%bde&e;pXSQOS($WofCPzs1rAw+|rM1oc9yR%~K$lg<+qd`i# z6g&Vgz#~u~@h~ZP2Qy;_qzF28^?mlt?ELO(70Bp9;wPzn^Aca+ZhK~%(>Ev9*|J`v zpXDpNEggb00jhU(mK%-Z)-}b6GYcHK`ZTZ1eRB`hSyN|bc5a(|Q*Yb+pf^Q38VuM- z>QM~QL`wrL36uC$;vv`s5n3D-Rh2hOjS^&_LO3k6XV@IQoeCi|Z6VYR3jrkU69VEu z+z~Uf@u`3~B*X+dB_9xLF(vCx?;EpvaPMO{5{{?(+TZK-;`&Ksnn2H@DU~HdK&2so z7?LabuQE9455fiEh`{peuIo50pQlV1_g|(ll`y8wOZLpft#BEiZ~OTigBO>tZ(3(B zj>|q7VQ63c)&7uMggg#js^5cOOY;AY%~Fiz1l99=DGmONfcerR2~ zd+YtAQ4->$Xsg;)+p-K{0pd_fwm+u3$s9dfWvuq#xB3_0 CBVaNB literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/exchange_round_robin.bin b/crates/krishiv-plan-codec/tests/golden/v4/exchange_round_robin.bin new file mode 100644 index 0000000000000000000000000000000000000000..8cbe4ff68cbcc5bd92c90be8ecd580d05c93aff5 GIT binary patch literal 754 zcmbtSy>8S%5Z+n)?7hh)arS~R7KL>cvP5$gN&!(SglI67NU(~1cUH_gviB6|Xpllj z!2|FDJOTw050ip-Ff(>QilAdx-)GP7=UK1vBbi)E{2*1RpW`#!?aobm_UhC)Q#Wf2 zi*jYQwL@?^LjO%ulv-iXxwblWdVv$yoRwDJ(YMh*Z<|8zUzoPsG~2GcKZKGU1sY5u z)g*;vs-%XFjA{BJ^AJpi2pvu;Ys+@2P=bv77>+9CX$~i^XF>=?TL}M}g#fY+2@!E5 z_QXELI4dHK2`PcHq@wg$ILQTs=xu8>qcRE<_61C;2aN zmij@sBpegie%bdui*Y{n$8>ly1tc8M<|PMa;sIFv^Ub=Dd+_3qR@Vb_75C+uf-nv( z{u+AJ7BL?JFMrr$T}Sr+9+sgV%PIO#%azkmkFRE>;q4M}l2_}EO8}O0^z){1nl6qr zRZag|$OqnuHtD9aF1;{aTCHupWnWI2^OKTrheDCeD}t!$&)O` vY1P@r*==3NumEu+l{|*Wko#KRJIv?V0m4IyklLTnt!$2-u?4$5`0f4$Xsutg literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/exchange_single.bin b/crates/krishiv-plan-codec/tests/golden/v4/exchange_single.bin new file mode 100644 index 0000000000000000000000000000000000000000..a61b36b3b6b9ca7e776654a2a1ef68f8942b3ba8 GIT binary patch literal 754 zcmbtSy>8S%5Z+n)?7hi_ID0`Di^94JSxR#iN&!(SglI67NU(~1cUH_gviB6|Xpllj z!2|FDJOTw050ip-Ff(>QilAdx-)GP7=UK1vBbi)E{2*1RpW`#!?aobm_UhC)Q#Wf2 zi*jYQwL@?^LjO%ulv-iXxwblWdVv$yoRwDJ)pyW8Z<|8zU6{7qG~2E`974&C0u3gS zYLY@SRZ>Go#x#AAc?c#$gbpW_wPm|hC_zSk3`dpnG>4PdGa-bcErfr~LI7C@gorp2 zx5OUBI4dHK2`PcHq@wh5LhFTs=xu8>m?#wMZfRB>yGO zQa=cngku8RFZ;e{G0vy{m<~^-fP@3uyyU=4JOGP-zF8M?4_^Gy>Uv&X7!!!p!kIYs|zxpErn@zty}yj>zr@@l@r)5US7 zs_9<~`M^8TCf!ulr5C14tF>*~vUn0Ji2Cnd{QhJA^V`=i!6G!))bBc%@7}mId6LCA vtvcH{yRGXO79ftKlE?5Ea$n2)2l+hPM|emPQu`CSoz2lRwqUmhzumt8M?qhr literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/expressions.bin b/crates/krishiv-plan-codec/tests/golden/v4/expressions.bin new file mode 100644 index 0000000000000000000000000000000000000000..989e70d581087785685dc60e9e1eee165ddfed74 GIT binary patch literal 1470 zcmbu9&u`pB6vsW|U*p-`?AuAwxK1eUf>5*)(q{9cB?u@zp^6IwQYC~$t*n!^>e|uT zt5PmqAr-YZ4jg)e3kMEdfD;l@5&bt17Y;peh=jzk=8boDm9)Jv*1UY?y?Jk(@AE9P zfCl#|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--4e-{y literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/file_scan.bin b/crates/krishiv-plan-codec/tests/golden/v4/file_scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..e2c43d89a626f568fdf7fdcc49f4f01c3ca43e8e GIT binary patch literal 495 zcmZWmy-ve05I(zZ5{_c}(DVjV;)s2+;NFOS1WXHFY*y&BRk?kFX;pOD&I(DHA zTf4tU>v~m*(KwBCBTIV2INBLm*AE(!X=Vd1tA+)2$u-MxsJAML+0LSzPQ}Xi7sc$j z^~@IlKw<)jIa3GliFibuAZWhG^Lz^{-3MUG6qTrkx*8@0FPMlS9OC7#&Z8X^5;w`Y zC#h&JX(B`u)(`;8g|jSrALr5=d^)2LMV7h@=Tf)yPiN>dFizg&E1-cDz#uTlwjhhO zU032P?mEo|Wox^d-M5vho+rJ&3m9wTtDmQ-(_R!in+))Ef$vv3+M0ecbmq=x2EQX6 c$I+=TSu_8r_|vC+_Wt%-+xG0x4jwZ11JlY#-v9sr literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/filter.bin b/crates/krishiv-plan-codec/tests/golden/v4/filter.bin new file mode 100644 index 0000000000000000000000000000000000000000..86164355bce9b6f8ca0aa8fd026278b800b0f87c GIT binary patch literal 967 zcmb_a&2H2%5cW8KUUy4VZxO-b5bX&?;<6`H30U=n5ElYRB(R$7ZfHoV#H+%M14!k@ zfd}9PcmyPnc$gk|hcJ_r6$K;?U~49s`DVWHj4j-d>r*b@a1p47@`2oLk4=`NFJAE5WKzy3gw@S3>%&ZHhyyrqZ+ZQC_L5>WXxan>tfFC#K0)^}5Y(^pM^X zpulJ@ycqOIa0M+IqxezcNHB?%&`PhUs=S#A$$^J%1p9?>6v6uGm@y`BGv=Oa#sIB7 zln`~uF0mcNC@Mr9p~NT{#jl|jY>enupBs~VNbfl;8V6%>7VP$Zc6LW|#k8Ne0z%MM zHwq6Fji%4U=hwH;^IuOhABs{lo72QRc9M7ePfb&cUr|m`<|t%8?z)beqUXnqM z*+{zg^MzGV4!>rh;KfWzFD;fUs{v*c>85pU6+T#NiW1)#<2RiHW#Uy)+4#h?aj~rG zCeQ9gc)Q&fCwld^`}XPM2Y(uv%9Jl#n{Hott*K5TIV#$!w$-{UBbWl)=Yr4SE~KvF r*Z0y%vMb>>B3Nw$T}~#_F;zxx2R81%l6#xbU&*~-1N}@c{2%V`YR+;5 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/hash_join.bin b/crates/krishiv-plan-codec/tests/golden/v4/hash_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..cc01eb7ecd2d1f720f91eef30d256db9e9c94927 GIT binary patch literal 1030 zcmbtSJ#W-N5Zzh(&U=%DIGaRqED9?WM3%ydqk%(k(G?IO8V*V%SYlt!iercKp2A0S z1xTSohz2TJ3it;UNR;p+=u)JlKW2U5MEDQ|vzl4Y%+9m+y%o00{e6y4xp3r7yn!>t zuIg{!*;HDk`4&1++*N~ABWUTOy^}|=6j&{EKiSlB1uHt=jx)I|m(bqr=aHPgO+_3fX_?%-fj2V^%Asd~{l)*JK5&2jp_M1xMe>2ahX3Q?vc&?!Qf7Mll(Ws$) zJ??4=>G)H1B|K;#RyxU6Px}B3Qm~WfTGC6`DoN>aGJf7_OBM8zOb2_a2uRwzA4k_b z0q*FdZD8)E_?^?!kLx; literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/hash_join_anti.bin b/crates/krishiv-plan-codec/tests/golden/v4/hash_join_anti.bin new file mode 100644 index 0000000000000000000000000000000000000000..d3b83ced6e44a12ab5fe0d091a5c763ca10aae18 GIT binary patch literal 879 zcmbtS!H&}~5Vf7Oai=Ytx?OEH;t=(OBJEZx7gPkh>Ioq(1dd2x8n=y_O-kZw;noY% za^t`U@B@4T2PBU0AA8t;DC4w?v;t0yHDk{_&scB9+;*9N!QegPj(7%7paKS#yV_Foy$Va^Hr&iD#&kg zDyZ}suLaa^a3*MpP9u08T0mt8Kuf44S*D612S%A|dvug=OBmRH)gpv&T!h#=O9(}; zhw7qr$t|*nn7QSm^-u#;yv1&#<)npZv&s!jc$;3;(CCa?{L<-EadLU$=Wxckcp&bB zgME0jQd%6sL3wvH7`$Fh7D^A~AUwW?_CjCT;Zm>GXM~1~k|FMoXu!gc>~G8s&+n)gsHdpLf3ewY3{!I-$w5%<#?+B@9hVVS zSz}E14D(Mn^Bhydi0K&Lb=zc_x+W(^-FBS((edru#4{5xVpnUMSEB#Flag3t*Z}*9 znrcDQ?NLqzeLDoGMagWgeM*NojVRBxz+bCnlHx}r?4C6iGMFcs4$fs6;F$74#g9Es z-OV@4TYuPGef{#e-gjgs(|4tg+BfDG`Jo4mq|9=itt zuIg{!*;HDk`4&1++*N~ABWUTOy^}|=6j&{EKiSlB1uHt=jx)I|m(bqr=aHP?DT8ZfBJ#0H>^GIh|7M<1&6r)T@mxXq|EjA9qftZq zdfe3#(($M2N_fyftaOsCp7sG6q+ln{wWOD>Rg%) z0^HF@%Nsr$9e#NCrn=&⩔MC-I+P2dB+butR+R3>uivwq%N=~7knKypkqtEy3m;q q=MXNJy7NhEF`P$BWsxacFe(2NyfQiZD|p43w0;H`)k&EqT=)+0|CyHn literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/hash_join_right_anti.bin b/crates/krishiv-plan-codec/tests/golden/v4/hash_join_right_anti.bin new file mode 100644 index 0000000000000000000000000000000000000000..8ce1e2de32fdbf4b813b30f01944ebb8bca7cbdd GIT binary patch literal 1030 zcmbtSJ#W-N5Zzh(&U=%DIGaRqEDGx=h%AK@M+1l8qAMUmG#r#ju*6QziercKp2A0S z1xTSohz2TJ3it;UNR;p+=u)JlKW2U5MEDQ|vzl4Y%)V#sn-#We`}-W9a^c9EcmrpO zUDe;dv!S#~^DT6ucvB5hji9BA_D&wfQed^v{bWPSI#zVP9cOY`E}^~K&m%dzr}}X( z9~ALw+483ZNKif(l>q#j;1UW-YQe40LQo+h6j({JEbcc2a^SA*K`RlKG;ROxG-FKA zV$42@i~(f}*onW)nGpAgN9(h1cr}#;Vf=yF*s>BASJcgGXnv%}6IB@35Iy-3l zIqnD{FUgDOkPi0;v6joYRGe!y8utbp+p%t_MtJ5&Yf;W0U3;oXwz4={-Br5xw=0(y z?#zhc(4PD9+1O#*U+-RV!E@@ILCAqMX>Zlg0&YoO!_H}8`KLtF9R?l*!x!IRWk~R| zj~Cuw{rc!h`@-v2#@@EqdBDTx{7-5bZ9b9jldqAR{Pt)xGU}M~BnSrOZbFrms?svX z%V>g8+^F^PILn9ZCF%J3RYL_VGp`)x|&e>2Za&6r)TF|MHef2XU6qgg}y zdc3J6q~lMuDd9mAvC>VpdfEqQ(geGCt|h&6t&)@;C*$X>j#NP}$#k%%ihxF&_v7fA zC%_$jw7lW7(cy=8Z>n`iWh#AG=Omb?iwnL6>(I3&Us>qR rhjR#*sY8F8v=+m8v{V+EvIUd!KglbTqrZ|@oJs3va#5X>3E{$b2@aW` literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/hash_join_right_outer.bin b/crates/krishiv-plan-codec/tests/golden/v4/hash_join_right_outer.bin new file mode 100644 index 0000000000000000000000000000000000000000..af55fe2bff5bea6e0548635b681d4894995f4a99 GIT binary patch literal 1030 zcmbtSJ#W-N5Zzh(&U=%DIGaRqEDGx=h%AK@M+1l8qAMUmG#r#ju*6QziercKp2A0S z1xTSohz2TJ3it;UNR;p+=u)JlKW2U5MEDQ|vzl4Y%)V#sn-#We`}-W9a^c9EcmrpO zUDe;dv!S#~^DT6ucvB5hji9BA_D&wfQed^v{bWPSI#zVP9cOY`E}^~K&m%dzr}}X( z9~ALw+483ZNKif(l>q#j;1UW-YQe40LQo+h6j({JEbcc2a^SA*K`RlKG;ROxG-FKA zV$42@i~(f}*onW)nGpAgN9(h1cr}#;Vf=yF*s>BASJcgGXnv%}6IB@35Iy-3l zIqnD{FUgDOkPi0;v6joYRGe!y8utbp+p%t_MtJ5&Yf;W0U3;oXwz4={-Br5xw=0(y z?#zhc(4PD9+1O#*U+-RV!4q}PAmqTBw6|(#0kuf z#|!VTetq<$ec|;hV{hB*JmBGT{wKAJHlN7%$=AqDetR?;8FkEg5(I;CH=#;ORcRUH zWi-JkZq#~t9OoDgV@72`%w|poWq1uHA|FqQ{WhiXznN#IX3Q?v7*|mKztdI3(X63; zJ>Jw3(($L-l<=U5Sm`EPJ?#TDX@cE6*OFekR!K^alkxLbN2;KgWIEVWML?s?`*C#5 z6X1?MTHf&4=q#j;1UW-YQe40LQo+h6j({JEbcc2a^SA*K`RlKG;ROxG-FKA zV$42@i~(f}*onW)nGpAgN9(h1cr}#;Vf=yF*s>BASJcgGXnv%}6IB@35Iy-3l zIqnD{FUgDOkPi0;v6joYRGe!y8utbp+p%t_MtJ5&Yf;W0U3;oXwz4={-Br5xw=0(y z?#zhc(4PD9+1O#*U+-RV!K>CegOCGj(%!0}1>BOnhMm*G@=uATI}AJshA+Os%8=k^ zA1}PW`t{M1_J!B4jJ<8I^MHrX`JdD>+I%A4Cto8s`R&nYWYjU|Ne~Rm-GnMBRi$N& zm(c{HxKZonahzi~j2V>$F`GFVl;Jg)h|ZDNXMUQQ^JELVx^mG^|TMrqzQKOTuXZCS|uqxPR7q$9jSs|lIdVi6#29KGZ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/limit.bin b/crates/krishiv-plan-codec/tests/golden/v4/limit.bin new file mode 100644 index 0000000000000000000000000000000000000000..9bf8f3a85f4ebb1b7f65cf2f5c7e248bb3f4f414 GIT binary patch literal 759 zcmbtSv2N5r5Zzh(&U^?@*f>+! z6$W{+G@HsHI31$@!RCclIBZ>09y>k9fwLz?tsm+K=$|$=*E?sXDb{w=7EiiVJXWB= z#8M3sh(}6lXvrBRZ&DAzq=?YspseenStyhsLqCH3QhA!w@%xDoLeUb!zhNPOygfol z9Ey8lhjP3tB#sCPfp*D9#7aykJMOdzhkJQf!;y0|QP;t47Z=wD@iZNW+gpPK(hK>o z6L7yP!a3oP!2Vm3aYx=yO)%-6Ohf8Eq{SvXX5uOM1mEuJ0x@>yq*U92E w!cp1Q*43M;ieL_6Un+SB&mr@*e7u*<(p`jSl%Z~aQuostdZx}f?7?mLHvu$gDgXcg literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/limit_offset.bin b/crates/krishiv-plan-codec/tests/golden/v4/limit_offset.bin new file mode 100644 index 0000000000000000000000000000000000000000..f691822722dce54d7a57cc5f9a6525ff170d4872 GIT binary patch literal 752 zcmbtSyKdA#6rDRBuV-&I#F-7kSQOS)$WofEPzs1rAw+|rM1oQ5-L+!u$evZ8qd^KC z1s}i<@Cg)1d`t@d!Q5*HqzF1@bdG(WXI@%`UVJI>gVcd}j?ZwnJGbrGt5fT2)vPhd ziJgyJ0y25c;~ zpFlj-(m+SXGmWTvlz{lPD;9>tSXg?_WvH0r5?*Es;9-u8K}lzyE5>0iRfqLdgEe%#+%9FG{R}w t)s3sSRTaSk#F5nU7#>5WjJ$u4&C`8^hZG^TKcU;{96ej-?DpWd`xjGlUr7J} literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/nested_loop_join.bin b/crates/krishiv-plan-codec/tests/golden/v4/nested_loop_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..290f90a0600955e5a429fd0595e8ed97a30e6180 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;-8K#+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`Pz>% literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/one_row.bin b/crates/krishiv-plan-codec/tests/golden/v4/one_row.bin new file mode 100644 index 0000000000000000000000000000000000000000..8413fa65c09ae3aa42fdd48e035a83d0e3883492 GIT binary patch literal 386 zcmZWiF;2uV5FGnVVhn=FP(%@hMM;^85>XDOu{92SxFY)K|<2^gi{own@Df|M?6)Wff literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/project.bin b/crates/krishiv-plan-codec/tests/golden/v4/project.bin new file mode 100644 index 0000000000000000000000000000000000000000..ee9ccc8ceec7fab75aaab3a0abdb7fbaa81386a5 GIT binary patch literal 953 zcmb_ay>8S%5Z+na=e<0FQtK5)YGtcQ|J3=%fH0>}tk4Ki|xFW|SXEb0+bvRH1&1k8q|d1mN|Ub~n_;V3{41f(z2w>NhP_Fqfh?q zQ2D`0YTs#b=5n(@88hi+w z>01L3^d_(M_xJb7@Lw-8>vD1}o6AHp^5ReOZ#8p;KS|C>W+be?Sg+SyBcG=r7`0!9 zKGoZ&%}ch(z`3~Q&sTW^j=_t2r7!!;ERM^Zl%(sm_@mcdZV_+u*pTc&LGU}Gl6UG*Dv?Pdi4 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/scan.bin b/crates/krishiv-plan-codec/tests/golden/v4/scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..0c2815dbb5a17ec35aa0282d4e5352d66a8146ab GIT binary patch literal 742 zcmbtSy>8S%5Z+n)&U=#!arS~R7KL>cvXtg3lp;i_5Te0QBEc&5-B~f~$lg<+qlpwc z3LYRYkVl|E;$c$o4ray(SKj(g;qFjT~nSoJ;#x&PYbIb>IdkbHFd6c&P`LS>rGqi_n~-8fd&&x zHA*0!D5;?(W0Jg1Jp_{?LW`r)+M-z~lpsStf`d|dn#1wCsSrYu3gO?f5J1)*AtVmP zJ+VVE&I*YmLPDS{`G{DFDMh#X+JxOhdEdj4aXeKw!EWytH&4?q@?QhF{SU$=;fTQY zo38740OwOrO!_C&n1+o>S!B;ld}to{^=@6jJy`gl>UK%4;=Wvy6NW{_UyBa2MZ}A< z_`MEwE!qG3feiInOwfN(ES-jGcr{B69~OwCtX!>K46vA?pVhU~bXlA!EBe1eKJkvV zN!F!x$+>Bha%JnL$e%|F!tT2lefr$}{C0gco(IO7>SOD&?K?Lmi_-`vWov6^H&qqE m9K?ZC@(7+m=4<(QFPo*i2!|A*X@5fZ(;0fk=Ir+1xBC~L`Cg>} literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/sort.bin b/crates/krishiv-plan-codec/tests/golden/v4/sort.bin new file mode 100644 index 0000000000000000000000000000000000000000..86d3bbfd4c07ef2109fa1af4cae3df5aca110f50 GIT binary patch literal 758 zcmbtSy>8S%5Z+nuo%bde;_L-sEDGx?WGT&6=n9BZAw+|rM1oc9yR%~K$lg<+qd`i# z6g&Vgz#~u~@h~ZP2Qy;_qzF28^nLcs?ELMnR)LHzBz}_GH!tx8?zU&PIel|toh|D% z`dPlR+tML86QFumXSvZhZe3HHI5WqQt55UF+&A}7oi%l4_RejSZ|ZHEANH^szCUhwwg!BjtFiul@brF0P+M(9a*h#|R> z|7w8C{UBTrjtH#3?z)Z(@_wp;asOo+Q}HngFWEB_55Wz5z8&Xt4qn{2z8RXMI4{R! zgkf;;*Wg2K5%NLsQoRguE$RPzV3u+$Ca9k0D`%h_e(lP@`vu}CE!G8S%5Z+nuo%be}#MukNSQOS($P&#}C^szCUhwwg!BV{?&*Zx6o7uQdsz2(%%5RB3gK!`CU zpX9$<;DSF0SA=5%>#w`6&8LG6d zouTLBY*EsU3;D=f8XIql%Ejlljf-_vH+l9fq?7Kxd*Q{$?&r6!U&eFaR<``mx^(x} z+exDWUswX3#e8NwXIk<{`So##lgx6St^mj1qDH?oDkxI;fMs=*h#FocJOZ0_Ld7$ zxpCkF_yInF0}@C0j~@CDVb)2JD&WMd=B;ODXV!0Kx$QFllEHh%9q|%gKwB+jIe#;j zT4u!*oH&`tWu}2@;evfy#EIZgSGr8cTJ)f%i+PfZ1F;YGqAX&ub0*7VRxDL=ScUv1 zr-DkK@mfIr24{jQG#bIH&;lw$0IHyt=6O;MI55gw+oQvjTf&6>w=F^l$3=*}MM5ZA zyQnUrOYV>zjG0p|qK6uw;wg3)k&_lin^kUL+?({OhDKx5;#W?m@{_A0p9$g@aLT!O zEFOV_y?D1wv^ao$^Nb*qxiPTCdgRxQ&E-xpyIt&xzyEX{RrCRKowDas^%=5gw-zggbu z!}|K`m(TT{BXgNuC>^zLjV$s*4;rb;3!N{s%%eR@hK#dg`jketVE1;TZrB0(1S5Fc R@6cY@1xx0!aa;7p{R_IBe6s)m literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v4/write_files.bin b/crates/krishiv-plan-codec/tests/golden/v4/write_files.bin new file mode 100644 index 0000000000000000000000000000000000000000..c5329870990ae179027f96f54f01b74b46c4b251 GIT binary patch literal 748 zcmbtSy>8S%5Z+n)&U=#!arS~R7KL>cvP5$gN{Ssf`c6(}?lQ+l4naZv( z$cv@fR1U%E5dHTyFSNpO>zeY|={b&^Jt=B^U*ALjw6VF~IWtYMwwtzi)Q93N1sY5& z)hK~@qNIkFj7jn;^$<*o2rZ7vx-Oc9LJ2bTBiJvMr#T$IoeCioZ6W*{76QoHBZS1E zxFdEb##te8L`VpfB_9wgF{S8MUz@OdDDQhXGLEL|I@s;~;`$(#??eLWh5VN+i~Jy* z6AlS%zwWw@WjLQYW70pF#w0nW%}e&o#7Ae@%iDDU_u$3HRX3w@75C+uoG^?l{u+1C z77>rn%kTG4*OL9epUF^<#RUE5#nNf0hF7!H@L_>C%F5N+#Q=*L`kA#()0J_itmrQb z`Pe(sCRvxYOU_K2l&jh{MgA;O5O&|a=+l?(=eMs{<9T3eQ+;e*wtefqWN{kdq-<;J r>P=NeFbA Date: Fri, 23 Oct 2026 10:17:59 +0200 Subject: [PATCH 08/16] exec: the build side a hash join prefers, from its inputs' statistics `planner/orient.rs` decides which input of a semi, anti or left outer hash join should be built: a pure function of the two inputs' estimated table bytes and where the estimates come from (`InputStats`), so the same function can later take statistics merged from materialized exchanges. The right input stays the build when it is small enough to broadcast; otherwise the left one is preferred when the right costs more than twice as much (both estimates from table statistics) or eight times as much (either from default estimates). A tie keeps the right input. Inner joins keep the join orderer's choice, a null-aware anti join (NOT IN) is never reoriented, and a hint decides as before. The joins that build their preserved side do not run yet, so the plan keeps the right build; EXPLAIN shows the preference on the join, e.g. `prefers RightSemi (probe: 600000 rows / 25.2 MiB, statistics; build: 59986052 rows / 1373.0 MiB, statistics)`. With TPC-H SF10 footer statistics it marks q04, q13, q21 (both subquery joins) and q22. Design: P2-T166 (PR 1a, part 3). Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/src/planner/build.rs | 17 ++ crates/krishiv-exec/src/planner/join.rs | 40 ++- crates/krishiv-exec/src/planner/mod.rs | 1 + crates/krishiv-exec/src/planner/orient.rs | 258 ++++++++++++++++++ crates/krishiv-exec/tests/planner/main.rs | 1 + .../krishiv-exec/tests/planner/orientation.rs | 136 +++++++++ 6 files changed, 452 insertions(+), 1 deletion(-) create mode 100644 crates/krishiv-exec/src/planner/orient.rs create mode 100644 crates/krishiv-exec/tests/planner/orientation.rs diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index 99b51cb2e..09b04a934 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -631,6 +631,23 @@ impl<'a> Planner<'a> { krishiv_optimizer::estimated_rows(plan, &stats) } + /// Whether every table `plan` reads reports statistics, so its estimate + /// does not rest on default row counts. + pub(crate) fn has_statistics(&self, plan: &Arc) -> bool { + use krishiv_optimizer::Statistics as _; + let stats = super::ProviderStatistics(Arc::clone(&self.cx.tables)); + let mut stack = vec![plan]; + while let Some(p) = stack.pop() { + if let PlanNode::Scan(scan) = &p.node + && stats.table(scan).is_none() + { + return false; + } + p.for_each_child(&mut |c| stack.push(c)); + } + true + } + fn scan( &mut self, plan: &Arc, diff --git a/crates/krishiv-exec/src/planner/join.rs b/crates/krishiv-exec/src/planner/join.rs index ba9a626eb..8982d6808 100644 --- a/crates/krishiv-exec/src/planner/join.rs +++ b/crates/krishiv-exec/src/planner/join.rs @@ -31,6 +31,7 @@ use krishiv_plan::{ use super::build::{Planned, Planner, unsupported}; use super::explain::ExplainNode; +use super::orient::{InputStats, StatsSource, orient}; use super::properties::Distribution; use crate::error::Result; use crate::join::JoinType; @@ -131,8 +132,23 @@ impl Planner<'_> { let build_exprs: Vec> = keys.iter().map(|k| Arc::clone(&k.1)).collect(); let (probe, probe_keys) = self.with_columns(probe, &probe_exprs)?; let (build, build_keys) = self.with_columns(build, &build_exprs)?; - let build_plan = if build_left { &logical.0 } else { &logical.1 }; + let (probe_plan, build_plan) = if build_left { + (&logical.1, &logical.0) + } else { + (&logical.0, &logical.1) + }; let build_rows = self.estimated_rows(build_plan); + // Which side the statistics say to build (design P2-T166). The + // joins that build their preserved side do not run yet: the plan + // keeps the right build and shows the preference. + let preferred = if choice.distribution.is_none() { + let probe_stats = self.input_stats(probe_plan, &probe.schema); + let build_stats = self.input_stats(build_plan, &build.schema); + let other = orient(join_type, &probe_stats, &build_stats); + (other != join_type).then_some((other, probe_stats, build_stats)) + } else { + None + }; // A hint decides, else the build's estimated size. The null-aware // anti join reads whether the whole build has a NULL key, so its // build is always broadcast. @@ -200,6 +216,13 @@ impl Planner<'_> { for p in published { detail.push_str(&format!(", {p}")); } + if let Some((other, probe_stats, build_stats)) = preferred { + detail.push_str(&format!( + ", prefers {other:?} (probe: {}; build: {})", + probe_stats.describe(), + build_stats.describe() + )); + } let (probe_schema, build_schema) = (Arc::clone(&probe.schema), Arc::clone(&build.schema)); let plan = HashJoinPlan { probe: Box::new(probe.node), @@ -245,6 +268,21 @@ fn residual_of(probe: &Planned, build: &Planned, residual: &[&Expr]) -> Option { + /// The statistics of a join input planned as `schema` from `plan`. + fn input_stats(&self, plan: &Arc, schema: &arrow_schema::Schema) -> InputStats { + let rows = self.estimated_rows(plan); + let source = if self.has_statistics(plan) { + StatsSource::Statistics + } else { + StatsSource::Default + }; + InputStats { + rows, + bytes: estimated_bytes(rows, schema), + source, + } + } + /// A join without an equality conjunct: a nested-loop join building the /// right input. fn nested_loop( diff --git a/crates/krishiv-exec/src/planner/mod.rs b/crates/krishiv-exec/src/planner/mod.rs index a5f109b16..50f715249 100644 --- a/crates/krishiv-exec/src/planner/mod.rs +++ b/crates/krishiv-exec/src/planner/mod.rs @@ -32,6 +32,7 @@ mod explain; mod file_table; mod join; pub(crate) mod nodes; +mod orient; mod properties; pub(crate) mod pushdown; mod runtime_filter; diff --git a/crates/krishiv-exec/src/planner/orient.rs b/crates/krishiv-exec/src/planner/orient.rs new file mode 100644 index 000000000..9ebce474e --- /dev/null +++ b/crates/krishiv-exec/src/planner/orient.rs @@ -0,0 +1,258 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Which input of a hash join builds the table (design P2-T166). +//! +//! The planner builds the right input by default, as Spark's hash joins +//! do. A semi, anti or left outer join whose right input is much larger +//! than its left is better built on the left input, as a right semi, right +//! anti or right outer join: the probe marks the build rows it matches and +//! the join emits the build rows by their marks at the end. Those joins do +//! not run yet; until they do, the planner keeps the right build and +//! `EXPLAIN` shows the preference (`prefers RightSemi (probe: ...; build: +//! ...)`). The decision is +//! a pure function of the two inputs' statistics ([`InputStats`]), so the +//! same function serves plan-time estimates now and statistics merged from +//! materialized exchanges later; it never reads process state, so every +//! task of a stage that calls it with the same statistics agrees. +//! +//! - A null-aware anti join (`NOT IN`) is never reoriented: it reads +//! whether the whole subquery side holds a NULL key. +//! - Inner joins keep the join orderer's choice. +//! - A build the planner would broadcast (at most +//! [`BROADCAST_MAX_BYTES`](super::join::BROADCAST_MAX_BYTES)) stays. +//! - Otherwise the left input is built when the right one costs more than +//! [`STATS_FACTOR`] times as much (both sides estimated from table +//! statistics), or [`DEFAULT_FACTOR`] times (either side from default +//! estimates). A tie keeps the right input. + +use crate::join::JoinType; + +use super::join::BROADCAST_MAX_BYTES; + +/// How much larger the right input's table must be to build the left one, +/// when both estimates rest on table statistics. +pub(crate) const STATS_FACTOR: f64 = 2.0; + +/// The same factor when either estimate rests on defaults. +pub(crate) const DEFAULT_FACTOR: f64 = 8.0; + +/// Where an input's estimate comes from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum StatsSource { + /// Some table below has no statistics: default row counts. + Default, + /// Every table below has statistics (row counts, column statistics). + Statistics, +} + +/// What is known of one join input. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct InputStats { + /// Estimated rows. + pub(crate) rows: f64, + /// Estimated bytes of the hash table built from this input. + pub(crate) bytes: f64, + /// Where the estimate comes from. + pub(crate) source: StatsSource, +} + +impl InputStats { + /// `rows / MiB, source` for `EXPLAIN`. + pub(crate) fn describe(&self) -> String { + let source = match self.source { + StatsSource::Default => "default", + StatsSource::Statistics => "statistics", + }; + format!( + "{:.0} rows / {:.1} MiB, {source}", + self.rows, + self.bytes / (1024.0 * 1024.0) + ) + } +} + +/// The join a hash join with a right build becomes when it builds its left +/// input instead, if it can. +pub(crate) fn flipped(t: JoinType) -> Option { + match t { + JoinType::LeftSemi => Some(JoinType::RightSemi), + JoinType::LeftAnti => Some(JoinType::RightAnti), + JoinType::LeftOuter => Some(JoinType::RightOuter), + JoinType::Inner + | JoinType::NullAwareAnti + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightOuter => None, + } +} + +/// The join type of a hash join whose right input builds by default, given +/// its inputs: `join_type` itself (build right), or the variant that builds +/// the left input. +pub(crate) fn orient(join_type: JoinType, left: &InputStats, right: &InputStats) -> JoinType { + let Some(other) = flipped(join_type) else { + return join_type; + }; + if right.bytes <= BROADCAST_MAX_BYTES { + return join_type; + } + let factor = + if left.source == StatsSource::Statistics && right.source == StatsSource::Statistics { + STATS_FACTOR + } else { + DEFAULT_FACTOR + }; + // NaN compares false: no flip. + if right.bytes > factor * left.bytes { + other + } else { + join_type + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const MIB: f64 = 1024.0 * 1024.0; + + fn stats(bytes: f64, source: StatsSource) -> InputStats { + InputStats { + rows: bytes / 32.0, + bytes, + source, + } + } + + fn known(bytes: f64) -> InputStats { + stats(bytes, StatsSource::Statistics) + } + + #[test] + fn a_much_larger_right_input_flips_semi_anti_and_left_outer_joins() { + let (small, large) = (known(20.0 * MIB), known(500.0 * MIB)); + for (t, want) in [ + (JoinType::LeftSemi, JoinType::RightSemi), + (JoinType::LeftAnti, JoinType::RightAnti), + (JoinType::LeftOuter, JoinType::RightOuter), + ] { + assert_eq!(orient(t, &small, &large), want, "{t:?}"); + // The other way round keeps the right build. + assert_eq!(orient(t, &large, &small), t, "{t:?}"); + } + } + + #[test] + fn inner_null_aware_and_flipped_joins_keep_their_build() { + let (small, large) = (known(1.0 * MIB), known(5000.0 * MIB)); + for t in [ + JoinType::Inner, + JoinType::NullAwareAnti, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::RightOuter, + ] { + assert_eq!(orient(t, &small, &large), t, "{t:?}"); + assert_eq!(flipped(t), None, "{t:?}"); + } + } + + #[test] + fn a_broadcastable_right_input_stays_the_build() { + let left = known(1.0); + assert_eq!( + orient(JoinType::LeftSemi, &left, &known(BROADCAST_MAX_BYTES)), + JoinType::LeftSemi + ); + assert_eq!( + orient(JoinType::LeftSemi, &left, &known(BROADCAST_MAX_BYTES + 1.0)), + JoinType::RightSemi + ); + } + + #[test] + fn the_factor_depends_on_where_the_estimates_come_from() { + let right = known(100.0 * MIB); + // Exactly the factor is a tie: the right input stays the build. + assert_eq!( + orient(JoinType::LeftAnti, &known(50.0 * MIB), &right), + JoinType::LeftAnti + ); + assert_eq!( + orient(JoinType::LeftAnti, &known(49.0 * MIB), &right), + JoinType::RightAnti + ); + // A default estimate on either side needs a ratio above 8. + for (l, r) in [ + (StatsSource::Default, StatsSource::Statistics), + (StatsSource::Statistics, StatsSource::Default), + ] { + let right = stats(100.0 * MIB, r); + assert_eq!( + orient(JoinType::LeftSemi, &stats(20.0 * MIB, l), &right), + JoinType::LeftSemi + ); + assert_eq!( + orient(JoinType::LeftSemi, &stats(12.0 * MIB, l), &right), + JoinType::RightSemi + ); + } + } + + #[test] + fn unusable_estimates_keep_the_right_build() { + for bad in [f64::NAN, f64::INFINITY] { + let left = known(bad); + assert_eq!( + orient(JoinType::LeftSemi, &left, &known(100.0 * MIB)), + JoinType::LeftSemi + ); + } + assert_eq!( + orient(JoinType::LeftSemi, &known(1.0), &known(f64::NAN)), + JoinType::LeftSemi + ); + } + + /// The decision is a function of its inputs only: any order of calls + /// over the same inputs gives the same answers. + #[test] + fn orientation_is_deterministic() { + let sources = [StatsSource::Default, StatsSource::Statistics]; + let types = [ + JoinType::Inner, + JoinType::LeftOuter, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::NullAwareAnti, + ]; + let mut cases = Vec::new(); + let mut x: u64 = 0x9e37_79b9_7f4a_7c15; + for i in 0..512 { + // A fixed xorshift sequence of sizes from 1 B to ~1 TiB. + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + #[expect(clippy::cast_precision_loss, reason = "test sizes")] + let (a, b) = ((x % (1 << 40)) as f64, ((x >> 20) % (1 << 40)) as f64); + let t = types[i % types.len()]; + let l = stats(a, sources[i % 2]); + let r = stats(b, sources[(i / 2) % 2]); + cases.push((t, l, r)); + } + let first: Vec<_> = cases.iter().map(|(t, l, r)| orient(*t, l, r)).collect(); + let reversed: Vec<_> = cases + .iter() + .rev() + .map(|(t, l, r)| orient(*t, l, r)) + .collect(); + assert!(first.iter().eq(reversed.iter().rev())); + assert!(first.iter().any(|t| matches!( + t, + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightOuter + ))); + assert!(first.iter().any(|t| matches!( + t, + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftOuter + ))); + } +} diff --git a/crates/krishiv-exec/tests/planner/main.rs b/crates/krishiv-exec/tests/planner/main.rs index 1e43b7c07..2c23ec442 100644 --- a/crates/krishiv-exec/tests/planner/main.rs +++ b/crates/krishiv-exec/tests/planner/main.rs @@ -13,6 +13,7 @@ mod exchanges; mod hints; mod listing; +mod orientation; mod physical; mod runtime_filters; mod semantics; diff --git a/crates/krishiv-exec/tests/planner/orientation.rs b/crates/krishiv-exec/tests/planner/orientation.rs new file mode 100644 index 000000000..5e0499770 --- /dev/null +++ b/crates/krishiv-exec/tests/planner/orientation.rs @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Which input a hash join prefers to build (design P2-T166): a semi, anti +//! or left outer join whose right input is much larger than its left +//! prefers to build the left one. Those joins do not run yet, so the plan +//! keeps the right build and `EXPLAIN` shows the preference. + +use std::collections::HashMap; +use std::sync::Arc; + +use krishiv_exec::planner::{MemoryTable, MemoryTables}; +use krishiv_optimizer::TableStats; +use krishiv_plan::analyzer::testing::context; +use krishiv_plan::analyzer::{MemoryCatalog, ResolutionConf, TableColumn}; + +use super::{Env, table}; + +/// `t` (10 rows), `big` and `mid` (the same rows, statistics claiming 100 +/// and 5 million), and `bare` (no statistics: a default estimate of one +/// million rows). +fn env() -> Env { + let (schema, batches) = table(10); + let columns: Vec = schema + .fields() + .iter() + .map(|f| TableColumn { + name: f.name().clone(), + data_type: krishiv_types::DataType::from_arrow(f.data_type()).unwrap(), + nullable: f.is_nullable(), + }) + .collect(); + let mut catalog = MemoryCatalog::new("spark_catalog"); + let mut tables = MemoryTables::new(); + for (name, rows) in [ + ("t", Some(10)), + ("big", Some(100_000_000)), + ("mid", Some(5_000_000)), + ("bare", None), + ] { + catalog = catalog.with_table(&["default"], name, columns.clone(), None); + let stats = rows.map(|rows| TableStats { + rows, + columns: HashMap::new(), + }); + let table = MemoryTable::new(Arc::clone(&schema), batches.clone()).with_statistics(stats); + tables = tables.with_table(&["spark_catalog", "default", name], Arc::new(table)); + } + Env { + ctx: context(catalog, ResolutionConf::default()), + tables: Arc::new(tables), + } +} + +/// The `HashJoin` lines of the plan of `sql`. +fn joins(env: &Env, sql: &str) -> Vec { + let plan = env.plan(sql).unwrap_or_else(|e| panic!("{sql}: {e}")); + plan.to_string() + .lines() + .filter_map(|l| l.find("HashJoin").map(|i| l[i..].to_owned())) + .collect() +} + +#[test] +fn a_much_larger_right_input_is_shown_as_the_side_not_to_build() { + let env = env(); + let cases = [ + ( + "SELECT a FROM t WHERE EXISTS (SELECT 1 FROM big WHERE big.a = t.a)", + "HashJoin LeftSemi, [a#0] = [a#3], RuntimeFilter #1: a#0 -> spark_catalog.default.t.a, \ + prefers RightSemi (probe: 10 rows / 0.0 MiB, statistics; build: 100000000 rows / 1144.4 MiB, statistics)", + ), + ( + "SELECT a FROM t WHERE NOT EXISTS (SELECT 1 FROM big WHERE big.a = t.a)", + "HashJoin LeftAnti, [a#0] = [a#3], \ + prefers RightAnti (probe: 10 rows / 0.0 MiB, statistics; build: 100000000 rows / 1144.4 MiB, statistics)", + ), + ]; + for (sql, want) in cases { + assert_eq!(joins(&env, sql), [want], "{sql}"); + } + let outer = joins( + &env, + "SELECT t.a, big.b FROM t LEFT JOIN big ON t.a = big.a", + ); + assert_eq!(outer.len(), 1); + assert!(outer[0].starts_with("HashJoin LeftOuter"), "{outer:?}"); + assert!( + outer[0].contains(", prefers RightOuter (probe: 10 rows"), + "{outer:?}" + ); +} + +#[test] +fn null_aware_inner_small_and_hinted_joins_show_no_preference() { + let env = env(); + for sql in [ + // `NOT IN` reads whether the whole subquery side holds a NULL. + "SELECT a FROM t WHERE a NOT IN (SELECT a FROM big)", + // Inner joins keep the join orderer's choice. + "SELECT t.a FROM t JOIN big ON t.a = big.a", + // A build small enough to broadcast stays. + "SELECT a FROM big WHERE EXISTS (SELECT 1 FROM t WHERE big.a = t.a)", + // A hint decides. + "SELECT /*+ SHUFFLE_HASH(big) */ t.a FROM t LEFT JOIN big ON t.a = big.a", + ] { + let js = joins(&env, sql); + assert!(!js.is_empty(), "{sql}"); + assert!(js.iter().all(|j| !j.contains("prefers")), "{sql}: {js:?}"); + } +} + +#[test] +fn a_default_estimate_needs_a_larger_ratio() { + let env = env(); + // 5M rows against a default 1M: 5x is enough with statistics on both + // sides (above 2x), not with a default estimate on one (8x). + let sql = "SELECT a FROM bare WHERE EXISTS (SELECT 1 FROM mid WHERE mid.a = bare.a)"; + assert!( + joins(&env, sql).iter().all(|j| !j.contains("prefers")), + "{sql}" + ); + let sql = "SELECT a FROM bare WHERE EXISTS (SELECT 1 FROM big WHERE big.a = bare.a)"; + let js = joins(&env, sql); + assert!( + js.iter().any(|j| j.contains( + "prefers RightSemi (probe: 1000000 rows / 11.4 MiB, default; build: 100000000" + )), + "{js:?}" + ); + let sql = "SELECT a FROM t WHERE EXISTS (SELECT 1 FROM mid WHERE mid.a = t.a)"; + assert!( + joins(&env, sql) + .iter() + .any(|j| j.contains("prefers RightSemi")), + "{sql}" + ); +} 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 09/16] feat(plan-codec): carry the Expand node (format v5) PhysNode::Expand (ROLLUP, CUBE, GROUPING SETS) and the NullOf project item it uses made the codec's matches non-exhaustive. The proto gains ExpandNode (layout, projections, nullability, subqueries), ExpandProjection and ProjectNullOf. The writer format version becomes 5. A plan raises its min_reader_version to 5 only when it holds an Expand or a NullOf item, so other plans stay readable by older readers and an older reader never meets a node it does not know. The decoder refuses an expand without a projection and a projection whose item count differs from the nullability list. The v1 to v4 goldens still decode; the v5 goldens add three Expand shapes (a rollup with a grouping id, a single projection, a computed column beside a nulled one with a subquery). The TPC-DS ratchet shrinks by the four grouping-set queries (q05, q14, q18, q22) that now plan and round-trip. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-plan-codec/src/decode.rs | 35 +++++++++ crates/krishiv-plan-codec/src/encode.rs | 38 +++++++-- crates/krishiv-plan-codec/src/lib.rs | 2 +- .../src/proto/krishiv.plan.v1.pb.rs | 42 +++++++++- .../tests/common/inventory.rs | 65 ++++++++++++++++ .../tests/golden/v5/aggregate.bin | Bin 0 -> 1730 bytes .../tests/golden/v5/aggregate_final.bin | Bin 0 -> 848 bytes .../tests/golden/v5/exchange_broadcast.bin | Bin 0 -> 754 bytes .../tests/golden/v5/exchange_hash.bin | Bin 0 -> 758 bytes .../tests/golden/v5/exchange_range.bin | Bin 0 -> 764 bytes .../tests/golden/v5/exchange_round_robin.bin | Bin 0 -> 754 bytes .../tests/golden/v5/exchange_single.bin | Bin 0 -> 754 bytes .../tests/golden/v5/expand_rollup.bin | Bin 0 -> 887 bytes .../tests/golden/v5/expand_single.bin | Bin 0 -> 762 bytes .../tests/golden/v5/expand_subquery.bin | Bin 0 -> 952 bytes .../tests/golden/v5/expressions.bin | Bin 0 -> 1470 bytes .../tests/golden/v5/file_scan.bin | Bin 0 -> 495 bytes .../tests/golden/v5/filter.bin | Bin 0 -> 967 bytes .../tests/golden/v5/hash_join.bin | Bin 0 -> 1030 bytes .../tests/golden/v5/hash_join_anti.bin | Bin 0 -> 879 bytes .../tests/golden/v5/hash_join_partitioned.bin | Bin 0 -> 1030 bytes .../tests/golden/v5/limit.bin | Bin 0 -> 759 bytes .../tests/golden/v5/limit_offset.bin | Bin 0 -> 752 bytes .../tests/golden/v5/nested_loop_join.bin | Bin 0 -> 866 bytes .../tests/golden/v5/one_row.bin | Bin 0 -> 386 bytes .../tests/golden/v5/project.bin | Bin 0 -> 953 bytes .../golden/v5/project_without_nullable.bin | Bin 0 -> 392 bytes .../tests/golden/v5/replicate_rows.bin | Bin 0 -> 879 bytes .../tests/golden/v5/scan.bin | Bin 0 -> 742 bytes .../tests/golden/v5/sort.bin | Bin 0 -> 758 bytes .../tests/golden/v5/top_n.bin | Bin 0 -> 760 bytes .../tests/golden/v5/union.bin | Bin 0 -> 870 bytes .../tests/golden/v5/window_aggregates.bin | Bin 0 -> 1790 bytes .../tests/golden/v5/window_ranking.bin | Bin 0 -> 860 bytes .../tests/golden/v5/write_files.bin | Bin 0 -> 748 bytes crates/krishiv-plan-codec/tests/ratchet.rs | 6 +- crates/krishiv-plan-codec/tests/roundtrip.rs | 72 +++++++++++++++++- proto/krishiv/plan/v1/plan.proto | 28 ++++++- 38 files changed, 270 insertions(+), 18 deletions(-) create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/aggregate_final.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_range.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_rollup.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_single.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/expressions.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/filter.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join_anti.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/limit.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/one_row.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/project.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/replicate_rows.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/scan.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/sort.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/top_n.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/union.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/v5/write_files.bin diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 552f70c87..80b6d2d0d 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -299,6 +299,20 @@ impl Decoder { count: self.expr_ref(r.count)?, subqueries: self.subqueries(&r.subqueries)?, }, + Kind::Expand(e) => { + expand_shape(e)?; + PhysNode::Expand { + input: self.boxed(e.input)?, + layout: layout(&e.layout), + projections: e + .projections + .iter() + .map(|p| p.items.iter().map(|i| self.item(i)).collect()) + .collect::>()?, + nullable: e.nullable.clone(), + subqueries: self.subqueries(&e.subqueries)?, + } + } Kind::Sort(s) => PhysNode::Sort { input: self.boxed(s.input)?, keys: sort_keys(&s.keys)?, @@ -364,6 +378,10 @@ impl Decoder { name: e.name.clone(), expr: self.expr_ref(e.expr)?, }, + Kind::NullOf(n) => ProjectItemPlan::NullOf { + index: usize_of(n.index)?, + name: n.name.clone(), + }, }) } @@ -545,6 +563,23 @@ fn aggregate(a: &pb::AggregateNode) -> Result { }) } +/// An expand has a projection, and every projection has one item per output +/// column. +fn expand_shape(e: &pb::ExpandNode) -> Result<()> { + if e.projections.is_empty() { + return Err(corrupt("an expand without a projection")); + } + if e.projections + .iter() + .any(|p| p.items.len() != e.nullable.len()) + { + return Err(corrupt( + "an expand projection whose items do not match the nullability", + )); + } + Ok(()) +} + fn window_plan(w: &pb::WindowNode) -> Result { Ok(WindowPlan { partition_by: usizes(&w.partition_by)?, diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index b1fa442f6..9c90308fd 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -243,10 +243,29 @@ impl Encoder { input: self.node_index(input)?, }), PhysNode::Exchange(exchange) => Kind::Exchange(self.exchange(exchange)?), - PhysNode::Expand { .. } => { - return Err(CodecError::Unsupported( - "an Expand node (not in this format yet)".to_owned(), - )); + PhysNode::Expand { + input, + layout: l, + projections, + nullable, + subqueries, + } => { + // An older reader has no `Expand` node to read. + self.min_reader = self.min_reader.max(5); + Kind::Expand(pb::ExpandNode { + input: self.node_index(input)?, + layout: layout(l), + projections: projections + .iter() + .map(|items| { + Ok(pb::ExpandProjection { + items: items.iter().map(|i| self.item(i)).collect::>()?, + }) + }) + .collect::>()?, + nullable: nullable.clone(), + subqueries: self.subqueries(subqueries)?, + }) } }; Ok(pb::Node { kind: Some(kind) }) @@ -281,10 +300,13 @@ impl Encoder { name: name.clone(), expr: self.expr(expr)?, }), - ProjectItemPlan::NullOf { .. } => { - return Err(CodecError::Unsupported( - "a NullOf project item (not in this format yet)".to_owned(), - )); + ProjectItemPlan::NullOf { index, name } => { + // An older reader has no such item to read. + self.min_reader = self.min_reader.max(5); + Kind::NullOf(pb::ProjectNullOf { + index: *index as u64, + name: name.clone(), + }) } }), }) diff --git a/crates/krishiv-plan-codec/src/lib.rs b/crates/krishiv-plan-codec/src/lib.rs index a451f1fff..0bf831492 100644 --- a/crates/krishiv-plan-codec/src/lib.rs +++ b/crates/krishiv-plan-codec/src/lib.rs @@ -54,7 +54,7 @@ pub use encode::{encode, to_message}; pub use error::{CodecError, Result}; /// The format version this crate writes. -pub const FORMAT_VERSION: u32 = 4; +pub const FORMAT_VERSION: u32 = 5; /// The `min_reader_version` this crate writes: readers older than this /// cannot understand the plan. diff --git a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs index 286966df8..7d11c857d 100644 --- a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs +++ b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs @@ -3,7 +3,7 @@ // This file is @generated by prost-build. #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalPlan { - /// The writer's format version (4). + /// The writer's format version (5). #[prost(uint32, tag = "1")] pub format_version: u32, /// The oldest reader version that understands every field in use. @@ -55,7 +55,7 @@ pub struct ExplainNode { pub struct Node { #[prost( oneof = "node::Kind", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16" )] pub kind: ::core::option::Option, } @@ -94,6 +94,10 @@ pub mod node { /// v4: a plan with a window needs a v4 reader. #[prost(message, tag = "15")] Window(super::WindowNode), + /// v5: a plan with an expand (ROLLUP, CUBE, GROUPING SETS) needs a v5 + /// reader. + #[prost(message, tag = "16")] + Expand(super::ExpandNode), } } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] @@ -130,7 +134,7 @@ pub struct BoolList { } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ProjectItem { - #[prost(oneof = "project_item::Kind", tags = "1, 2")] + #[prost(oneof = "project_item::Kind", tags = "1, 2, 3")] pub kind: ::core::option::Option, } /// Nested message and enum types in `ProjectItem`. @@ -141,6 +145,9 @@ pub mod project_item { Column(super::ProjectColumn), #[prost(message, tag = "2")] Expr(super::ProjectExpr), + /// v5: only inside an expand; a plan holding one needs a v5 reader. + #[prost(message, tag = "3")] + NullOf(super::ProjectNullOf), } } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] @@ -157,6 +164,14 @@ pub struct ProjectExpr { #[prost(uint32, tag = "2")] pub expr: u32, } +/// v5: a column of nulls in the Arrow type of an input column. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ProjectNullOf { + #[prost(uint64, tag = "1")] + pub index: u64, + #[prost(string, tag = "2")] + pub name: ::prost::alloc::string::String, +} #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct SubqueryPlan { #[prost(uint64, tag = "1")] @@ -190,6 +205,27 @@ pub struct ReplicateRowsNode { #[prost(message, repeated, tag = "4")] pub subqueries: ::prost::alloc::vec::Vec, } +/// v5: each input row once per projection (ROLLUP, CUBE, GROUPING SETS). +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ExpandNode { + #[prost(uint32, tag = "1")] + pub input: u32, + #[prost(uint64, repeated, tag = "2")] + pub layout: ::prost::alloc::vec::Vec, + /// The projections: the same output columns, each computed one way. + #[prost(message, repeated, tag = "3")] + pub projections: ::prost::alloc::vec::Vec, + /// Whether each output column is nullable. + #[prost(bool, repeated, tag = "4")] + pub nullable: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "5")] + pub subqueries: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ExpandProjection { + #[prost(message, repeated, tag = "1")] + pub items: ::prost::alloc::vec::Vec, +} #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct SortKey { #[prost(uint64, tag = "1")] diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index 123fa6257..2f673d7af 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -307,6 +307,34 @@ fn join(join_type: JoinType) -> HashJoinPlan { } } +/// `ROLLUP (a, b)` over a two-column input: the key columns are kept or +/// nulled and the last column is the grouping id. +fn expand_rollup() -> PhysNode { + let kept = |index, name: &str| ProjectItemPlan::Column { + index, + name: Some(name.to_owned()), + }; + let nulled = |index, name: &str| ProjectItemPlan::NullOf { + index, + name: name.to_owned(), + }; + let id = |n: i64| ProjectItemPlan::Expr { + name: "grouping_id".to_owned(), + expr: lit(Literal::Long(n)), + }; + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![ + vec![kept(0, "a"), kept(1, "b"), id(0)], + vec![kept(0, "a"), nulled(1, "b"), id(1)], + vec![nulled(0, "a"), nulled(1, "b"), id(3)], + ], + nullable: vec![true, true, false], + subqueries: Vec::new(), + } +} + /// One node of every variant, each as the root of its own plan. pub(crate) fn nodes() -> Vec<(&'static str, PhysNode)> { let keys = vec![SortKey::asc(0), SortKey::desc(1)]; @@ -383,6 +411,43 @@ pub(crate) fn nodes() -> Vec<(&'static str, PhysNode)> { subqueries: vec![subquery()], }, ), + ("expand_rollup", expand_rollup()), + ( + "expand_single", + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![vec![ProjectItemPlan::Column { + index: 0, + name: None, + }]], + nullable: vec![false], + subqueries: Vec::new(), + }, + ), + ( + "expand_subquery", + PhysNode::Expand { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + projections: vec![ + vec![ProjectItemPlan::Expr { + name: "shifted".to_owned(), + expr: call( + Builtin::Add, + vec![col(1, "a", DataType::Integer), lit(Literal::Integer(1))], + DataType::Integer, + ), + }], + vec![ProjectItemPlan::NullOf { + index: 0, + name: "shifted".to_owned(), + }], + ], + nullable: vec![true], + subqueries: vec![subquery()], + }, + ), ( "sort", PhysNode::Sort { diff --git a/crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin b/crates/krishiv-plan-codec/tests/golden/v5/aggregate.bin new file mode 100644 index 0000000000000000000000000000000000000000..35a39060157d4948f6e7d071a38ec8e4588b88ea GIT binary patch literal 1730 zcmcgt&2G~`5Z+n;ZMG?OHwBEv!OE$M#3d)x15`bsiVLB~BD$)b#Hwpo_C^giE=c9Z zfd}9PcmxheJOUE0&;#!fX5zL=)bJNXXEiJD%#J_L&bMRX_&hx0@(mZBdM=;I)#^l- zM=uYx(OEH&UYzvxC^Hg_@}+%T#EBBptxTC78r7DLDUOm{-BWj^eNq;&T07NcGAu?_ z@?h!;R|F`~As0>rVMA~Q6)qalUdNK4J5oX=oixvrvMVG9?%M%urNUC=4iB1)F@b|I z`(kDcaJ7Nz93+G_k8~jH8tGr^6#TqS zl#BC6xm;)B%*ly57yA0yeCBh>h3nreC%h#&@wa!37+4|$xn)QE1_fC+YLJK zscdq*Y47+VI#TQSf_jF!jY{kN$z(E9J{h5{mf7;vbR_#F`QOjS{r1Ad7+W_ncCCrk zmyRtYXLS{~eDl-zj&4KRPm{h;kkya6uVCDj(&?pxp$P%HEot|P!YKR^7@cO=yBU9I z?I;}$)7(U-x{A_4UX)3^8wl_xpRM5a+sXG&A3wNlPv<&&Rhi!E++Uy@#08aLZh)116NHrAB& z61^;6n009ptoG6USZBFbIBsoI99uodk*!biNMHOmCf;CSTR-Hhdmw)?oadJ!%Tv+$jQ2;!sr zFmnwN0DlRK+n<4w*%GMc0K1I;Ro+Fwjr=CPB;6-v|83WG+{!gLMsy!!Ol^(nu*e}q zmdC?n-)`4=jKRVqQH0DxR~)X$5a`I9rzdi3xCmLX8-^I literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_broadcast.bin new file mode 100644 index 0000000000000000000000000000000000000000..44186479a51888279684b067d8549bebdeb9267d GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M?28vloQ1D6FfHr8HNe6cD9Chz3K61gqG0XT{i&y{AA&gOqeB zcmQ62N1#CBVN&o8X2uRk5p?Y8`|SDsJnL0nDB}x>pQH-(OMHQ|?##5OZ%&LgRkOw* zFIHw-Sp=&?bl)|3p%wa_ZOap@7wFmMw5auceGlEUw#oJGxoL|{v+at5Ar$W@&|qSz zyaeKjk{UWPCdsSRK`<#Ibm*0JU9?Mu5@hH`a9Aowb2xrG6+$T5Lbx|91dz2)2#G^+ zN96;98cACus8U{^^;h&zM91{773(R@?YXC z^@DIhI3lq9y6<}y<9zCm$?#Dra&Ee$T-QxoO*I<*{xfX#c70- uva1_gZ>uVT1&BkbPYel}0{5FS&6)c%C-rgL;mowM74-|k;p=wGY= literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_hash.bin new file mode 100644 index 0000000000000000000000000000000000000000..87676f17bdb8c6e75c1bcc692a3a0b666b25651d GIT binary patch literal 758 zcmbtSy>8S%5Z+n;?af_?vloQ1D6FfHr8HNe6cD9Chz3K61gqG0XT{i&y{AA&gOqeB zcmQ62N1#CBVN&o8X2uRk5p?Y6`|SDO*w@1c9vHo4wCH*K+Lwq0>BxZQL5dv3__)B}^@%j8pYpM;kjn2Cqr4Sc?x7jO+;ym56iG-q*L&dCYm z;Nq{r$KE31gW%;3GUj!p|L=hrs&>5_6?H*JwWixh8S%5Z+na=e@~=bM}HT7KL>cvXtg3l!8R55Te0QBEc&5-B~eqWbY}^(I6#V ziabDGAdf(S#KWZE9n6d!kRs^V)%V#mv-7*Fyg)|h62D62>o@osx7$-TRlh5)+c$TAL~cxo;G!+r)Q?g*Y&2&_j*&br9gv; zr1D~j#!6~vNf^g(69>U0h|r=}R8`(A6iSeR8^V5}9L?tF!$b(7Xba)qu@FGgE+HTe z!~-!U8=ne@LqbfTQ}PM15)-m+^}aEy2lqaPCE;kIZvCBJFK!Pa-2`eHO{gpx0=fzU zh#|R@|0;u%{veza4hby3>$;B9@_EXHasOo+Q3)g3ykyTz+zOX*b>GkD7`(WAb=Nw3 zaa{Jv2t)hgul9%BBII%Ka{E2_wIu)V*bK#3jM06SFRg}h_%%xnpB9K-TCCPK0$9w@ zP3zifdP3F|CEdD^&zvJ|;&oBk_{_9%v8w7O&t8T!=kB``e!1v=e!ICI&3#ju@^fp` z?R)Phjgk<@MO)Rj+LUDoa}fJd$wPPnsjKDF-E@}hAUr1vW&2}#n9R^IRmN%ueye{0 D;~ro$ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_round_robin.bin new file mode 100644 index 0000000000000000000000000000000000000000..c4f0eabdd9a59858b29e2e6c03d64928e4e5353c GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M*I;vloQ1D6FfHC7P>H3KFG4hz3K61gqG0XT{i&y{AA&gA}?H zJOD4iBTz)j!=&IH%#0n7BIwxF_u2FNdDg4EP{tP$KT8$pSNIaQyED_CzCAJ4RLu&5 zyjYq|Wf80n(S6_Kg;wZywk=Pro}*`*)1uZ7^nG;C+9ub#=cX;z&890J521KVfd&&x z!MvKlpsSlf`d{yn#1wCsSrZZ7Q(${A%LuXLP#8n z9kEL>&I*YmLPDS{`H)zNDMh!2+JxO>`Ow3XaXeKw!QS8(H;1umeYF$IxkwZWE zscXUo;fTQYo4)T^jPt2KCc~4-Ct;s9FF7z155VGI@74v}gBO3ax*eFSxG&e_gmGx` z*U+Q3i1-kAxx*goIt)x;WO9 z75!@=pE$?bB{PF#pKMzc8s*jz`w(s1UEKVbw vlwIA}dQ(*q%t0JTC6C}4WUiKv_On^Khwzjlr1mFtFP))d>YUvU{C58WYo1@V literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin b/crates/krishiv-plan-codec/tests/golden/v5/exchange_single.bin new file mode 100644 index 0000000000000000000000000000000000000000..b9c7afa741c7cac571b6fd349b8e0010bb9f5a35 GIT binary patch literal 754 zcmbtSy>8S%5Z+n;?M*Jk*$cv06xLP9Qktt!iV&qjhz3K61gqG0XT{i&y{AA&6Df2Q zJV0I`k3fOM!=&IH%#0n7BIwxF_u2FNdDg4EP{x-MFQf|eD}0ID-I-}m-<}w2s%C{j zUM$U~vItg(=)Q0ALM!w;+mLnz)-puxma zc?rZ5B{g(pOp@2BgJ4ob=+G)CCJc?;Gk5F=5YLODuht9g>Y|K2q0^Z5E6&t zp4g!nXNANOAt6wfd_=6ol%iWhZNl!ceCT1xIG(DTV0Z9~o2RjAeKn1x773(ZIdPHa71AHP2cw{#`)AAli|taldw;lmmHXh2Vn89ck2S~!HYjy-44uE+?Q)|!Z@_} zYv@s1M0^Om++mM(9oheTScZBmCg{E>mR3VGzM7?m4+}&uD_3h911x6fW=&%?T^wu5 zivG2bPn=_Il66_zSJfK?K`(7i_-`v uWmh-0-c(fta}Wno$s>3MnXBdFy=<25A{k~D>H(^r5aL4Mh*Xx^Nvs+>vNviG7ledV zZX9?3UVum7fW*V}z&nH)yDF-nCw4WTJwM;9zws(B^x`v#AEgTPb9{!|&C1lrheyVm zqFP{(W^=PHEP~Y`x^Jp9(+d5@*7=duWAtowoR#{vzJ>0pu2MZXF?F`A)=hT5O~qRZ zG?-W_FM+tHq=tr^Uh-n%AeadvH0b4JnblK;5@hH`u$L=Gb2@(27eXl7Lbw+w1dz8w z77~Wyh8R$et3twvEFq&Rd52JmK4rIBZo=V?-uAHM9Q4(BFl^)E{89Yw1DtC5iuf*i z=+K8_fF3J&jYB=atwmj}S9!TO%u^cNpQ8A>BA<)muG6(e!Po&12~0lAf9Yj)KgrI> z_Q}|Px!G)3AJ>y4Oxh=dPg;H2ykv_EJTun&SD&s;=asD$+|gxq06TSbm`StzVSAQIs(iCqSJKMf;QD^Za!d~8# zl`Ypr5y2S5o>X!l9>dJl^4`vDG#Mg1qzs*QnYuX{p<~LF!wy`AiLon~w<}xpe*8S%5Z+na=e@}tarS~R7KL>cvP5$gN-Sv^P3*2j6JAL<9_o-}o)w@*!zuj@^lKj}@;mI4hX zlFExA8Y`)xC1D)DNgM=|AVP~?QB`@fP$)qLZU}pYax|Nx_Y)z6qAi4b%R&H2yM%x^ z5ckA3+1M)}4hb=VyyPQdB_?Fu>SJS858-_dOTytq-S|7bUfk?Q!8uGdbrHEJQOrTa zkbIHv-;fLBHM9c*DN)BTp)UBv0B>*U@=2C zt!u040a;U&bmu}obq=+O*F|OHQ`5%9s;Zkjdl4!Kx@#vqyX<~meg8I^`=&DGr`D#^ zJFh2=k`Tv5Th+GOlw}BXJ$q8g19%RptL5X}be8NOJR=KL`(wJF%+N7a#%c$CtA7C; CJzw$w literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin b/crates/krishiv-plan-codec/tests/golden/v5/expand_subquery.bin new file mode 100644 index 0000000000000000000000000000000000000000..1f9e9d7f857fbc0e6b84da5b9dd8dc000ed138bf GIT binary patch literal 952 zcmb_a&5G1O5bo+Elb)I#C)>MXY&?WHNyN*XWI-MFB;v)LR^UrvJN_I!s<1=i{VJv-B2gtW~Ub?w~f@o^ip3z1m+{W zg7=`QCNsjJ9$?rmi@D3QC^SPHn7;m{7z3nF)z{ZI3H4tqo~{eyA3m!%k%^bSlYeUr zx8Xa{G0~KWb~mOCSbg*2mv2`;zI^`F8;7>C<%`zMhUYeV zX3`k@MO)RbIxNc=#vpd2l2f5*)(q{9cB?u@zp^6IwQYC~$t*n!^>e|uT zt5PmqAr-YZ4jg)e3kMEdfD;l@5&bt17Y;peh=jzk=8boDm9)Jv*1UY?y?Jk(@AE9P zfd=;}zekxXKI9+pR&gASZ+?0Zl~I}>ad#LWMw7JUq!hrdTlp{+j5|d+P7X@Z<#w6h zj5Bdwta0mjoDW6&b~KJh`J{+nuR}qD5kaDWGCL%J$EY9$F1+xgp2Z> z4;ZI}g5{I_gqcExgBwl5Fc?0DVQFg^1Xn9K0WBCyMjK;l7ohcVhB$E--9XEXCPo`| zZm8U|^twigi>poc#O>66?wCA`5FvkIO=C}=FjyteC$6Z*N;~uQe@e-Vyoe=lI z*R%!aH4Pxpv<;!A9k{4z7rL7E;FXZ{zM?OZUy1z~=Y5=4aH{nuRaGgu)qYf%hxK1& z9F*R{M-5tAWmF55{9m8Vb5#vBRNI(txyn*?Rn8FStnJ2gZJ(WO`06uI!>XTIUV-cX z`?*9|W8iV?{rIpHB%NV&D9C36&TW|-jmm(K!7{gGo|gizP8lUBe(eUmY+e;nI7+fI zyd4!`a+KxcczDfc1ghUmfBK;M`_~`7b-HepMd_WQl&v%OMg~2fdr6VyWj0AupLB__ zPZ_;J-XqcybaO>6_d1-sjS)O;5nAmnb2G|@%59Q`8#A~en}Tz^%B%se`CL^2B5h&% z6_Jjx!l#uf@eXWagQ~V{iAMjaLN$nVh5Z()G1c(su51YJUG(5MxhpS-jcWl1)v!l@ zO#Z$1=AG}Z|Md80F+Z#=TCy?v{rfb{!;^0w)@436jmlLg-}vXR>#E#nIz6dXQSWw4 zBp2*+hphwgsRXg~1!=J9!yo_r#YT_ytYCb=*qWfs@iB41^c0)VraZ6l`f0_^XuST? zB6M2g3mXgft0=CE|F8+R@Ab~9i(X7_LC2?EvbIV3FXK)}w!1Cn2|J&Z@i?HYDcAN_ j@x42Y3pYxWsECK?ZiPgy^aOVf@(h2NeG_w+RPX--6Rp<0 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin b/crates/krishiv-plan-codec/tests/golden/v5/file_scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..50724761fef2212a8852bd0822410443c91392fa GIT binary patch literal 495 zcmZWmy-ve05I(zU!k}@Ff^nG_H`@2uFMYYNOMBWIYj@6NLCuVv(i>%s~Tj~B`5M0i#t|J?m zpw<38+EnY(nM{*VHPU!6j>5f(u|vNhm}17~qN zt!KUj00I-hSuk|~pNL1qDT3yUJkNKq(tQ9Xby10GsHRj4EA#jse zc#?|tk|snnVFLlMS~!cM_i!$~!J{(@r^r&5;#}%h{^>Mb2FA&od>v?D1<(i#vMtD9 z?bM}n9`&taqp~%9#qK*&R(tVaXnn>y`0BkRvC0i2YvK{UF7Sg&R~bEw$5!8&RO5Ff c(?6FU-2Mz^w|weO_m|lHj+tDGsfgf@kZayi!-y6>yK6I#W9*rpZ_Jy3KF&klvD9 zaifLwVy;J0DBhwmiXSBoaFYP=7QCXW@@6JMa1pu@-!G)22-Z)>j4_FuG51_E#?jhC z2~mga65BzHqC(UWN{oV0;u>nn#)xk9xiPti^q#|_aWIx=!EWzoXLqzvO#6u|dC0ZZ zjlu&(qv9Tdm77;#1D{g%opsm#3~0 r*Z0y%vJ3n+B3Nw$T}~$8m?|T;!#D1~l6#xbU&*~-1N}@c{U7e{Zj^Eb literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin b/crates/krishiv-plan-codec/tests/golden/v5/hash_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..ddfde900465a97dd93e490ca1dcf1663622258d3 GIT binary patch literal 1030 zcmbtSzi-n(6ux)1Fix424Oz#t!VovQ7>< zUe0$I6PLtAZX*x(N3jyCe5E+oZZ_|YHn(Hdl+EDGkJhqSI=c2;k!)pgw7x6V;BQwh zFWjAHq;ui*D{XH(J@9Gpoc@U|&CMtDee?}!SLoc#GXe%Sjv2R-NnX@HV}!q0kw(`K8^h?|HWpr(M*|~sb=z_BN5{8s6IV|_i&?L2UWNYuPD)~punFdq zWU2(swnsS?^lb=WMd@s=JW4~HMpP6^;ICCO&G4fUcF!0K>Ce+#`RB6oaZE*-#E)H0 zoy|AHU4PhIef{#eF|cJWvv-w>IydGQd4UVfw8{&Wud>Xg14>4Wvr`;XWD0htAN7JR W(1(cNYrmoWpa+J`V;wf=b@&$mJbirt literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin b/crates/krishiv-plan-codec/tests/golden/v5/hash_join_partitioned.bin new file mode 100644 index 0000000000000000000000000000000000000000..048d0341b790566fb4f1ddffe0fdfd5ede05a9fa GIT binary patch literal 1030 zcmbtSzi-n(6ux)1Fix424Oz#t!VovQ7>< zUe0$I6PLtAZX*x(N3jyCe5E+oZZ_|YHn(Hdl+EDGkJhqSI=c2;k!)pgw7x6V;BQwh zFWjAHq;ui*D{XH(J?hioIsFq`nwwAP`{*0!I=?-cOf)A#7qRnumcYdTk8{=e!<#AsEy zc|Gnc0qOKp_60m>ac+gl)Nv8ZgS@9r|`2YFw^k#js%m*IXN7nhHdSvHAw*9IwMpXI+! z!2Ny_E(k{i_Fs{VJMw;Nf@%L`8d3KVEjHOR6Hmb>_;OPhat$^>Ag`nk1E)8FGv zRnrF-@}YOEO}eQXm!6w0t=5fg%i?LQAlm)#;t!v8zrKI_I$DILG4=b-<+B^VPM%~j wj;pS*uG!Xg3=0s4QpqED0=ci{{eyg-?IS#*40ZdHx}D9@Gflx^53a+%0R$>&D*ylh literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin b/crates/krishiv-plan-codec/tests/golden/v5/limit_offset.bin new file mode 100644 index 0000000000000000000000000000000000000000..d55e972c842f15a43345d50a9d77725d3eeb3e44 GIT binary patch literal 752 zcmbtSyKdA#6rDSM&E9N?GaH1lD6FlJr8HZi6d_855DkVB2}ZGZ*NU+tdscytCQ|4q z_<;OCK7j&>k4eEln0xJj6hX&~&avp?FJy1`|u= zB@j=P)Xl4hYF zgiFE^f$cYa-?J3wQ+G^;FOyHQK208SU?yHZOMksv7jO?AykvE|U9RH3T$2;VEsMXl zJZg)GkAR0e%(1Q``+twhP>;m~-515uYN*Cvv()fmf#_xBYHeeH#SGo7X{@GaV@+An zO$+(NIo2jwm$gmKO_!9bx@n92d88ohzdO;V&;8GD*H`{LFtw>Zb~f9-^JcO*jc`(S tbz|#IRYfodaUhjEf@hGqT0Y*(X6Y`%Aw@{-Pw0L+L&wxPyB+xL{smTIUrGP~ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin b/crates/krishiv-plan-codec/tests/golden/v5/nested_loop_join.bin new file mode 100644 index 0000000000000000000000000000000000000000..0f49b3f684db38f274f53e354754e1257d3ef418 GIT binary patch literal 866 zcmbtSy>8S%5Z+na=k?^m`SyY^7KL>cvXtg3lmenu2+?4u5E86n-<=g>NA{jVx->|k zqu>E}0Um(@iHAwSI~+5<=tKaCf?a)|J%8Wqj;*{<#+MS`OXch5_zZWuQ`4TjJ~q}= z%>wNqj3~xgA#p@V2$Uu75i2pI=uTgouzOS9YgjUlhU&^6_kMBpD3)5pU#-_`?#%hr43qxJG@w%qX!DXiGx2TsB%f~9`K-Z< zZ=kMk&sD6;H92830rAHKHoZl}6XE6ddsEku{lAklq%jN7eOk<|hHCR_<{IA35WTEi zENu)ho1mLDjn(uQSyNW@;e~wQ9BGp*%i1Pqrc262-LysiBvKHrzd6yxhxL!IUp^0} zzNt-h-q~#T#xIn`X@o)9)s3xJRTaS$#DP@u5FSJ3YI%Pzo1|lehZLc1+tBTFf{v+k cb~~_j|5M!Crv5JO`Pii!N#-+}ZkuRfS z$fPWS7>I9weL3G?Xu70#ne5`b>}iMvPKdx!O4eyEX`i!E@*6{?ntc=Pv+wYq&+%sn z$=RH%uumOuL;}Yq@zm=*qi))Du%e$wzqaC`HfFi%Hrr4MQ6C$%8pmMsrVYO9r%FOz qEUv6~+inQ%-cPRE4C7tfT$frD@oP~&y~Ll-kN4~>_k-^rr|=8UmMiH1 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/project.bin b/crates/krishiv-plan-codec/tests/golden/v5/project.bin new file mode 100644 index 0000000000000000000000000000000000000000..cb7da8071df616756d098deac40e248a474db82e GIT binary patch literal 953 zcmb_ay>8S%5Z+na=e<0FQtK5)YGtcQ|J3=%fH0>}tk4Ki|xFW>gSKb0+bvRH1&1k8q|d1mN|Ub~n_;V3{41f(z2w>NhP_Fqfh?JNXtN;K2 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin b/crates/krishiv-plan-codec/tests/golden/v5/project_without_nullable.bin new file mode 100644 index 0000000000000000000000000000000000000000..91115ce8affabf703ee8dab06fc19e701b060b4c GIT binary patch literal 392 zcmZWiF;2uV5FGn_i7^Nw-xX0rVO>d?iV{&Gs;jTDl-ymUD3ORIpreG2f(P&coAdHKh zB_Urz#gI{11Thfb{`ykB!O(O`A2Qj+ciB@BGc1X~SxnY(&S{^rk@Fiv#hQEr?UV2H zpU?4U2+7%;t8h#ma6|&fM)B0kBcpEXZLp%9M!&Vy=wD2oh#f@ zpuvPvc@cyoB{j69jH2hUgJ5DrXwfUGDsK{n5~S}2uwN)gvpak>7D6c6Lb#VG1dz5z z;uHGfn%E&9pYjOPtBM9ZSG(TQoJVo}vi zo;?oeao0bb;N8dduWw(!45mXGqXjO~W|AQPEbltyX0jz!b#3RB{H7Aa%98 ey_ZhnU4#eZp==x1^>~7gsWLV@urU+kuKOE#TV@6T literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/scan.bin b/crates/krishiv-plan-codec/tests/golden/v5/scan.bin new file mode 100644 index 0000000000000000000000000000000000000000..8d48e76e5683a6f6a4c8cde8f70b6fc8c4904fb2 GIT binary patch literal 742 zcmbtSy>8S%5Z+na=e@~=ID0`Di^94JSxR#iN&!(SglI67NU(~1cUH_gviB6|Xpllj z!2|FDJOTw050ip-Ff(>QilAdg-)GPNj=d^~WO6C-gH)k@j?ZwnJ2&mwt5f4l-K;Sz z%9YvH4#DXN{WncfYK4=|wbiN93k+OyR$6^m-$DPpZ3?}2VcK%jY`gMs2qilTG?+xH zAcbVAq=t@+Y5F4b5KM*$9R`)PWxG@;K}LQIN0st4hm+ScA%r3o!oOxAfUE;TL>!4* zVvk~+6%ogTlt5YX9mWTvjd{lPD;9%Y~8zXo#qAB0Q7F@f!u zec$r{&ZnN34o{{D4V#d%$bp&o&^++-&AO0#u<$|E^^#o0eYvI}jEjoD79D4cm=|a9 z2OaA=vj6u38S1f|qW`p9ISuvrYE~NFE)j#gT5nteu$-fxH;vPDS)8eA`oBUx@J_Tz zH8S%5Z+na=e@~=ID0`Di^94JSxR#ix&oq95u%Ahi3F?IcW1@ek-eutM}w4f zDR=;0fJdM}%EP4K9n6d!kRs^V(f8Riv-7vR$_r$4De<#ZzCOWY+-}cJbN2SsSX0(3 z^s{_vHl;V(^mXyPZy7hN^ySRN3O%|wS2-;N;Kn%&X z{8s~9?uKwlI3Te8rt3N`$or`V#{HLRM8!uWykyTzJOnrJ<$j#cIe2m7>TYO`;=CM_ z5r)CVUxN>|MaT!i%k5=|Yf1m#12dFkK1TOdzO)+3;nyrRe3&D8X|Y<{2w*-%H?3={ z>G4=olyu)hK64JWiPuGC;|tTq#j2{CJbM|^Id?yt@Y7ZI^ZU22qnU3iQ+{l1x_$5E zq)`&$xM-`|R-3X6VFqF$l{|nWNL?+T?xoXY7vVX1sM;Ub!(@t%sWLV@@Z08S%5Z+na=e@}#arS~R7KL>cvP5%Llp;h|Aw+|rM1oc9yR%~K$lg;(mnKru zQSbn~0FOX{#KWZE9n6d!kRs^V(f8Riv-7vR$_r$4E%8b!U%$pzxZR$a=Jef(v8Jq7 z=x6!TY)Xq@b%5@NI?J`fQEQvx#OfJ(wm!`({YXDV_pGTiy?bt&d|hwa{AurswiIYE zkyKs`(O5|hEh*#pP2wP!1QA;FimJ++xk3pta6>pOl%v@ly`Kmn6p0Y-9SZ@Z?Gpmx zKKI3&ab+9e+oD={JORv#O)c?j=wSW*@fb?fi-c5(YG+F4Ar3_-gJ0thjN z>rA8Lz`kAj!m%MjO+{=Y|ND93z^?#p~>HI&1zS!(z=NA%KSwYCw!e2Q*b z*H+W>v8E{L#)Uj^7TUz?qO$S1Y2#v5)lHtg2{X4$MqnYqGPIz%?|uF{{mwz BU%UVS literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/union.bin b/crates/krishiv-plan-codec/tests/golden/v5/union.bin new file mode 100644 index 0000000000000000000000000000000000000000..32f0afadd3ea26360418e4dff225fd2f263f5917 GIT binary patch literal 870 zcmbtS!EVz)5Zztd@n(`xXVa>##lgx6St^mj1qDH?oDkxI;fMs=*lDb|cJOXgxaERW zZXEaket=KlfW#60qlf-Om~~R53OF&VdFz?knf04l?szP?VDO%CS3HMj&`}Fn&0kHV zmU(#!Zjw&rGS@)0@W4JRlT>hMC|zX}E&AZ-a-J6AK#m>@w3XPk?N z;sLnWi#N+uiv!qK_eO)k$#ODJ^*|1yyVs#T(c4mQsqC~!lf#A7vwzX9Nsl{xwX(ZE ze;Bzs3fvP1Ed1ZliaxTx5jOO`qh6pMqZ<3=dc8IT&3WVmVZ9q@AlU{kBdkUPrfG)o zr<-}MsbR!4if@~2vP@l*lb~)IP5x-~=5FGf#v8HgwaHUx|L=e#))+Ryewt3Tp!w!g zP6d5C1aRW)bfyDJhd6+^EVaPLs%4hr`y%YVH5M|QWrYsUr3!H}WtApRd``XfH_Kmr zSYLhp^10D>Wg+u-O2?fWBa4H`hi0aVQWwiS_i3M!5##KbKBloP*qz RJG2+|z>-B`+!nod{{p@Ne6#=n literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin b/crates/krishiv-plan-codec/tests/golden/v5/window_aggregates.bin new file mode 100644 index 0000000000000000000000000000000000000000..c9a55711561e6751d820d9fdadc9e50e00a4abb3 GIT binary patch literal 1790 zcmcgt&2G~`5Z+lk$!42UcU!<%94wzuBrZ82AE4@~LR>IC7PYI|Nvs;XvNuXW98jTB zxpCkDcmW=P0}>C@1Md)Ky=j!v@{=O5t66z>Jf3g9of)%+V|%BRpHS8k&-hcmTph~t zV0T|CnH2-xij%$^Wr~9m4!2(vaUvLRR;o<*l?Zr46$eQ!?ugsmJ}iq^tc+!u42w~f z+^<#M5(5F!qpabB*Je~eg`2j&-L*K#F6U73Mw;hI*<+jnb!->HlvzTny}b@0gyAH_ zp00!d?$%KpTpe2a~_$c@my^AX&9n>z>w!S_#=JhueZZcqi=DF=G%^;d2H)bhodAM%$oM(CGBA>uF$U1U<+%o7$agMSMIgvSha` z##QhZ!i-T52s&Pv1|9m;aZ5kzp{Zc#T$4#rv|dwbKpY_mT7y`cf{^(j30{XoW*HDW zzmNXbm3LvBgQf30qclfZiInqoPm(e-yW=*KQL zz{yw3ee-_u{mbW1&7dW7nZ2x3w0z+Y5_w&hx6>*wR6fcw7Xlz5Wpo3!AhHF$w;ru^ YS2;XH4SJiYZg$tWCG%JxTX25-11?3X4gdfE literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin b/crates/krishiv-plan-codec/tests/golden/v5/window_ranking.bin new file mode 100644 index 0000000000000000000000000000000000000000..f8a20245a49e3d9fb349b5b10d66cb20427aadae GIT binary patch literal 860 zcmbtSv5wO~5ZzgClFeK$@of%aRuq;iWC;nSLRY{k6+$!^N+ejtPGZH_k-gCY6%EqS zQSbr$0G~jC#K&C0KZIH5Aa{a}UA^a>H#@VM@yL!Z^f|{5T=vv6e2Rz7+SIF;Ghl0>*HxM+iQUH5`OK;aJGNS7rMjhVqFC2eszzs~&X(1t$?mo;eIP-B z(Oh;y(0$1jG-UL{=dp)iVnk@Llb2;yPb6~Sz6jtXm!5KN{c6A%le8HVm&_QDbxiPy zeKuqx^0`|+aX<(Obc^36mTW-YgElsF@-Dp1Vab>d1;0T9Xi65!NIOALOW`fvB^2O3>fXNs|QdL&bh+2~u zG*yh>_om8(%bdD;W|}ZxlvSOjj{^z*_PZCn|G546?dz9r)H9_i-ZnNlyf&YSjsxuH uOwMJRlFf_Q!Q29;0W<)X5(FmVW{2=xufY literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/golden/v5/write_files.bin b/crates/krishiv-plan-codec/tests/golden/v5/write_files.bin new file mode 100644 index 0000000000000000000000000000000000000000..0e6927a4e89736bd79b3c5c33d9214f733087356 GIT binary patch literal 748 zcmbtSy>8S%5Z+na=e@~=ID0`Di^94JS)#cLrGO|ELNpjkBv{41I}_uL>^%iK8l=!s z@Bq93k3fOM!=&IH%#0n7BIwxF_u2FN?bxe=NG2B&KS&koXZRF%yED_CzC1C`)OL+w zQLfCkb_h;K=)bW=sTEE-*H$M^FEDWSv~2WkeGC1w))sp2+_dG!ZoBf{5K49wXfTOX zK?=!KNevwt)AV`fA(#viIt;3&Dchw&2{Q6yIINVXIh?$j2_Y11A^a;A0?0ZbM8uJ} zA@(T7SrKteNC}iB?+`07qv+01o3eW>A9^@4j%VsJ+#men@?j!hixjd?@?Wwn@`G?e zI3lq9qVIc_;e6_h>F{KlkmQ6mFF7z1ADv}CU#| Vec { diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index 62958449e..b9f9fd449 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -282,11 +282,10 @@ fn an_offset_limit_needs_a_v3_reader() { fn a_window_needs_a_v4_reader() { // A `Window` node sets `min_reader_version` to 4: an older reader has // no such node. Plans without one stay readable by older readers. - assert_eq!(FORMAT_VERSION, 4); assert_eq!(message("window_ranking").min_reader_version, 4); assert_eq!(message("window_aggregates").min_reader_version, 4); for (name, plan) in cases() { - if !name.starts_with("window") { + if !name.starts_with("window") && !name.starts_with("expand") { let m = to_message(&plan).unwrap(); assert!(m.min_reader_version <= 3, "{name} needs a newer reader"); } @@ -337,6 +336,75 @@ fn a_corrupt_window_is_refused() { ); } +#[test] +fn an_expand_needs_a_v5_reader() { + // An `Expand` node (and the `NullOf` item inside it) sets + // `min_reader_version` to 5: an older reader has neither. Plans + // without one keep the older minimum, and a v4 reader still reads them. + assert_eq!(FORMAT_VERSION, 5); + for name in ["expand_rollup", "expand_single", "expand_subquery"] { + assert_eq!(message(name).min_reader_version, 5, "{name}"); + } + for (name, plan) in cases() { + if !name.starts_with("expand") { + let m = to_message(&plan).unwrap(); + assert!(m.min_reader_version <= 4, "{name} needs a newer reader"); + } + } +} + +#[test] +fn an_expand_plan_survives_the_wire_with_every_field() { + for name in ["expand_rollup", "expand_single", "expand_subquery"] { + let (_, plan) = cases().into_iter().find(|(n, _)| n == name).unwrap(); + assert_eq!(round_trip(&plan), plan, "{name}"); + } +} + +#[test] +fn a_loose_null_of_item_still_needs_a_v5_reader() { + let plan = plan_of(PhysNode::Project { + input: Box::new(PhysNode::OneRow), + layout: Vec::new(), + items: vec![krishiv_exec::physical::ProjectItemPlan::NullOf { + index: 0, + name: "x".to_owned(), + }], + nullable: None, + subqueries: Vec::new(), + }); + assert_eq!(to_message(&plan).unwrap().min_reader_version, 5); + assert_eq!(round_trip(&plan), plan); +} + +#[test] +fn a_corrupt_expand_is_refused() { + use pb::node::Kind; + let edit = |f: &dyn Fn(&mut pb::ExpandNode)| { + let mut m = message("expand_rollup"); + for node in &mut m.nodes { + if let Some(Kind::Expand(e)) = &mut node.kind { + f(e); + } + } + m + }; + corrupt( + &edit(&|e| e.projections.clear()), + "an expand without a projection", + ); + corrupt( + &edit(&|e| { + e.nullable.pop(); + }), + "do not match the nullability", + ); + corrupt( + &edit(&|e| e.projections[0].items[0].kind = None), + "a project item kind", + ); +} + #[test] fn a_newer_writer_that_stays_compatible_is_read() { // N reads N+1 when the writer only added fields: unknown fields (here diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 5945829d1..dfa056764 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -19,7 +19,7 @@ syntax = "proto3"; package krishiv.plan.v1; message PhysicalPlan { - // The writer's format version (4). + // The writer's format version (5). uint32 format_version = 1; // The oldest reader version that understands every field in use. uint32 min_reader_version = 2; @@ -69,6 +69,9 @@ message Node { ExchangeNode exchange = 14; // v4: a plan with a window needs a v4 reader. WindowNode window = 15; + // v5: a plan with an expand (ROLLUP, CUBE, GROUPING SETS) needs a v5 + // reader. + ExpandNode expand = 16; } } @@ -98,6 +101,8 @@ message ProjectItem { oneof kind { ProjectColumn column = 1; ProjectExpr expr = 2; + // v5: only inside an expand; a plan holding one needs a v5 reader. + ProjectNullOf null_of = 3; } } @@ -111,6 +116,12 @@ message ProjectExpr { uint32 expr = 2; } +// v5: a column of nulls in the Arrow type of an input column. +message ProjectNullOf { + uint64 index = 1; + string name = 2; +} + message SubqueryPlan { uint64 id = 1; SubqueryValueKind kind = 2; @@ -139,6 +150,21 @@ message ReplicateRowsNode { repeated SubqueryPlan subqueries = 4; } +// v5: each input row once per projection (ROLLUP, CUBE, GROUPING SETS). +message ExpandNode { + uint32 input = 1; + repeated uint64 layout = 2; + // The projections: the same output columns, each computed one way. + repeated ExpandProjection projections = 3; + // Whether each output column is nullable. + repeated bool nullable = 4; + repeated SubqueryPlan subqueries = 5; +} + +message ExpandProjection { + repeated ProjectItem items = 1; +} + message SortKey { uint64 column = 1; bool descending = 2; From 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 10/16] 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 4d3c0afc2..819288ff0 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 2081ab8f0..b9cd29a42 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 552f70c87..391aa0104 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 43ad12dd6..a483c7988 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 1eab60ba4..51e28e034 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 000000000..31dd6bcfb --- /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 a451f1fff..bde8c5a05 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 286966df8..787ab2b62 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 000000000..4f74770a2 --- /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 44a870029..e266973e5 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 5945829d1..18db6879e 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 4d6377879..b8db3d153 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 2b1f50732..4407042ce 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 000000000..75366be89 --- /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 000000000..1322549c3 --- /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 1b6935f11..7f86c0180 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 11/16] 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 391aa0104..93bfee3c9 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 4f74770a2..650816bbd 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 2b3faaabfadaf55668df5fb7446a0b465797e5c2 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:42:35 +0200 Subject: [PATCH 12/16] optimizer: join conditions count the keys a filtered input keeps Join ordering estimated each join conjunct from its inputs' unfiltered column statistics, while the conjuncts that filter one input only scaled that input's rows. A dimension filtered to one year (365 of 73,049 `d_date_sk` values) joined to a fact table holding about five years of dates was then estimated at 1/200 of the fact rows instead of about a fifth. Once literal equalities on columns without a distinct count got the default selectivity, that underestimate decided join orders: in TPC-DS q07, q26, q27 and q61 the fact-and-date join looked smaller than the filtered demographics or promotion dimension, so the fact side was built and the dimension probed, and the Bloom filter the dimension build gave the fact scan was lost. Before estimating join conjuncts, each input's column distinct counts are now capped at the rows its filters keep, as the estimate of a filter already does (the containment assumption). In those queries the fact table is again the probe side of every dimension join, each dimension build filters the fact scan, and the most selective dimension joins first. TPC-H q10 still builds orders, not lineitem. New test: TPC-DS q26's shape over footer statistics (the demographics join below the date join); default-estimate TPC-H snapshots q10, q17, q21 and two join-order snapshots change. Design: P2-T166 (PR 1a, review of the benchmark A/B). Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .../planner/snapshots/planner__tpch__q10.snap | 30 ++--- crates/krishiv-optimizer/src/join_order.rs | 20 ++- .../krishiv-optimizer/tests/it/join_order.rs | 77 ++++++++++++ ...d_dimensions_count_the_keys_they_keep.snap | 16 +++ ...der__large_trees_are_ordered_greedily.snap | 114 +++++++++--------- ...__positional_roots_keep_their_columns.snap | 8 +- .../tests/it/snapshots/it__tpch__q10.snap | 22 ++-- .../tests/it/snapshots/it__tpch__q17.snap | 26 ++-- .../tests/it/snapshots/it__tpch__q21.snap | 30 ++--- 9 files changed, 227 insertions(+), 116 deletions(-) create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__join_order__filtered_dimensions_count_the_keys_they_keep.snap diff --git a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q10.snap b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q10.snap index 8cb5322a4..9399c5c6c 100644 --- a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q10.snap +++ b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q10.snap @@ -9,20 +9,20 @@ TopN 20, [revenue#37 DESC NULLS LAST] +- HashAggregate keys=[c_custkey#0, c_name#1, c_acctbal#5, c_phone#4, n_name#34, c_address#2, c_comment#7], functions=[sum((l_extendedprice#22 * (1 - l_discount#23)))] +- Exchange hashpartitioning(c_custkey#0, c_name#1, c_acctbal#5, c_phone#4, n_name#34, c_address#2, c_comment#7) +- Project [c_custkey#0, c_name#1, c_acctbal#5, c_phone#4, n_name#34, c_address#2, c_comment#7, (l_extendedprice#22 * (1 - l_discount#23))] - +- Project [c_custkey#0, c_name#1, c_address#2, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23, n_name#34] - +- HashJoin Inner, [c_nationkey#3] = [n_nationkey#33], RuntimeFilter #3: c_nationkey#3 -> spark_catalog.default.customer.c_nationkey - :- Project [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23] - : +- HashJoin Inner, [c_custkey#0] = [o_custkey#9], RuntimeFilter #2: c_custkey#0 -> spark_catalog.default.customer.c_custkey - : :- Scan spark_catalog.default.customer [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7] + +- Project [l_extendedprice#22, l_discount#23, c_custkey#0, c_name#1, c_address#2, c_phone#4, c_acctbal#5, c_comment#7, n_name#34] + +- HashJoin Inner, [o_custkey#9] = [c_custkey#0], RuntimeFilter #3: o_custkey#9 -> spark_catalog.default.orders.o_custkey + :- Project [l_extendedprice#22, l_discount#23, o_custkey#9] + : +- HashJoin Inner, [l_orderkey#17] = [o_orderkey#8], RuntimeFilter #1: l_orderkey#17 -> spark_catalog.default.lineitem.l_orderkey + : :- Project [l_orderkey#17, l_extendedprice#22, l_discount#23] + : : +- Filter (l_returnflag#25 == 'R') + : : +- Scan spark_catalog.default.lineitem [l_orderkey#17, l_extendedprice#22, l_discount#23, l_returnflag#25], PushedFilters: [(l_returnflag#25 == 'R')] : +- BroadcastExchange - : +- Project [l_extendedprice#22, l_discount#23, o_custkey#9] - : +- HashJoin Inner, [l_orderkey#17] = [o_orderkey#8], RuntimeFilter #1: l_orderkey#17 -> spark_catalog.default.lineitem.l_orderkey - : :- Project [l_orderkey#17, l_extendedprice#22, l_discount#23] - : : +- Filter (l_returnflag#25 == 'R') - : : +- Scan spark_catalog.default.lineitem [l_orderkey#17, l_extendedprice#22, l_discount#23, l_returnflag#25], PushedFilters: [(l_returnflag#25 == 'R')] - : +- BroadcastExchange - : +- Project [o_orderkey#8, o_custkey#9] - : +- Filter ((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01')) - : +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12], PushedFilters: [((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01'))] + : +- Project [o_orderkey#8, o_custkey#9] + : +- Filter ((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01')) + : +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12], PushedFilters: [((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01'))] +- BroadcastExchange - +- Scan spark_catalog.default.nation [n_nationkey#33, n_name#34] + +- Project [c_custkey#0, c_name#1, c_address#2, c_phone#4, c_acctbal#5, c_comment#7, n_name#34] + +- HashJoin Inner, [c_nationkey#3] = [n_nationkey#33], RuntimeFilter #2: c_nationkey#3 -> spark_catalog.default.customer.c_nationkey + :- Scan spark_catalog.default.customer [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7] + +- BroadcastExchange + +- Scan spark_catalog.default.nation [n_nationkey#33, n_name#34] diff --git a/crates/krishiv-optimizer/src/join_order.rs b/crates/krishiv-optimizer/src/join_order.rs index 9cd27440b..4ad247a85 100644 --- a/crates/krishiv-optimizer/src/join_order.rs +++ b/crates/krishiv-optimizer/src/join_order.rs @@ -400,7 +400,7 @@ impl Graph { .collect() }) .collect(); - let mut kept = Vec::new(); + let mut joins = Vec::new(); for (p, r) in preds.into_iter().zip(&refs) { let mut on: Vec = r .iter() @@ -418,6 +418,24 @@ impl Graph { })); continue; } + joins.push((p, on)); + } + // A filtered input keeps at most as many distinct values of a column + // as it keeps rows, as the estimate of a filter does: a join with a + // filtered input then counts the distinct keys that input keeps + // (containment), not its whole table's. Without it, a dimension + // filtered to a year (365 of 73,049 dates) joined to a fact table + // holding five years of dates was estimated 40x too small, which + // made the fact side look like the smaller input to build. + for (i, ids) in attrs.iter().enumerate() { + for id in ids { + if let Some(c) = all.cols.get_mut(id) { + c.ndv = c.ndv.min(rows[i]); + } + } + } + let mut kept = Vec::new(); + for (p, on) in joins { let sel = conditional(&p, selectivity(&p, &all), &on, &attrs, |i| &ests[order[i]]); let edge = !has_subquery(&p); kept.push(Pred { diff --git a/crates/krishiv-optimizer/tests/it/join_order.rs b/crates/krishiv-optimizer/tests/it/join_order.rs index f13de51d4..39bcf4fae 100644 --- a/crates/krishiv-optimizer/tests/it/join_order.rs +++ b/crates/krishiv-optimizer/tests/it/join_order.rs @@ -317,6 +317,56 @@ const TPCDS_SF1: &[(&str, u64, &[Column])] = &[ ("ca_country", "string", f64::NAN, f64::NAN, 1_525, None), ], ), + // Footer statistics as Parquet gives them (TPC-DS SF1): row counts and + // bounds; distinct counts only from integral bounds (`max - min + 1`, + // at most the non-null rows), none for strings. + ( + "catalog_sales", + 1_441_548, + &[ + ( + "cs_sold_date_sk", + "bigint", + 2_450_815.0, + 2_452_648.0, + 7_191, + Some(1_834), + ), + ( + "cs_bill_cdemo_sk", + "bigint", + 5.0, + 1_920_765.0, + 7_183, + Some(1_434_365), + ), + ], + ), + ( + "date_dim", + 73_049, + &[ + ( + "d_date_sk", + "bigint", + 2_415_022.0, + 2_488_070.0, + 0, + Some(73_049), + ), + ("d_year", "bigint", 1_900.0, 2_100.0, 0, Some(201)), + ], + ), + ( + "customer_demographics", + 1_920_800, + &[ + ("cd_demo_sk", "bigint", 1.0, 1_920_800.0, 0, Some(1_920_800)), + ("cd_gender", "string", f64::NAN, f64::NAN, 0, None), + ("cd_marital_status", "string", f64::NAN, f64::NAN, 0, None), + ("cd_education_status", "string", f64::NAN, f64::NAN, 0, None), + ], + ), ]; #[derive(Debug)] @@ -478,3 +528,30 @@ fn literal_equalities_without_a_distinct_count_use_the_default() { let cols = "SELECT l_orderkey FROM lineitem WHERE l_commitdate = l_receiptdate"; assert!(close(rows(cols, &footers), 1.0), "{}", rows(cols, &footers)); } + +/// TPC-DS q26's shape: a fact table joined to a dimension filtered to one +/// year and to one filtered by three string equalities. A join with a +/// filtered input counts the distinct keys the filter keeps (at most its +/// rows), so the year of dates joins about a fifth of the fact rows (the +/// fact holds five years), not 1/200 of them as the whole 73,049-date +/// domain would say. The demographics join, the more selective one, then +/// comes first and the fact table stays the larger (probe) input of every +/// join, so each dimension's build can filter the fact scan. +#[test] +fn filtered_dimensions_count_the_keys_they_keep() { + let sql = "SELECT count(*) FROM catalog_sales, customer_demographics, date_dim \ + WHERE cs_sold_date_sk = d_date_sk AND cs_bill_cdemo_sk = cd_demo_sk \ + AND cd_gender = 'M' AND cd_marital_status = 'S' \ + AND cd_education_status = 'College' AND d_year = 2000"; + let plan = support::resolve(sql, &tpcds_ctx()); + let optimized = optimize_with(&plan, &TpcdsFooters).unwrap(); + let text = optimized.root().to_string(); + let line = |needle: &str| { + text.lines() + .position(|l| l.contains("Join Inner") && l.contains(needle)) + .unwrap_or_else(|| panic!("no join on {needle}:\n{text}")) + }; + // The date join is above (printed before) the demographics join. + assert!(line("d_date_sk") < line("cd_demo_sk"), "{text}"); + insta::assert_snapshot!(text); +} diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__filtered_dimensions_count_the_keys_they_keep.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__filtered_dimensions_count_the_keys_they_keep.snap new file mode 100644 index 000000000..6c6616f51 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__filtered_dimensions_count_the_keys_they_keep.snap @@ -0,0 +1,16 @@ +--- +source: crates/krishiv-optimizer/tests/it/join_order.rs +expression: text +--- +Aggregate [], [count(1) AS count(1)#8] ++- Project [cs_sold_date_sk#0] + +- Join Inner, (cs_sold_date_sk#0 == d_date_sk#6) + :- Project [cs_sold_date_sk#0] + : +- Join Inner, (cs_bill_cdemo_sk#1 == cd_demo_sk#2) + : :- Scan spark_catalog.default.catalog_sales [cs_sold_date_sk#0, cs_bill_cdemo_sk#1] + : +- Project [cd_demo_sk#2] + : +- Filter (((cd_gender#3 == 'M') and (cd_marital_status#4 == 'S')) and (cd_education_status#5 == 'College')) + : +- Scan spark_catalog.default.customer_demographics [cd_demo_sk#2, cd_gender#3, cd_marital_status#4, cd_education_status#5] + +- Project [d_date_sk#6] + +- Filter (d_year#7 == 2000L) + +- Scan spark_catalog.default.date_dim [d_date_sk#6, d_year#7] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__large_trees_are_ordered_greedily.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__large_trees_are_ordered_greedily.snap index 236e82895..51f6c65f8 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__large_trees_are_ordered_greedily.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__large_trees_are_ordered_greedily.snap @@ -3,60 +3,60 @@ source: crates/krishiv-optimizer/tests/it/join_order.rs expression: optimized.root().to_string() --- Project [b#34] -+- Join Inner, (CAST(a#33 AS BIGINT) == c#32) - :- Project [a#33, b#34] - : +- Shift +33 - : +- Scan spark_catalog.default.t [a#0, b#1, c#2] (shared #1) - +- Project [c#32] - +- Join Inner, (CAST(a#30 AS BIGINT) == c#29) - :- Project [a#30, c#32] - : +- Shift +30 - : +- Shared #1, printed above - +- Project [c#29] - +- Join Inner, (CAST(a#27 AS BIGINT) == c#26) - :- Project [a#27, c#29] - : +- Shift +27 - : +- Shared #1, printed above - +- Project [c#26] - +- Join Inner, (CAST(a#24 AS BIGINT) == c#23) - :- Project [a#24, c#26] - : +- Shift +24 - : +- Shared #1, printed above - +- Project [c#23] - +- Join Inner, (CAST(a#21 AS BIGINT) == c#20) - :- Project [a#21, c#23] - : +- Shift +21 - : +- Shared #1, printed above - +- Project [c#20] - +- Join Inner, (CAST(a#18 AS BIGINT) == c#17) - :- Project [a#18, c#20] - : +- Shift +18 - : +- Shared #1, printed above - +- Project [c#17] - +- Join Inner, (CAST(a#15 AS BIGINT) == c#14) - :- Project [a#15, c#17] - : +- Shift +15 - : +- Shared #1, printed above - +- Project [c#14] - +- Join Inner, (CAST(a#12 AS BIGINT) == c#11) - :- Project [a#12, c#14] - : +- Shift +12 - : +- Shared #1, printed above - +- Project [c#11] - +- Join Inner, (CAST(a#9 AS BIGINT) == c#8) - :- Project [a#9, c#11] - : +- Shift +9 - : +- Shared #1, printed above - +- Project [c#8] - +- Join Inner, (CAST(a#6 AS BIGINT) == c#5) - :- Project [a#6, c#8] - : +- Shift +6 - : +- Shared #1, printed above - +- Project [c#5] - +- Join Inner, (CAST(a#3 AS BIGINT) == c#2) - :- Project [a#3, c#5] - : +- Shift +3 - : +- Shared #1, printed above - +- Project [c#2] - +- Filter (b#1 == 'x') - +- Shared #1, printed above ++- Join Inner, (CAST(a#3 AS BIGINT) == c#2) + :- Project [a#3, b#34] + : +- Join Inner, (CAST(a#33 AS BIGINT) == c#32) + : :- Project [a#3, c#32] + : : +- Join Inner, (CAST(a#30 AS BIGINT) == c#29) + : : :- Project [a#3, c#29] + : : : +- Join Inner, (CAST(a#27 AS BIGINT) == c#26) + : : : :- Project [a#3, c#26] + : : : : +- Join Inner, (CAST(a#24 AS BIGINT) == c#23) + : : : : :- Project [a#3, c#23] + : : : : : +- Join Inner, (CAST(a#21 AS BIGINT) == c#20) + : : : : : :- Project [a#3, c#20] + : : : : : : +- Join Inner, (CAST(a#18 AS BIGINT) == c#17) + : : : : : : :- Project [a#3, c#17] + : : : : : : : +- Join Inner, (CAST(a#15 AS BIGINT) == c#14) + : : : : : : : :- Project [a#3, c#14] + : : : : : : : : +- Join Inner, (CAST(a#12 AS BIGINT) == c#11) + : : : : : : : : :- Project [a#3, c#11] + : : : : : : : : : +- Join Inner, (CAST(a#9 AS BIGINT) == c#8) + : : : : : : : : : :- Project [a#3, c#8] + : : : : : : : : : : +- Join Inner, (CAST(a#6 AS BIGINT) == c#5) + : : : : : : : : : : :- Project [a#3, c#5] + : : : : : : : : : : : +- Shift +3 + : : : : : : : : : : : +- Scan spark_catalog.default.t [a#0, b#1, c#2] (shared #1) + : : : : : : : : : : +- Project [a#6, c#8] + : : : : : : : : : : +- Shift +6 + : : : : : : : : : : +- Shared #1, printed above + : : : : : : : : : +- Project [a#9, c#11] + : : : : : : : : : +- Shift +9 + : : : : : : : : : +- Shared #1, printed above + : : : : : : : : +- Project [a#12, c#14] + : : : : : : : : +- Shift +12 + : : : : : : : : +- Shared #1, printed above + : : : : : : : +- Project [a#15, c#17] + : : : : : : : +- Shift +15 + : : : : : : : +- Shared #1, printed above + : : : : : : +- Project [a#18, c#20] + : : : : : : +- Shift +18 + : : : : : : +- Shared #1, printed above + : : : : : +- Project [a#21, c#23] + : : : : : +- Shift +21 + : : : : : +- Shared #1, printed above + : : : : +- Project [a#24, c#26] + : : : : +- Shift +24 + : : : : +- Shared #1, printed above + : : : +- Project [a#27, c#29] + : : : +- Shift +27 + : : : +- Shared #1, printed above + : : +- Project [a#30, c#32] + : : +- Shift +30 + : : +- Shared #1, printed above + : +- Project [a#33, b#34] + : +- Shift +33 + : +- Shared #1, printed above + +- Project [c#2] + +- Filter (b#1 == 'x') + +- Shared #1, printed above diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__positional_roots_keep_their_columns.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__positional_roots_keep_their_columns.snap index cd26fc387..b17196efe 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__positional_roots_keep_their_columns.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__positional_roots_keep_their_columns.snap @@ -3,11 +3,11 @@ source: crates/krishiv-optimizer/tests/it/join_order.rs expression: optimized.root().to_string() --- Project [a#0, b#1, c#2, a#3, d#4, a#5, b#6, c#7] -+- Join Inner, (c#7 == CAST(a#3 AS BIGINT)) - :- Shift +5 ++- Join Inner, (a#0 == a#3) + :- Filter (b#1 == 'y') : +- Scan spark_catalog.default.t [a#0, b#1, c#2] (shared #1) - +- Join Inner, (a#0 == a#3) - :- Filter (b#1 == 'y') + +- Join Inner, (c#7 == CAST(a#3 AS BIGINT)) + :- Shift +5 : +- Shared #1, printed above +- Filter (d#4 == 'z') +- Scan spark_catalog.default.u [a#3, d#4] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q10.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q10.snap index 8b8b6f795..1f3afb8b2 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q10.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q10.snap @@ -5,17 +5,17 @@ expression: optimized.root().to_string() Limit 20 +- Sort [revenue#37 DESC NULLS LAST], true +- Aggregate [c_custkey#0, c_name#1, c_acctbal#5, c_phone#4, n_name#34, c_address#2, c_comment#7], [c_custkey#0, c_name#1, sum((l_extendedprice#22 * (1 - l_discount#23))) AS revenue#37, c_acctbal#5, n_name#34, c_address#2, c_phone#4, c_comment#7] - +- Project [n_name#34, c_custkey#0, c_name#1, c_address#2, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23] - +- Join Inner, (c_nationkey#3 == n_nationkey#33) - :- Scan spark_catalog.default.nation [n_nationkey#33, n_name#34] - +- Project [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7, l_extendedprice#22, l_discount#23] - +- Join Inner, (c_custkey#0 == o_custkey#9) - :- Scan spark_catalog.default.customer [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7] - +- Project [l_extendedprice#22, l_discount#23, o_custkey#9] - +- Join Inner, (l_orderkey#17 == o_orderkey#8) - :- Project [l_orderkey#17, l_extendedprice#22, l_discount#23] - : +- Filter (l_returnflag#25 == 'R') - : +- Scan spark_catalog.default.lineitem [l_orderkey#17, l_extendedprice#22, l_discount#23, l_returnflag#25] + +- Project [l_extendedprice#22, l_discount#23, n_name#34, c_custkey#0, c_name#1, c_address#2, c_phone#4, c_acctbal#5, c_comment#7] + +- Join Inner, (l_orderkey#17 == o_orderkey#8) + :- Project [l_orderkey#17, l_extendedprice#22, l_discount#23] + : +- Filter (l_returnflag#25 == 'R') + : +- Scan spark_catalog.default.lineitem [l_orderkey#17, l_extendedprice#22, l_discount#23, l_returnflag#25] + +- Project [n_name#34, c_custkey#0, c_name#1, c_address#2, c_phone#4, c_acctbal#5, c_comment#7, o_orderkey#8] + +- Join Inner, (c_nationkey#3 == n_nationkey#33) + :- Scan spark_catalog.default.nation [n_nationkey#33, n_name#34] + +- Project [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7, o_orderkey#8] + +- Join Inner, (c_custkey#0 == o_custkey#9) + :- Scan spark_catalog.default.customer [c_custkey#0, c_name#1, c_address#2, c_nationkey#3, c_phone#4, c_acctbal#5, c_comment#7] +- Project [o_orderkey#8, o_custkey#9] +- Filter ((o_orderdate#12 >= DATE '1993-10-01') and (o_orderdate#12 < DATE '1994-01-01')) +- Scan spark_catalog.default.orders [o_orderkey#8, o_custkey#9, o_orderdate#12] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q17.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q17.snap index ad5e14a57..e175ed4d7 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q17.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q17.snap @@ -4,19 +4,19 @@ expression: optimized.root().to_string() --- Aggregate [], [(sum(l_extendedprice#5) / 7.0) AS avg_yearly#42] +- Project [l_extendedprice#5] - +- Join Inner, ((CAST(l_quantity#4 AS DECIMAL(21,7)) < (0.2 * avg(l_quantity))#41) and (p_partkey#16 == l_partkey#1)) - :- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5] (shared #1) - +- Project [(0.2 * avg(l_quantity))#41, p_partkey#16] - +- Join Inner, (p_partkey#16 == l_partkey#26) - :- Aggregate [l_partkey#26], [(0.2 * avg(l_quantity#29)) AS (0.2 * avg(l_quantity))#41, l_partkey#26] - : +- Join LeftSemi, (l_partkey#26 == p_partkey#59) - : :- Project [l_partkey#26, l_quantity#29] - : : +- Shift +25 - : : +- Shared #1, printed above - : +- Shift +43 - : +- Project [p_partkey#16] - : +- Filter ((p_brand#19 == 'Brand#23') and (p_container#22 == 'MED BOX')) - : +- Scan spark_catalog.default.part [p_partkey#16, p_brand#19, p_container#22] (shared #2) + +- Join Inner, ((p_partkey#16 == l_partkey#26) and (CAST(l_quantity#4 AS DECIMAL(21,7)) < (0.2 * avg(l_quantity))#41)) + :- Aggregate [l_partkey#26], [(0.2 * avg(l_quantity#29)) AS (0.2 * avg(l_quantity))#41, l_partkey#26] + : +- Join LeftSemi, (l_partkey#26 == p_partkey#59) + : :- Project [l_partkey#26, l_quantity#29] + : : +- Shift +25 + : : +- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5] (shared #1) + : +- Shift +43 + : +- Project [p_partkey#16] + : +- Filter ((p_brand#19 == 'Brand#23') and (p_container#22 == 'MED BOX')) + : +- Scan spark_catalog.default.part [p_partkey#16, p_brand#19, p_container#22] (shared #2) + +- Project [l_quantity#4, l_extendedprice#5, p_partkey#16] + +- Join Inner, (p_partkey#16 == l_partkey#1) + :- Shared #1, printed above +- Project [p_partkey#16] +- Filter ((p_brand#19 == 'Brand#23') and (p_container#22 == 'MED BOX')) +- Shared #2, printed above diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q21.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q21.snap index d581af153..fa8f7c6f1 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q21.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q21.snap @@ -8,22 +8,22 @@ Limit 100 +- Project [s_name#1] +- Join LeftAnti, ((l_orderkey#52 == l_orderkey#7) and (NOT (l_suppkey#54 == l_suppkey#9))) :- Join LeftSemi, ((l_orderkey#36 == l_orderkey#7) and (NOT (l_suppkey#38 == l_suppkey#9))) - : :- Project [s_name#1, l_orderkey#7, l_suppkey#9] - : : +- Join Inner, (s_nationkey#3 == n_nationkey#32) - : : :- Project [n_nationkey#32] - : : : +- Filter (n_name#33 == 'SAUDI ARABIA') - : : : +- Scan spark_catalog.default.nation [n_nationkey#32, n_name#33] - : : +- Project [s_name#1, s_nationkey#3, l_orderkey#7, l_suppkey#9] + : :- Project [l_orderkey#7, l_suppkey#9, s_name#1] + : : +- Join Inner, (o_orderkey#23 == l_orderkey#7) + : : :- Project [o_orderkey#23] + : : : +- Filter (o_orderstatus#25 == 'F') + : : : +- Scan spark_catalog.default.orders [o_orderkey#23, o_orderstatus#25] + : : +- Project [l_orderkey#7, l_suppkey#9, s_name#1] : : +- Join Inner, (s_suppkey#0 == l_suppkey#9) - : : :- Scan spark_catalog.default.supplier [s_suppkey#0, s_name#1, s_nationkey#3] - : : +- Project [l_orderkey#7, l_suppkey#9] - : : +- Join Inner, (o_orderkey#23 == l_orderkey#7) - : : :- Project [l_orderkey#7, l_suppkey#9] - : : : +- Filter (l_receiptdate#19 > l_commitdate#18) - : : : +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_commitdate#18, l_receiptdate#19] (shared #1) - : : +- Project [o_orderkey#23] - : : +- Filter (o_orderstatus#25 == 'F') - : : +- Scan spark_catalog.default.orders [o_orderkey#23, o_orderstatus#25] + : : :- Project [l_orderkey#7, l_suppkey#9] + : : : +- Filter (l_receiptdate#19 > l_commitdate#18) + : : : +- Scan spark_catalog.default.lineitem [l_orderkey#7, l_suppkey#9, l_commitdate#18, l_receiptdate#19] (shared #1) + : : +- Project [s_suppkey#0, s_name#1] + : : +- Join Inner, (s_nationkey#3 == n_nationkey#32) + : : :- Scan spark_catalog.default.supplier [s_suppkey#0, s_name#1, s_nationkey#3] + : : +- Project [n_nationkey#32] + : : +- Filter (n_name#33 == 'SAUDI ARABIA') + : : +- Scan spark_catalog.default.nation [n_nationkey#32, n_name#33] : +- Project [l_orderkey#36, l_suppkey#38] : +- Shift +29 : +- Shared #1, printed above From 4c05a6ea5de743b5be224cb87cf40963b144d237 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Fri, 23 Oct 2026 19:20:24 +0200 Subject: [PATCH 13/16] optimizer: join on the equalities every disjunct of a condition holds Join ordering takes a condition's conjuncts as its join edges. TPC-DS q13 joins store_sales, customer_demographics and household_demographics through one `or` whose three disjuncts each hold `cd_demo_sk = ss_cdemo_sk` and `ss_hdemo_sk = hd_demo_sk`: a single condition over three inputs, so two of them could only meet in a cross product, which joins the parts largest first. With the demographics filter estimated at its default selectivity, the filtered customer_demographics became the larger part and was crossed with the fact side (hundreds of millions of pairs instead of the household side's thousands). Before ordering, the column equalities every disjunct holds are now conjuncts of their own and the `or` keeps the rest: `(a = b and x) or (a = b and y)` is `a = b and (x or y)` (distributivity, valid in three-valued logic; a column equality cannot raise and is deterministic). q13 then joins every dimension to the fact table by key, with no cross product; q19 and q85 conditions are rewritten the same way (q85 also joins its two demographics instances on the shared marital and education status). New test: q13's shape over TPC-H statistics plans without a cross product (it does not without this change), and the rewrite keeps the rows of a three-input `or`, of a disjunct holding only the shared equalities, and of equalities only some disjuncts hold. Design: P2-T166 (PR 1a, review of the benchmark A/B). Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .../planner/snapshots/planner__tpch__q19.snap | 2 +- crates/krishiv-optimizer/src/expr.rs | 43 +++++++++++++++++++ crates/krishiv-optimizer/src/join_order.rs | 4 +- .../krishiv-optimizer/tests/it/join_order.rs | 31 +++++++++++++ ...very_disjunct_holds_join_their_inputs.snap | 13 ++++++ .../it/snapshots/it__join_order__q19.snap | 2 +- .../tests/it/snapshots/it__tpch__q19.snap | 2 +- 7 files changed, 93 insertions(+), 4 deletions(-) create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__join_order__equalities_every_disjunct_holds_join_their_inputs.snap diff --git a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q19.snap b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q19.snap index 4f36987a7..3e8fae454 100644 --- a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q19.snap +++ b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q19.snap @@ -7,7 +7,7 @@ Project [sum((l_extendedprice#5 * (1 - l_discount#6)))#4611686018427387919 AS re +- Exchange SinglePartition +- Project [(l_extendedprice#5 * (1 - l_discount#6))] +- Project [l_extendedprice#5, l_discount#6] - +- HashJoin Inner, [l_partkey#1] = [p_partkey#16], residual: ((((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#12')) and ((p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')) and (l_quantity#4 >= 1.00))) and (((l_quantity#4 <= 11.00) and ((p_size#21 >= 1) and (p_size#21 <= 5))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#23')) and ((p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')) and (l_quantity#4 >= 10.00))) and (((l_quantity#4 <= 20.00) and ((p_size#21 >= 1) and (p_size#21 <= 10))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#34')) and ((p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')) and (l_quantity#4 >= 20.00))) and (((l_quantity#4 <= 30.00) and ((p_size#21 >= 1) and (p_size#21 <= 15))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))), RuntimeFilter #1: l_partkey#1 -> spark_catalog.default.lineitem.l_partkey + +- HashJoin Inner, [l_partkey#1] = [p_partkey#16], residual: ((((((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (l_quantity#4 >= 1.00)) and (l_quantity#4 <= 11.00)) and (p_size#21 >= 1)) and (p_size#21 <= 5)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (l_quantity#4 >= 10.00)) and (l_quantity#4 <= 20.00)) and (p_size#21 >= 1)) and (p_size#21 <= 10)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (l_quantity#4 >= 20.00)) and (l_quantity#4 <= 30.00)) and (p_size#21 >= 1)) and (p_size#21 <= 15)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))), RuntimeFilter #1: l_partkey#1 -> spark_catalog.default.lineitem.l_partkey :- Filter ((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) : +- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5, l_discount#6, l_shipinstruct#13, l_shipmode#14], PushedFilters: [((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))] +- BroadcastExchange diff --git a/crates/krishiv-optimizer/src/expr.rs b/crates/krishiv-optimizer/src/expr.rs index 140cbbaf4..a27aaa749 100644 --- a/crates/krishiv-optimizer/src/expr.rs +++ b/crates/krishiv-optimizer/src/expr.rs @@ -235,6 +235,49 @@ pub(crate) fn implied(pred: &Arc, side: &Attrs) -> Option> { or_all(out) } +/// `pred` with the column equalities every one of its disjuncts holds +/// taken out as conjuncts of their own: `(a = b and x) or (a = b and y)` is +/// `a = b` and `x or y` (distributivity, which holds in three-valued logic +/// too). A plain column equality cannot raise and is deterministic, so +/// evaluating it once instead of in each disjunct changes nothing. Join +/// ordering needs it: an `or` over three inputs whose disjuncts all join +/// the same keys (TPC-DS q13) is otherwise one condition over all three, +/// and the inputs meet only in a cross product. +pub(crate) fn common_equalities(pred: &Arc) -> Vec> { + let parts = disjuncts(pred); + if parts.len() < 2 || has_subquery(pred) { + return vec![Arc::clone(pred)]; + } + let each: Vec>> = parts.iter().map(conjuncts).collect(); + let is_column = |e: &Expr| matches!(e.node, ExprNode::Resolved(ResolvedExpr::Column(_))); + let common: Vec> = each[0] + .iter() + .filter(|c| matches!(builtin(c), Some((Builtin::Eq, [l, r], _)) if is_column(l) && is_column(r))) + .filter(|c| each[1..].iter().all(|d| d.iter().any(|x| same(x, c)))) + .cloned() + .collect(); + if common.is_empty() { + return vec![Arc::clone(pred)]; + } + let mut rest = Vec::with_capacity(each.len()); + for d in &each { + let left: Vec> = d + .iter() + .filter(|x| !common.iter().any(|c| same(x, c))) + .cloned() + .collect(); + match and_all(left) { + Some(e) => rest.push(e), + // A disjunct of only common equalities: the `or` always holds + // once they do. + None => return common, + } + } + let mut out = common; + out.extend(or_all(rest)); + out +} + /// `a and b and ...`, left-deep; `None` for no predicates. pub(crate) fn and_all(preds: impl IntoIterator>) -> Option> { preds diff --git a/crates/krishiv-optimizer/src/join_order.rs b/crates/krishiv-optimizer/src/join_order.rs index 4ad247a85..98cc9a545 100644 --- a/crates/krishiv-optimizer/src/join_order.rs +++ b/crates/krishiv-optimizer/src/join_order.rs @@ -23,7 +23,8 @@ use krishiv_plan::{ }; use crate::expr::{ - Attrs, and_all, column, conjuncts, has_subquery, implied, is_deterministic, references, + Attrs, and_all, column, common_equalities, conjuncts, has_subquery, implied, is_deterministic, + references, }; use crate::stats::{Est, Estimator, Statistics, equi_columns, selectivity}; @@ -228,6 +229,7 @@ impl Reorder<'_> { if leaves.len() < 2 { return self.children(root); } + let preds = preds.iter().flat_map(common_equalities).collect(); let ests: Vec> = leaves.iter().map(|l| self.est.plan(l)).collect(); let graph = Graph::new(&leaves, &ests, preds); if graph.parts(false) < graph.parts(true) { diff --git a/crates/krishiv-optimizer/tests/it/join_order.rs b/crates/krishiv-optimizer/tests/it/join_order.rs index 39bcf4fae..8263adddf 100644 --- a/crates/krishiv-optimizer/tests/it/join_order.rs +++ b/crates/krishiv-optimizer/tests/it/join_order.rs @@ -555,3 +555,34 @@ fn filtered_dimensions_count_the_keys_they_keep() { assert!(line("d_date_sk") < line("cd_demo_sk"), "{text}"); insta::assert_snapshot!(text); } + +/// TPC-DS q13's shape: one `or` over three inputs whose disjuncts all hold +/// the same key equalities. The equalities are join conditions of their +/// own, so the inputs meet in equi-joins; as one condition over all three +/// inputs they could only meet in a cross product (q13 crossed a filtered +/// customer_demographics with the fact side). +#[test] +fn equalities_every_disjunct_holds_join_their_inputs() { + let sql = "SELECT count(*) FROM lineitem, orders, customer \ + WHERE (l_orderkey = o_orderkey AND o_custkey = c_custkey \ + AND l_quantity > 10 AND c_acctbal > 0) \ + OR (l_orderkey = o_orderkey AND o_custkey = c_custkey \ + AND l_quantity < 5 AND o_totalprice > 100)"; + let plan = support::resolve(sql, &support::ctx()); + let optimized = optimize_with(&plan, &fixture()).unwrap(); + no_cross_product(optimized.root(), sql); + insta::assert_snapshot!(optimized.root().to_string()); + // The same rows, and a disjunct of only shared equalities leaves just + // the equalities. + same_rows( + "SELECT t.a, u.d, x.b FROM t, u, t x \ + WHERE (t.a = u.a AND u.a = x.a AND t.b = 'p') OR (t.a = u.a AND u.a = x.a AND x.c > 1)", + ); + let only = same_rows( + "SELECT t.a, u.d, x.b FROM t, u, t x \ + WHERE (t.a = u.a AND u.a = x.a) OR (t.a = u.a AND u.a = x.a AND x.c > 1)", + ); + assert!(!only.root().to_string().contains(" or "), "{}", only.root()); + // Equalities only some disjuncts hold stay inside the `or`. + same_rows("SELECT t.a, u.d FROM t, u WHERE (t.a = u.a AND t.b = 'p') OR (t.c = u.a)"); +} diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__equalities_every_disjunct_holds_join_their_inputs.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__equalities_every_disjunct_holds_join_their_inputs.snap new file mode 100644 index 000000000..4ebeb28ce --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__equalities_every_disjunct_holds_join_their_inputs.snap @@ -0,0 +1,13 @@ +--- +source: crates/krishiv-optimizer/tests/it/join_order.rs +expression: optimized.root().to_string() +--- +Aggregate [], [count(1) AS count(1)#33] ++- Project [l_orderkey#0] + +- Join Inner, ((l_orderkey#0 == o_orderkey#16) and (((l_quantity#4 > 10.00) and (c_acctbal#30 > 0.00)) or ((l_quantity#4 < 5.00) and (o_totalprice#19 > 100.00)))) + :- Filter ((l_quantity#4 > 10.00) or (l_quantity#4 < 5.00)) + : +- Scan spark_catalog.default.lineitem [l_orderkey#0, l_quantity#4] + +- Project [o_orderkey#16, o_totalprice#19, c_acctbal#30] + +- Join Inner, (o_custkey#17 == c_custkey#25) + :- Scan spark_catalog.default.orders [o_orderkey#16, o_custkey#17, o_totalprice#19] + +- Scan spark_catalog.default.customer [c_custkey#25, c_acctbal#30] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__q19.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__q19.snap index 560ab2df4..716e9ff25 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__q19.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__q19.snap @@ -4,7 +4,7 @@ expression: optimized.root().to_string() --- Aggregate [], [sum((l_extendedprice#5 * (1 - l_discount#6))) AS revenue#25] +- Project [l_extendedprice#5, l_discount#6] - +- Join Inner, ((((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#12')) and ((p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')) and (l_quantity#4 >= 1.00))) and (((l_quantity#4 <= 11.00) and ((p_size#21 >= 1) and (p_size#21 <= 5))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#23')) and ((p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')) and (l_quantity#4 >= 10.00))) and (((l_quantity#4 <= 20.00) and ((p_size#21 >= 1) and (p_size#21 <= 10))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#34')) and ((p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')) and (l_quantity#4 >= 20.00))) and (((l_quantity#4 <= 30.00) and ((p_size#21 >= 1) and (p_size#21 <= 15))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) + +- Join Inner, ((p_partkey#16 == l_partkey#1) and ((((((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (l_quantity#4 >= 1.00)) and (l_quantity#4 <= 11.00)) and (p_size#21 >= 1)) and (p_size#21 <= 5)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (l_quantity#4 >= 10.00)) and (l_quantity#4 <= 20.00)) and (p_size#21 >= 1)) and (p_size#21 <= 10)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (l_quantity#4 >= 20.00)) and (l_quantity#4 <= 30.00)) and (p_size#21 >= 1)) and (p_size#21 <= 15)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) :- Filter ((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) : +- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5, l_discount#6, l_shipinstruct#13, l_shipmode#14] +- Filter ((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 5)) or ((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (p_size#21 >= 1)) and (p_size#21 <= 10))) or ((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 15))) diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap index b22844ab2..e0e9775a2 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap @@ -4,7 +4,7 @@ expression: optimized.root().to_string() --- Aggregate [], [sum((l_extendedprice#5 * (1 - l_discount#6))) AS revenue#25] +- Project [l_extendedprice#5, l_discount#6] - +- Join Inner, ((((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#12')) and ((p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')) and (l_quantity#4 >= 1.00))) and (((l_quantity#4 <= 11.00) and ((p_size#21 >= 1) and (p_size#21 <= 5))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#23')) and ((p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')) and (l_quantity#4 >= 10.00))) and (((l_quantity#4 <= 20.00) and ((p_size#21 >= 1) and (p_size#21 <= 10))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#34')) and ((p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')) and (l_quantity#4 >= 20.00))) and (((l_quantity#4 <= 30.00) and ((p_size#21 >= 1) and (p_size#21 <= 15))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) + +- Join Inner, ((p_partkey#16 == l_partkey#1) and ((((((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (l_quantity#4 >= 1.00)) and (l_quantity#4 <= 11.00)) and (p_size#21 >= 1)) and (p_size#21 <= 5)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (l_quantity#4 >= 10.00)) and (l_quantity#4 <= 20.00)) and (p_size#21 >= 1)) and (p_size#21 <= 10)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (l_quantity#4 >= 20.00)) and (l_quantity#4 <= 30.00)) and (p_size#21 >= 1)) and (p_size#21 <= 15)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) :- Filter ((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 5)) or ((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (p_size#21 >= 1)) and (p_size#21 <= 10))) or ((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 15))) : +- Scan spark_catalog.default.part [p_partkey#16, p_brand#19, p_size#21, p_container#22] +- Filter ((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) From 10e470e83824b829f41a2a9f5a5368be12f76c19 Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:07:00 +0200 Subject: [PATCH 14/16] optimizer: take shared equalities out of an `or` only to connect inputs The previous change took the column equalities every disjunct holds out of each `or` condition. Where other conditions already join the inputs, that only adds join keys: TPC-DS q85 joined its two demographics instances on two string columns as well, and ran slower. The equalities are now taken out only when some inputs are joined by nothing but a condition over three or more of them, which is the case they exist for (TPC-DS q13). q19 and q85 plan as before that change. New test: q85's shape (an `or` over three inputs other conditions join) keeps its shared equality inside the `or`. Design: P2-T166 (PR 1a, review of the benchmark A/B). Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- .../planner/snapshots/planner__tpch__q19.snap | 2 +- crates/krishiv-optimizer/src/join_order.rs | 29 +++++++++++++++++-- .../krishiv-optimizer/tests/it/join_order.rs | 17 +++++++++++ .../it/snapshots/it__join_order__q19.snap | 2 +- ..._in_the_or_when_the_inputs_are_joined.snap | 13 +++++++++ .../tests/it/snapshots/it__tpch__q19.snap | 2 +- 6 files changed, 60 insertions(+), 5 deletions(-) create mode 100644 crates/krishiv-optimizer/tests/it/snapshots/it__join_order__shared_equalities_stay_in_the_or_when_the_inputs_are_joined.snap diff --git a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q19.snap b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q19.snap index 3e8fae454..4f36987a7 100644 --- a/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q19.snap +++ b/crates/krishiv-exec/tests/planner/snapshots/planner__tpch__q19.snap @@ -7,7 +7,7 @@ Project [sum((l_extendedprice#5 * (1 - l_discount#6)))#4611686018427387919 AS re +- Exchange SinglePartition +- Project [(l_extendedprice#5 * (1 - l_discount#6))] +- Project [l_extendedprice#5, l_discount#6] - +- HashJoin Inner, [l_partkey#1] = [p_partkey#16], residual: ((((((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (l_quantity#4 >= 1.00)) and (l_quantity#4 <= 11.00)) and (p_size#21 >= 1)) and (p_size#21 <= 5)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (l_quantity#4 >= 10.00)) and (l_quantity#4 <= 20.00)) and (p_size#21 >= 1)) and (p_size#21 <= 10)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (l_quantity#4 >= 20.00)) and (l_quantity#4 <= 30.00)) and (p_size#21 >= 1)) and (p_size#21 <= 15)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))), RuntimeFilter #1: l_partkey#1 -> spark_catalog.default.lineitem.l_partkey + +- HashJoin Inner, [l_partkey#1] = [p_partkey#16], residual: ((((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#12')) and ((p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')) and (l_quantity#4 >= 1.00))) and (((l_quantity#4 <= 11.00) and ((p_size#21 >= 1) and (p_size#21 <= 5))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#23')) and ((p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')) and (l_quantity#4 >= 10.00))) and (((l_quantity#4 <= 20.00) and ((p_size#21 >= 1) and (p_size#21 <= 10))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#34')) and ((p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')) and (l_quantity#4 >= 20.00))) and (((l_quantity#4 <= 30.00) and ((p_size#21 >= 1) and (p_size#21 <= 15))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))), RuntimeFilter #1: l_partkey#1 -> spark_catalog.default.lineitem.l_partkey :- Filter ((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) : +- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5, l_discount#6, l_shipinstruct#13, l_shipmode#14], PushedFilters: [((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))] +- BroadcastExchange diff --git a/crates/krishiv-optimizer/src/join_order.rs b/crates/krishiv-optimizer/src/join_order.rs index 98cc9a545..104f4ba58 100644 --- a/crates/krishiv-optimizer/src/join_order.rs +++ b/crates/krishiv-optimizer/src/join_order.rs @@ -229,9 +229,17 @@ impl Reorder<'_> { if leaves.len() < 2 { return self.children(root); } - let preds = preds.iter().flat_map(common_equalities).collect(); let ests: Vec> = leaves.iter().map(|l| self.est.plan(l)).collect(); - let graph = Graph::new(&leaves, &ests, preds); + let mut graph = Graph::new(&leaves, &ests, preds.clone()); + if graph.binary_parts() > graph.parts(true) { + // Some inputs are joined only by a condition over three or more + // of them, so two of them would meet in a cross product. Take + // the column equalities every disjunct of such a condition holds + // as conditions of their own (TPC-DS q13). Only then: elsewhere + // they would add join keys the inputs do not need. + let preds = preds.iter().flat_map(common_equalities).collect(); + graph = Graph::new(&leaves, &ests, preds); + } if graph.parts(false) < graph.parts(true) { // Inputs related only through a conjunct with a subquery: as // written they meet in a join on it; re-planned they would meet @@ -475,6 +483,23 @@ impl Graph { (0..self.n()).filter(|&x| find(&mut root, x) == x).count() } + /// The connected parts through join conditions over exactly two inputs. + fn binary_parts(&self) -> usize { + let mut root: Vec = (0..self.n()).collect(); + fn find(root: &mut [usize], mut x: usize) -> usize { + while root[x] != x { + root[x] = root[root[x]]; + x = root[x]; + } + x + } + for p in self.preds.iter().filter(|p| p.edge && p.on.len() == 2) { + let (a, b) = (find(&mut root, p.on[0]), find(&mut root, p.on[1])); + root[a] = b; + } + (0..self.n()).filter(|&x| find(&mut root, x) == x).count() + } + fn n(&self) -> usize { self.leaves.len() } diff --git a/crates/krishiv-optimizer/tests/it/join_order.rs b/crates/krishiv-optimizer/tests/it/join_order.rs index 8263adddf..0071e54af 100644 --- a/crates/krishiv-optimizer/tests/it/join_order.rs +++ b/crates/krishiv-optimizer/tests/it/join_order.rs @@ -586,3 +586,20 @@ fn equalities_every_disjunct_holds_join_their_inputs() { // Equalities only some disjuncts hold stay inside the `or`. same_rows("SELECT t.a, u.d FROM t, u WHERE (t.a = u.a AND t.b = 'p') OR (t.c = u.a)"); } + +/// TPC-DS q85's shape: an `or` over three inputs that other conditions +/// already join. Its shared equality is not needed to connect them, so it +/// stays inside the `or` instead of becoming an extra join key (q85 would +/// otherwise join its two demographics instances on two string columns). +#[test] +fn shared_equalities_stay_in_the_or_when_the_inputs_are_joined() { + let sql = "SELECT count(*) FROM lineitem, orders, customer \ + WHERE l_orderkey = o_orderkey AND o_custkey = c_custkey \ + AND ((o_orderpriority = c_mktsegment AND l_quantity > 10) \ + OR (o_orderpriority = c_mktsegment AND l_quantity < 5))"; + let plan = support::resolve(sql, &support::ctx()); + let optimized = optimize_with(&plan, &fixture()).unwrap(); + let text = optimized.root().to_string(); + assert_eq!(text.matches("(o_orderpriority#").count(), 2, "{text}"); + insta::assert_snapshot!(text); +} diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__q19.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__q19.snap index 716e9ff25..560ab2df4 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__q19.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__q19.snap @@ -4,7 +4,7 @@ expression: optimized.root().to_string() --- Aggregate [], [sum((l_extendedprice#5 * (1 - l_discount#6))) AS revenue#25] +- Project [l_extendedprice#5, l_discount#6] - +- Join Inner, ((p_partkey#16 == l_partkey#1) and ((((((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (l_quantity#4 >= 1.00)) and (l_quantity#4 <= 11.00)) and (p_size#21 >= 1)) and (p_size#21 <= 5)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (l_quantity#4 >= 10.00)) and (l_quantity#4 <= 20.00)) and (p_size#21 >= 1)) and (p_size#21 <= 10)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (l_quantity#4 >= 20.00)) and (l_quantity#4 <= 30.00)) and (p_size#21 >= 1)) and (p_size#21 <= 15)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) + +- Join Inner, ((((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#12')) and ((p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')) and (l_quantity#4 >= 1.00))) and (((l_quantity#4 <= 11.00) and ((p_size#21 >= 1) and (p_size#21 <= 5))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#23')) and ((p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')) and (l_quantity#4 >= 10.00))) and (((l_quantity#4 <= 20.00) and ((p_size#21 >= 1) and (p_size#21 <= 10))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#34')) and ((p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')) and (l_quantity#4 >= 20.00))) and (((l_quantity#4 <= 30.00) and ((p_size#21 >= 1) and (p_size#21 <= 15))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) :- Filter ((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) : +- Scan spark_catalog.default.lineitem [l_partkey#1, l_quantity#4, l_extendedprice#5, l_discount#6, l_shipinstruct#13, l_shipmode#14] +- Filter ((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 5)) or ((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (p_size#21 >= 1)) and (p_size#21 <= 10))) or ((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 15))) diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__shared_equalities_stay_in_the_or_when_the_inputs_are_joined.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__shared_equalities_stay_in_the_or_when_the_inputs_are_joined.snap new file mode 100644 index 000000000..98a4f5deb --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__join_order__shared_equalities_stay_in_the_or_when_the_inputs_are_joined.snap @@ -0,0 +1,13 @@ +--- +source: crates/krishiv-optimizer/tests/it/join_order.rs +expression: text +--- +Aggregate [], [count(1) AS count(1)#33] ++- Project [l_orderkey#0] + +- Join Inner, ((((o_orderpriority#21 == c_mktsegment#31) and (l_quantity#4 > 10.00)) or ((o_orderpriority#21 == c_mktsegment#31) and (l_quantity#4 < 5.00))) and (l_orderkey#0 == o_orderkey#16)) + :- Filter ((l_quantity#4 > 10.00) or (l_quantity#4 < 5.00)) + : +- Scan spark_catalog.default.lineitem [l_orderkey#0, l_quantity#4] + +- Project [o_orderkey#16, o_orderpriority#21, c_mktsegment#31] + +- Join Inner, (o_custkey#17 == c_custkey#25) + :- Scan spark_catalog.default.orders [o_orderkey#16, o_custkey#17, o_orderpriority#21] + +- Scan spark_catalog.default.customer [c_custkey#25, c_mktsegment#31] diff --git a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap index e0e9775a2..b22844ab2 100644 --- a/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap +++ b/crates/krishiv-optimizer/tests/it/snapshots/it__tpch__q19.snap @@ -4,7 +4,7 @@ expression: optimized.root().to_string() --- Aggregate [], [sum((l_extendedprice#5 * (1 - l_discount#6))) AS revenue#25] +- Project [l_extendedprice#5, l_discount#6] - +- Join Inner, ((p_partkey#16 == l_partkey#1) and ((((((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (l_quantity#4 >= 1.00)) and (l_quantity#4 <= 11.00)) and (p_size#21 >= 1)) and (p_size#21 <= 5)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (l_quantity#4 >= 10.00)) and (l_quantity#4 <= 20.00)) and (p_size#21 >= 1)) and (p_size#21 <= 10)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (l_quantity#4 >= 20.00)) and (l_quantity#4 <= 30.00)) and (p_size#21 >= 1)) and (p_size#21 <= 15)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) + +- Join Inner, ((((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#12')) and ((p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG')) and (l_quantity#4 >= 1.00))) and (((l_quantity#4 <= 11.00) and ((p_size#21 >= 1) and (p_size#21 <= 5))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON')))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#23')) and ((p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK')) and (l_quantity#4 >= 10.00))) and (((l_quantity#4 <= 20.00) and ((p_size#21 >= 1) and (p_size#21 <= 10))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) or ((((p_partkey#16 == l_partkey#1) and (p_brand#19 == 'Brand#34')) and ((p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG')) and (l_quantity#4 >= 20.00))) and (((l_quantity#4 <= 30.00) and ((p_size#21 >= 1) and (p_size#21 <= 15))) and ((l_shipmode#14 IN ('AIR', 'AIR REG')) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))))) :- Filter ((((((p_brand#19 == 'Brand#12') and (p_container#22 IN ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 5)) or ((((p_brand#19 == 'Brand#23') and (p_container#22 IN ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK'))) and (p_size#21 >= 1)) and (p_size#21 <= 10))) or ((((p_brand#19 == 'Brand#34') and (p_container#22 IN ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG'))) and (p_size#21 >= 1)) and (p_size#21 <= 15))) : +- Scan spark_catalog.default.part [p_partkey#16, p_brand#19, p_size#21, p_container#22] +- Filter ((((((l_quantity#4 >= 1.00) and (l_quantity#4 <= 11.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON')) or ((((l_quantity#4 >= 10.00) and (l_quantity#4 <= 20.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) or ((((l_quantity#4 >= 20.00) and (l_quantity#4 <= 30.00)) and (l_shipmode#14 IN ('AIR', 'AIR REG'))) and (l_shipinstruct#13 == 'DELIVER IN PERSON'))) From 72eab5c57144726a41bad7b3b98e986e6ca6c92c Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:18:38 +0200 Subject: [PATCH 15/16] review: null-safe NULL selectivity, WITH references as default estimates, exact orientation tests - `c <=> NULL` holds exactly where `c` is NULL: its selectivity is now the column's null fraction, as for `c IS NULL`, instead of `1 / ndv` or the default (the estimator sees the form the optimizer keeps). - `has_statistics` counts a `WITH` reference as a default-sized input, as the estimator reads it, so orientation takes the larger factor. - The orientation tests for joins without a preference assert each join's exact line (type, key order, filters) instead of only the absence of a preference. Design: P2-T166 (PR 1a, review). Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/src/planner/build.rs | 13 +++--- .../krishiv-exec/tests/planner/orientation.rs | 42 +++++++++++++++---- crates/krishiv-optimizer/src/stats.rs | 9 ++++ .../krishiv-optimizer/tests/it/join_order.rs | 20 +++++++++ 4 files changed, 71 insertions(+), 13 deletions(-) diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index 09b04a934..5cabc1717 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -632,16 +632,19 @@ impl<'a> Planner<'a> { } /// Whether every table `plan` reads reports statistics, so its estimate - /// does not rest on default row counts. + /// does not rest on default row counts. A `WITH` reference counts as + /// default-sized, as the estimator treats it. pub(crate) fn has_statistics(&self, plan: &Arc) -> bool { use krishiv_optimizer::Statistics as _; let stats = super::ProviderStatistics(Arc::clone(&self.cx.tables)); let mut stack = vec![plan]; while let Some(p) = stack.pop() { - if let PlanNode::Scan(scan) = &p.node - && stats.table(scan).is_none() - { - return false; + match &p.node { + PlanNode::Scan(scan) if stats.table(scan).is_none() => return false, + // The estimator reads a reference to a `WITH` definition + // as a default-sized relation, not through its definition. + PlanNode::CteRef(_) => return false, + _ => {} } p.for_each_child(&mut |c| stack.push(c)); } diff --git a/crates/krishiv-exec/tests/planner/orientation.rs b/crates/krishiv-exec/tests/planner/orientation.rs index 5e0499770..f40789f21 100644 --- a/crates/krishiv-exec/tests/planner/orientation.rs +++ b/crates/krishiv-exec/tests/planner/orientation.rs @@ -92,19 +92,30 @@ fn a_much_larger_right_input_is_shown_as_the_side_not_to_build() { #[test] fn null_aware_inner_small_and_hinted_joins_show_no_preference() { let env = env(); - for sql in [ + for (sql, want) in [ // `NOT IN` reads whether the whole subquery side holds a NULL. - "SELECT a FROM t WHERE a NOT IN (SELECT a FROM big)", + ( + "SELECT a FROM t WHERE a NOT IN (SELECT a FROM big)", + "HashJoin NullAwareAnti, [a#0] = [a#3]", + ), // Inner joins keep the join orderer's choice. - "SELECT t.a FROM t JOIN big ON t.a = big.a", + ( + "SELECT t.a FROM t JOIN big ON t.a = big.a", + "HashJoin Inner, [a#0] = [a#3], RuntimeFilter #1: a#0 -> spark_catalog.default.t.a", + ), // A build small enough to broadcast stays. - "SELECT a FROM big WHERE EXISTS (SELECT 1 FROM t WHERE big.a = t.a)", + ( + "SELECT a FROM big WHERE EXISTS (SELECT 1 FROM t WHERE big.a = t.a)", + "HashJoin LeftSemi, [a#0] = [a#3], \ + RuntimeFilter #1: a#0 -> spark_catalog.default.big.a (bloom)", + ), // A hint decides. - "SELECT /*+ SHUFFLE_HASH(big) */ t.a FROM t LEFT JOIN big ON t.a = big.a", + ( + "SELECT /*+ SHUFFLE_HASH(big) */ t.a FROM t LEFT JOIN big ON t.a = big.a", + "HashJoin LeftOuter, [a#0] = [a#3]", + ), ] { - let js = joins(&env, sql); - assert!(!js.is_empty(), "{sql}"); - assert!(js.iter().all(|j| !j.contains("prefers")), "{sql}: {js:?}"); + assert_eq!(joins(&env, sql), [want], "{sql}"); } } @@ -134,3 +145,18 @@ fn a_default_estimate_needs_a_larger_ratio() { "{sql}" ); } + +/// A `WITH` reference reads as a default-sized relation (the estimator +/// does not look through it), so it is not taken as resting on statistics. +#[test] +fn a_cte_reference_is_a_default_estimate() { + let env = env(); + let sql = "WITH c AS (SELECT a FROM big) \ + SELECT a FROM t WHERE EXISTS (SELECT 1 FROM c WHERE c.a = t.a)"; + let js = joins(&env, sql); + assert_eq!(js.len(), 1, "{js:?}"); + assert!( + js[0].ends_with("prefers RightSemi (probe: 10 rows / 0.0 MiB, statistics; build: 1000000 rows / 11.4 MiB, default)"), + "{js:?}" + ); +} diff --git a/crates/krishiv-optimizer/src/stats.rs b/crates/krishiv-optimizer/src/stats.rs index 06f208af8..7d93fa7b9 100644 --- a/crates/krishiv-optimizer/src/stats.rs +++ b/crates/krishiv-optimizer/src/stats.rs @@ -331,6 +331,10 @@ fn column_ndv(e: &Expr, est: &Est) -> Option { column(e, est).map(|c| c.ndv) } +fn is_null_literal(e: &Expr) -> bool { + matches!(literal(e), Some(Literal::Null(_))) +} + fn literal(e: &Expr) -> Option<&Literal> { match &e.node { ExprNode::Literal(l) => Some(l), @@ -362,6 +366,11 @@ pub(crate) fn selectivity(pred: &Expr, est: &Est) -> f64 { x + y - x * y } (Some((Builtin::Not, [a], _)), _) => 1.0 - selectivity(a, est), + // `c <=> NULL` holds exactly where `c` is NULL, as `c IS NULL`. + (Some((Builtin::EqNullSafe, [a, b], _)), _) if is_null_literal(a) || is_null_literal(b) => { + let c = if is_null_literal(a) { b } else { a }; + column(c, est).map_or(DEFAULT_SELECTIVITY, |c| c.nulls) + } (Some((Builtin::Eq | Builtin::EqNullSafe, [a, b], _)), _) => { match (column(a, est), column(b, est)) { (Some(x), Some(y)) => 1.0 / x.ndv.max(y.ndv), diff --git a/crates/krishiv-optimizer/tests/it/join_order.rs b/crates/krishiv-optimizer/tests/it/join_order.rs index 0071e54af..5918e1dc9 100644 --- a/crates/krishiv-optimizer/tests/it/join_order.rs +++ b/crates/krishiv-optimizer/tests/it/join_order.rs @@ -525,6 +525,26 @@ fn literal_equalities_without_a_distinct_count_use_the_default() { let only = without_distinct(&["l_returnflag"]); assert!(close(rows(eq, &only), lineitem / 3.0)); // An equality between two columns of one input keeps `1 / ndv`. + // `c <=> NULL` is `c IS NULL`: the column's null fraction (none here). + // Optimized first: folding makes the coerced NULL a literal. + let folded = |sql: &str, stats: &Fixture| { + let plan = optimize_with(&support::resolve(sql, &ctx), stats).unwrap(); + krishiv_optimizer::estimated_rows(plan.root(), stats) + }; + let is_null = folded( + "SELECT l_orderkey FROM lineitem WHERE l_returnflag IS NULL", + &known, + ); + for sql in [ + "SELECT l_orderkey FROM lineitem WHERE l_returnflag <=> NULL", + "SELECT l_orderkey FROM lineitem WHERE NULL <=> l_returnflag", + ] { + let (k, f) = (folded(sql, &known), folded(sql, &footers)); + assert!( + close(k, is_null) && close(f, is_null), + "{sql}: {k} {f} vs {is_null}" + ); + } let cols = "SELECT l_orderkey FROM lineitem WHERE l_commitdate = l_receiptdate"; assert!(close(rows(cols, &footers), 1.0), "{}", rows(cols, &footers)); } From 1de5b2e141189764e02894816f1a852d21f7332a Mon Sep 17 00:00:00 2001 From: techysneha14 <343708768+shehatathe@users.noreply.github.com> Date: Sat, 24 Oct 2026 16:14:40 +0200 Subject: [PATCH 16/16] test(exec): assert the exact join when a default estimate keeps the orientation Signed-off-by: techysneha14 <343708768+shehatathe@users.noreply.github.com> --- crates/krishiv-exec/tests/planner/orientation.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/krishiv-exec/tests/planner/orientation.rs b/crates/krishiv-exec/tests/planner/orientation.rs index f40789f21..04f77a2fa 100644 --- a/crates/krishiv-exec/tests/planner/orientation.rs +++ b/crates/krishiv-exec/tests/planner/orientation.rs @@ -125,8 +125,10 @@ fn a_default_estimate_needs_a_larger_ratio() { // 5M rows against a default 1M: 5x is enough with statistics on both // sides (above 2x), not with a default estimate on one (8x). let sql = "SELECT a FROM bare WHERE EXISTS (SELECT 1 FROM mid WHERE mid.a = bare.a)"; - assert!( - joins(&env, sql).iter().all(|j| !j.contains("prefers")), + // No preference: `bare` probes a `LeftSemi` that builds `mid`. + assert_eq!( + joins(&env, sql), + ["HashJoin LeftSemi, [a#0] = [a#3], RuntimeFilter #1: a#0 -> spark_catalog.default.bare.a"], "{sql}" ); let sql = "SELECT a FROM bare WHERE EXISTS (SELECT 1 FROM big WHERE big.a = bare.a)";