From ff203c9bf722f3eafff0b28387c51fe2a61d3a35 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Wed, 21 Oct 2026 18:32:46 +0000 Subject: [PATCH 1/8] fix(planning): keep grouped shared-subplan families to equal input rows The grouped family key of shared_subplan_materialization compared group keys, aggregate expressions and predicate-wildcarded sources only, so two aggregates whose inputs differed by a Filter or Join between the aggregate and the scan were merged and one branch read the other's aggregate. A CTE read twice, once under `WHERE col IS NOT NULL`, returned the unfiltered groups. The key now includes the input shape, computed against the columns the operators above read: sorts and column-only projections are skipped, a join keeps its type, keys, residual, null equality and the read output columns with their side, and filters and computed projections add the columns they read. Below a limit, window, union or any other operator, and for floating statistical aggregates, the shape is exact. TPC-DS q27's rollup branches, which differ only in column pruning and join output order, still merge; applied merges match the previous build on TPC-H and TPC-DS. Co-Authored-By: Claude Opus 5.5 --- .../shared_subplan_materialization.rs | 294 ++++++++++++++++++ .../shared_subplan_materialization/mod.rs | 15 +- .../signature.rs | 162 +++++++++- .../src/plan/identity.rs | 14 +- 4 files changed, 469 insertions(+), 16 deletions(-) diff --git a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/shared_subplan_materialization.rs b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/shared_subplan_materialization.rs index 96851626c..b64c817f1 100644 --- a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/shared_subplan_materialization.rs +++ b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/shared_subplan_materialization.rs @@ -676,6 +676,18 @@ fn t_shared_subplan_rejects_cases() { "optimizer_contracts::shared_subplan_materialization::native_optimizer_does_not_combine_scalar_branches_with_different_local_filters", native_optimizer_does_not_combine_scalar_branches_with_different_local_filters, ), + ( + "optimizer_contracts::shared_subplan_materialization::native_optimizer_does_not_combine_grouped_branches_with_different_local_filters", + native_optimizer_does_not_combine_grouped_branches_with_different_local_filters, + ), + ( + "optimizer_contracts::shared_subplan_materialization::native_optimizer_does_not_combine_grouped_branches_over_differently_ordered_limits", + native_optimizer_does_not_combine_grouped_branches_over_differently_ordered_limits, + ), + ( + "optimizer_contracts::shared_subplan_materialization::native_optimizer_does_not_combine_grouped_branches_reading_different_join_sides", + native_optimizer_does_not_combine_grouped_branches_reading_different_join_sides, + ), ( "optimizer_contracts::shared_subplan_materialization::native_optimizer_rejects_grouped_materialization_when_source_predicate_collides_with_aggregate_alias", native_optimizer_rejects_grouped_materialization_when_source_predicate_collides_with_aggregate_alias, @@ -738,6 +750,88 @@ fn native_optimizer_does_not_combine_scalar_branches_with_different_local_filter assert_eq!(diagnostics.applied_count, 0); } +/// One branch filters the shared source above the reader and the other does +/// not; one aggregate image cannot answer both. +fn native_optimizer_does_not_combine_grouped_branches_with_different_local_filters() { + let plan = repeated_grouped_count_with_and_without_local_filter_plan(); + let options = NativeOptimizerOptions { + shared_subplan_materialization: SharedSubplanMaterializationOptions::default(), + ..NativeOptimizerOptions::disabled() + }; + + let result = optimize_query_plan( + plan.clone(), + &options, + rapids_analysis_context(PlanningPolicy::default()), + local_dimension_pruning_resolver(), + ) + .unwrap(); + + assert_eq!(result.plan, plan); + assert_eq!(result.change, NativeOptimizerChange::Unchanged); + let diagnostics = &result.diagnostics.shared_subplan_materialization; + assert_eq!(diagnostics.considered_count, 2); + assert_eq!(diagnostics.candidate_count, 0); + assert_eq!(diagnostics.applied_count, 0); +} + +/// `ORDER BY a LIMIT 10` and `ORDER BY a DESC LIMIT 10` keep different rows, +/// so the sort below a limit is part of what the aggregate groups. +fn native_optimizer_does_not_combine_grouped_branches_over_differently_ordered_limits() { + assert_grouped_branches_stay_apart(repeated_grouped_count_over_ordered_limits_plan()); +} + +/// Both branches emit a join column named `x`, one from the left input and one +/// from the right; grouping by `x` groups different values. +fn native_optimizer_does_not_combine_grouped_branches_reading_different_join_sides() { + assert_grouped_branches_stay_apart(repeated_grouped_count_over_join_sides_plan()); +} + +fn assert_grouped_branches_stay_apart(plan: QueryPlan) { + let options = NativeOptimizerOptions { + shared_subplan_materialization: SharedSubplanMaterializationOptions::default(), + ..NativeOptimizerOptions::disabled() + }; + + let result = optimize_query_plan( + plan.clone(), + &options, + rapids_analysis_context(PlanningPolicy::default()), + local_dimension_pruning_resolver(), + ) + .unwrap(); + + assert_eq!(result.plan, plan); + let diagnostics = &result.diagnostics.shared_subplan_materialization; + assert_eq!(diagnostics.candidate_count, 0); + assert_eq!(diagnostics.applied_count, 0); +} + +/// Branches that join the same sources on the same keys but prune columns at +/// different points and order the join output from different sides group the +/// same rows, as TPC-DS q27's rollup branches do. +#[test] +fn native_optimizer_combines_grouped_join_branches_that_differ_only_in_column_pruning() { + let plan = repeated_grouped_join_aggregate_with_different_column_pruning_plan(); + let options = NativeOptimizerOptions { + shared_subplan_materialization: SharedSubplanMaterializationOptions::default(), + ..NativeOptimizerOptions::disabled() + }; + + let result = optimize_query_plan( + plan, + &options, + rapids_analysis_context(PlanningPolicy::default()), + local_dimension_pruning_resolver(), + ) + .unwrap(); + + let diagnostics = &result.diagnostics.shared_subplan_materialization; + assert_eq!(diagnostics.candidate_count, 1); + assert_eq!(diagnostics.applied_count, 1); + assert_eq!(result.plan.nodes().iter().filter(|node| matches!(node.kind, PlanNode::Aggregate(_))).count(), 1); +} + fn native_optimizer_rejects_grouped_materialization_when_source_predicate_collides_with_aggregate_alias() { let plan = repeated_grouped_aggregate_predicate_alias_collision_plan(); let options = NativeOptimizerOptions { @@ -1271,6 +1365,206 @@ fn estimated_sales_source_without_reader_predicate(bytes: u64) -> SourceNode { )) } +fn repeated_grouped_count_with_and_without_local_filter_plan() -> QueryPlan { + let mut builder = QueryPlanBuilder::new(); + let aggregate_schema = + Arc::new(Schema::new(vec![Field::new("g", DataType::Int32, false), Field::new("n", DataType::Int64, true)])); + let mut aggregates = Vec::new(); + for filtered in [true, false] { + let source = builder.add_source(estimated_local_source_with_stable_id( + "labels", + "labels", + &["g", "label"], + 512 * 1024 * 1024, + 10_000_000, + )); + let input = if filtered { + builder.add_filter(source, Expr::is_not_null(Expr::column("label"))).unwrap() + } else { + source + }; + let aggregate = builder + .add_aggregate( + input, + AggregateSpec::new( + vec![Expr::column("g")], + vec![AggregateExpr::new(AggregateFunction::Count, Expr::literal(LiteralValue::Int64(1)), "n")], + ), + Arc::clone(&aggregate_schema), + ) + .unwrap(); + aggregates.push(aggregate); + } + let union = builder.add_union(aggregates, aggregate_schema).unwrap(); + builder.finish_gpu_resident(union).unwrap() +} + +fn repeated_grouped_count_over_ordered_limits_plan() -> QueryPlan { + let mut builder = QueryPlanBuilder::new(); + let aggregate_schema = schema(&["g", "n"]); + let mut aggregates = Vec::new(); + for descending in [false, true] { + let source = builder.add_source(estimated_local_source_with_stable_id( + "ranked", + "ranked", + &["g", "a"], + 512 * 1024 * 1024, + 10_000_000, + )); + let sorted = + builder.add_sort(source, SortSpec::new(vec![SortExpr::new(Expr::column("a"), descending, false)])).unwrap(); + let limited = builder.add_limit(sorted, LimitSpec::fetch(10)).unwrap(); + aggregates.push( + builder + .add_aggregate( + limited, + AggregateSpec::new( + vec![Expr::column("g")], + vec![AggregateExpr::new(AggregateFunction::Count, Expr::literal(LiteralValue::Int64(1)), "n")], + ), + Arc::clone(&aggregate_schema), + ) + .unwrap(), + ); + } + let union = builder.add_union(aggregates, aggregate_schema).unwrap(); + builder.finish_gpu_resident(union).unwrap() +} + +fn repeated_grouped_count_over_join_sides_plan() -> QueryPlan { + let mut builder = QueryPlanBuilder::new(); + let aggregate_schema = schema(&["x", "n"]); + let mut aggregates = Vec::new(); + for side in [JoinSide::Left, JoinSide::Right] { + let left = builder.add_source(estimated_local_source_with_stable_id( + "join_left", + "join_left", + &["k", "x"], + 512 * 1024 * 1024, + 10_000_000, + )); + let right = builder.add_source(estimated_local_source_with_stable_id( + "join_right", + "join_right", + &["k", "x"], + 512 * 1024 * 1024, + 10_000_000, + )); + let index = usize::from(side == JoinSide::Right); + let joined = builder + .add_join( + left, + right, + JoinSpec::new( + JoinType::Inner, + vec![JoinKey::new(ColumnRef::input(0, "k", "k"), ColumnRef::input(1, "k", "k"))], + ) + .with_output_projection(vec![JoinOutputColumn::new( + side, + ColumnRef::input(index, "x", "x"), + "x", + )]), + schema(&["x"]), + ) + .unwrap(); + aggregates.push( + builder + .add_aggregate( + joined, + AggregateSpec::new( + vec![Expr::column("x")], + vec![AggregateExpr::new(AggregateFunction::Count, Expr::literal(LiteralValue::Int64(1)), "n")], + ), + Arc::clone(&aggregate_schema), + ) + .unwrap(), + ); + } + let union = builder.add_union(aggregates, aggregate_schema).unwrap(); + builder.finish_gpu_resident(union).unwrap() +} + +fn repeated_grouped_join_aggregate_with_different_column_pruning_plan() -> QueryPlan { + let mut builder = QueryPlanBuilder::new(); + let aggregate_schema = schema(&["customer_id", "store_sales"]); + let mut aggregates = Vec::new(); + for prune_after_join in [true, false] { + let fact = builder.add_source(estimated_local_source_with_stable_id( + "customer_sales", + "customer_sales", + &["customer_id", "dim_key", "sales"], + 512 * 1024 * 1024, + 10_000_000, + )); + let dimension = builder.add_source(estimated_local_source_with_stable_id( + "customer_dim", + "customer_dim", + &["dim_key"], + 1024 * 1024, + 10_000, + )); + let mut output_projection = vec![ + JoinOutputColumn::left(ColumnRef::input(0, "customer_id", "customer_id"), "customer_id"), + JoinOutputColumn::left(ColumnRef::input(0, "sales", "sales"), "sales"), + ]; + let mut order_side = JoinSide::Right; + if prune_after_join { + output_projection.push(JoinOutputColumn::left(ColumnRef::input(0, "dim_key", "dim_key"), "dim_key")); + order_side = JoinSide::Left; + } + let join_schema = + schema(&output_projection.iter().map(|output| output.output_name.as_str()).collect::>()); + let joined = builder + .add_join( + fact, + dimension, + JoinSpec::new( + JoinType::Inner, + vec![JoinKey::new( + ColumnRef::input(0, "dim_key", "dim_key"), + ColumnRef::input(1, "dim_key", "dim_key"), + )], + ) + .with_output_projection(output_projection) + .with_datafusion_order_preserved_side(order_side), + join_schema, + ) + .unwrap(); + let input = if prune_after_join { + builder + .add_projection( + joined, + ["customer_id", "sales"] + .into_iter() + .map(|name| { + ProjectionOutput::new( + Expr::column(name).alias(name), + Field::new(name, DataType::Int32, false), + ) + }) + .collect(), + ) + .unwrap() + } else { + joined + }; + aggregates.push( + builder + .add_aggregate( + input, + AggregateSpec::new( + vec![Expr::column("customer_id")], + vec![AggregateExpr::new(AggregateFunction::Sum, Expr::column("sales"), "store_sales")], + ), + Arc::clone(&aggregate_schema), + ) + .unwrap(), + ); + } + let union = builder.add_union(aggregates, aggregate_schema).unwrap(); + builder.finish_gpu_resident(union).unwrap() +} + fn repeated_grouped_aggregate_predicate_alias_collision_plan() -> QueryPlan { let mut builder = QueryPlanBuilder::new(); let aggregate_schema = schema(&["county", "quarter", "year", "store_sales"]); diff --git a/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/mod.rs b/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/mod.rs index c0c34e106..913d51baf 100644 --- a/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/mod.rs +++ b/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/mod.rs @@ -2,13 +2,14 @@ //! repeated aggregates over the same local Parquet sources. //! //! The rule matches two or more aggregate nodes whose group keys, aggregate -//! expressions and source set agree except for the sources' reader predicates (a -//! grouped family), or scalar aggregates without `FILTER` clauses over sources -//! that differ only in reader predicate (a scalar filtered family). In the -//! rewrite, the family's representative subtree is cloned once with its sources -//! widened to the union of the branch predicates and columns; each original -//! aggregate becomes a filter over that one aggregate (grouped) or a projection -//! that renames its branch's `FILTER`ed aggregate outputs (scalar). As a +//! expressions, input operators and source set agree except for the sources' +//! reader predicates (a grouped family), or scalar aggregates without `FILTER` +//! clauses over sources that differ only in reader predicate (a scalar +//! filtered family). In the rewrite, the family's representative subtree is +//! cloned once with its sources widened to the union of the branch predicates +//! and columns; each original aggregate becomes a filter over that one +//! aggregate (grouped) or a projection that renames its branch's `FILTER`ed +//! aggregate outputs (scalar). As a //! precondition, a grouped branch predicate must be evaluable after the //! aggregate, over group columns the source emits unchanged; a scalar branch //! predicate must reference source columns the widened subtree can carry to the diff --git a/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/signature.rs b/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/signature.rs index 30a3032e7..c5ca8031d 100644 --- a/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/signature.rs +++ b/crates/algeon-query-planning/src/optimizer/shared_subplan_materialization/signature.rs @@ -5,7 +5,9 @@ //! missing columns, scan limits. The abstract fingerprint adds what the source //! reads and emits. Both wildcard the reader predicate, and each has an exact //! twin that includes it. A grouped family keys on abstract fingerprints, so -//! members must read the same columns; a scalar filtered family keys on rowset +//! members must read the same columns, and on the operators between the +//! aggregate and its sources, so a filter, join or computed projection that +//! differs between members keeps them apart; a scalar filtered family keys on rowset //! fingerprints, because its rewrite re-reads the source with the union of the //! branches' columns. Signature strings are compared for equality only and are //! never shown in reports except through `family_key`. @@ -24,8 +26,8 @@ use super::{ source_widening::{local_parquet_source_node_occurrences_are_unique, scalar_source_occurrences_are_unique}, }; use crate::{ - expr::{AggregateExpr, AggregateFunction, AggregateSpec, CaseWhen, Expr}, - plan::{NodeId, PlanNode, QueryPlan}, + expr::{AggregateExpr, AggregateFunction, AggregateSpec, CaseWhen, ColumnQualifier, Expr}, + plan::{JoinNode, JoinSide, JoinSpec, NodeId, PlanNode, ProjectionNode, QueryPlan}, source::{ LocalParquetFile, LocalParquetSource, ParquetReaderDecimalType, ParquetSourceLocation, ResolvedParquetSource, SourceColumnRef, SourceNode, @@ -41,10 +43,10 @@ pub(super) struct AggregateMaterializationSignature { /// unchanged under the same name. A grouped branch predicate may be /// re-applied after the aggregate only over these. pub(super) post_aggregate_filter_columns_by_source: BTreeMap>, - /// Hash of the family signature: group keys, aggregate expressions and - /// predicate-wildcarded sources (grouped), or input shape and rowsets - /// (scalar). Equal keys mean the aggregates differ at most in reader - /// predicates. + /// Hash of the family signature: group keys, aggregate expressions, input + /// operators and predicate-wildcarded sources (grouped), or input shape + /// and rowsets (scalar). Equal keys mean the aggregates differ at most in + /// reader predicates. pub(super) family_key: String, pub(super) group_exprs: Vec, pub(super) aggregate_exprs: Vec, @@ -159,9 +161,22 @@ pub(super) fn aggregate_materialization_signature( source_fingerprints.join("|") ) } else { + // Every member reads the representative's aggregate, so members must + // compute the same rows apart from the reader predicates the rewrite + // factors. Over projections and sorts of one source, the canonical + // expressions already name source columns; any other input must match + // operator for operator. + let input_shape = if aggregate_alias_canonicalization_input_is_rowset_preserving(plan, input) { + "source_columns".to_string() + } else { + grouped_input_shape(plan, input, &aggregate_read_columns(spec))? + }; + let canonical_group_exprs = + spec.group_exprs.iter().map(|expr| aggregate_input_expr_signature(plan, input, expr)).collect::>(); format!( - "groups={};aggr={};sources={}", + "groups={};canonical_groups={};aggr={};input={input_shape};sources={}", group_exprs.join(","), + canonical_group_exprs.join(","), aggregate_exprs.join(","), source_fingerprints.join("|") ) @@ -185,6 +200,137 @@ pub(super) fn aggregate_materialization_signature( }) } +/// Columns of a node's output that the operators above it read. +enum ReadColumns { + /// Every column, by name and position, and the row order: an operator + /// above depends on more than the named values. + All, + /// Only these columns, and only as an unordered set of rows. + Named(BTreeSet), +} + +/// The columns a grouped aggregate reads, or [`ReadColumns::All`] when one of +/// its functions depends on input order. Floating statistical aggregates +/// replay DataFusion's input order to stay bit-exact. +fn aggregate_read_columns(spec: &AggregateSpec) -> ReadColumns { + if spec.aggr_exprs.iter().any(|aggregate| aggregate.function.is_floating_statistical()) { + return ReadColumns::All; + } + let exprs = spec + .group_exprs + .iter() + .chain(spec.aggr_exprs.iter().flat_map(|aggregate| std::iter::once(&aggregate.expr).chain(&aggregate.filter))); + ReadColumns::Named(exprs.flat_map(Expr::columns).map(|column| column.name().to_string()).collect()) +} + +/// Operator structure below a grouped aggregate, with each local Parquet +/// source reduced to its predicate-wildcarded abstract fingerprint. +/// +/// While only named columns of an unordered row set are read, the shape +/// describes those columns and nothing else: sorts and projections that only +/// select columns are left out, a join keeps its type, keys, residual +/// predicate and null equality plus the output columns read above it (with +/// the side and input column each comes from), and a filter or computed +/// projection adds the columns it reads. Branches over the same rows differ in +/// column pruning and in physical join choices that order the output (TPC-DS +/// q27's rollup branches); neither reaches a read column. Below an operator +/// that depends on order or positions (a limit, window or union), and below +/// anything else, the shape is exact operator for operator. +fn grouped_input_shape(plan: &QueryPlan, node_id: NodeId, read: &ReadColumns) -> Option { + let ReadColumns::Named(read) = read else { + return exact_input_shape(plan, node_id); + }; + let node = plan.node(node_id)?; + let (fingerprint, inputs) = match &node.kind { + PlanNode::Source(_) => return exact_input_shape(plan, node_id), + PlanNode::Sort(sort) => return grouped_input_shape(plan, sort.input, &ReadColumns::Named(read.clone())), + PlanNode::Projection(projection) if projection_only_selects_columns(projection) => { + return grouped_input_shape(plan, projection.input, &ReadColumns::Named(read.clone())); + }, + PlanNode::Projection(projection) => { + let input_read = projection.expressions().flat_map(Expr::columns).map(|column| column.name().to_string()); + (node.kind.detached_fingerprint(), vec![(projection.input, ReadColumns::Named(input_read.collect()))]) + }, + PlanNode::Filter(filter) => { + let mut input_read = read.clone(); + input_read.extend(filter.predicate.columns().map(|column| column.name().to_string())); + (node.kind.detached_fingerprint(), vec![(filter.input, ReadColumns::Named(input_read))]) + }, + PlanNode::Join(join) => { + let spec = &join.spec; + let mut outputs = spec + .output_projection + .iter() + .filter(|output| read.contains(&output.output_name)) + .cloned() + .collect::>(); + outputs.sort_by(|left, right| left.output_name.cmp(&right.output_name)); + let mut side_reads = [BTreeSet::new(), BTreeSet::new()]; + for output in &outputs { + let side = usize::from(output.side == JoinSide::Right); + side_reads[side].insert(output.column.name().to_string()); + } + for key in &spec.keys { + side_reads[0].insert(key.left.name().to_string()); + side_reads[1].insert(key.right.name().to_string()); + } + for column in spec.residual_predicate.iter().flat_map(Expr::columns) { + match column.qualifier() { + Some(ColumnQualifier::Input { index }) if *index < 2 => { + side_reads[*index].insert(column.name().to_string()); + }, + _ => side_reads.iter_mut().for_each(|side| { + side.insert(column.name().to_string()); + }), + } + } + let mut rows = JoinSpec::new(spec.join_type.clone(), spec.keys.clone()) + .with_null_equality(spec.null_equality) + .with_output_projection(outputs); + if let Some(predicate) = &spec.residual_predicate { + rows = rows.with_residual_predicate(predicate.clone()); + } + let [left_read, right_read] = side_reads; + ( + PlanNode::Join(JoinNode { inputs: join.inputs, spec: rows }).detached_fingerprint(), + vec![(join.inputs[0], ReadColumns::Named(left_read)), (join.inputs[1], ReadColumns::Named(right_read))], + ) + }, + _ => return exact_input_shape(plan, node_id), + }; + let inputs = inputs + .iter() + .map(|(input, input_read)| grouped_input_shape(plan, *input, input_read)) + .collect::>>()?; + Some(format!("{fingerprint}({})", inputs.join(","))) +} + +/// Operator-for-operator shape of the subtree at `node_id`, with each local +/// Parquet source reduced to its predicate-wildcarded abstract fingerprint. +fn exact_input_shape(plan: &QueryPlan, node_id: NodeId) -> Option { + let node = plan.node(node_id)?; + if let PlanNode::Source(source) = &node.kind + && let SourceNode::LocalParquet(source) = source.as_ref() + { + let fingerprint = local_source_materialization_signature(source).abstract_fingerprint; + return Some(format!("local_parquet({fingerprint})")); + } + let inputs = node.input_ids().iter().map(|input| exact_input_shape(plan, *input)).collect::>>()?; + Some(format!("{}({})", node.kind.detached_fingerprint(), inputs.join(","))) +} + +/// Whether every output of `projection` is an unqualified input column under +/// its own name, possibly behind aliases. +fn projection_only_selects_columns(projection: &ProjectionNode) -> bool { + projection.outputs.iter().all(|output| { + let mut expr = &output.expr; + while let Expr::Alias { expr: inner, .. } = expr { + expr = inner; + } + matches!(expr, Expr::Column(column) if column.qualifier().is_none() && column.name() == output.field.name()) + }) +} + /// Group output columns that are a source column under the same name, keyed /// by that source's abstract fingerprint. A predicate over such a column /// selects the same groups before or after the aggregate. diff --git a/crates/algeon-query-planning/src/plan/identity.rs b/crates/algeon-query-planning/src/plan/identity.rs index 7832479c2..2a154d866 100644 --- a/crates/algeon-query-planning/src/plan/identity.rs +++ b/crates/algeon-query-planning/src/plan/identity.rs @@ -25,7 +25,7 @@ use arrow::{ record_batch::RecordBatch, }; -use super::{Annotations, PlanNode, QueryNode, QueryPlan}; +use super::{Annotations, NodeId, PlanNode, QueryNode, QueryPlan, clone_plan_node_with_inputs}; use crate::source::SourceNode; mod impls; @@ -310,6 +310,18 @@ impl QueryPlan { } } +impl PlanNode { + /// Digest of this operator alone: every input id is hashed as one + /// placeholder, so operators in different subtrees that compute the same + /// thing from their inputs compare equal. + pub(crate) fn detached_fingerprint(&self) -> PlanFingerprint { + let detached = clone_plan_node_with_inputs(self, |_| NodeId(0)); + let mut hasher = PlanHasher::new(); + detached.stable_hash(&mut hasher); + PlanFingerprint(hasher.finish_digest()) + } +} + impl QueryNode { fn masked(&self, slots: &DeferredGraphSourceFactSlots) -> QueryNode { let mut node = self.clone(); From 1ab47a23bbdcfa645c54c2b4144340af32cb9956 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Wed, 21 Oct 2026 18:32:56 +0000 Subject: [PATCH 2/8] fix(chunk): keep each source read within the backend's column element limit A FixedSizeList column with 2.9M rows (5.9e9 elements) failed every scan, filtered or limited, because one planned source chunk spanned enough row groups that the cuDF shim's concatenation of reader passes exceeded cudf::size_type. The client saw "result size computation overflowed the platform range" with no column or remedy. BackendCapabilities gains max_column_elements (cuDF: i32::MAX). Chunk, batch, retry and scan-limit planning keep every read's largest column under it, counted from Parquet footer leaf value counts, or rows times fixed list widths without a footer. A row group that alone exceeds the limit fails planning with a typed size_overflow naming the column and the row-group size that fits. libcudf concatenation is pre-checked per column, and the client message names the column, element count, limit and the row count that fits at its width. One relation input above the limit is still unsupported and now fails with that message. Co-Authored-By: Claude Opus 5.5 --- .config/nextest.toml | 3 + .config/policies/public-api/algeon-cudf.txt | 2 + .../public-api/algeon-query-engine.txt | 1 + .../public-api/algeon-query-planning.txt | 4 + .../cudf/algeon-cudf/src/core/column/owned.rs | 5 + .../src/frame/dataframe/concat_ops.rs | 11 +- .../cudf/algeon-cudf/src/ops/concatenate.rs | 236 +++++++++++++++++- crates/algeon-datafusion/src/error/client.rs | 42 +++- crates/algeon-datafusion/src/server/error.rs | 50 ++++ crates/algeon-query-engine/src/chunk/tests.rs | 162 +++++++++++- .../exec/bounded/aggregate/source_adjacent.rs | 3 + .../src/exec/bounded/chunks/retry.rs | 11 + .../src/exec/bounded/chunks/source_read.rs | 3 + .../src/exec/bounded/facts.rs | 3 + .../src/exec/bounded/join/leaf/batching.rs | 3 + .../src/exec/source/cudf_contract/tests.rs | 9 + .../src/exec/source/lakehouse_cache/mod.rs | 3 + .../src/exec/source/metrics.rs | 3 + .../algeon-query-engine/src/rapids_backend.rs | 5 + .../src/source_chunk/local_parquet_plan.rs | 5 +- crates/algeon-query-planning/src/backend.rs | 8 + .../src/chunk/assembly.rs | 162 +++++++++++- .../src/chunk/elements.rs | 129 ++++++++++ crates/algeon-query-planning/src/chunk/mod.rs | 12 +- .../algeon-query-planning/src/chunk/model.rs | 10 + .../src/chunk/planner.rs | 24 ++ .../src/chunk/row_group.rs | 16 +- .../src/chunk/scan_limit.rs | 5 + .../algeon-query-runtime/src/native_error.rs | 19 ++ .../runtime/source_identity/local_parquet.rs | 3 + website/docs/cuml/ml-inputs.mdx | 13 + website/docs/cuvs/vector-inputs.mdx | 18 ++ website/docs/design/data-path.mdx | 8 +- website/docs/integration/handling-errors.mdx | 11 + website/docs/reference/data-type-support.md | 9 +- 35 files changed, 989 insertions(+), 22 deletions(-) create mode 100644 crates/algeon-query-planning/src/chunk/elements.rs diff --git a/.config/nextest.toml b/.config/nextest.toml index 331c3af2a..004f1366e 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -501,6 +501,9 @@ binary(/_cpu$/) | internal::resolve::tests::validate_no_duplicate_names_rejects_dup | io::tests::test_collect_column_names_rejects_negative_count | ops::reshape::tests::interleave_output_size_overflow_is_structured + | ops::concatenate::tests::concatenated_vector_frames_overflow_at_the_list_child_with_the_width_row_cap + | ops::concatenate::tests::concatenation_up_to_the_width_row_cap_is_accepted_and_one_row_more_is_rejected + | ops::concatenate::tests::concatenated_flat_rows_overflow_at_the_row_count | runtime::device_health::tests::cross_device_teardown_suspends_instead_of_quarantining | diagnostics::memory::tests::fixed_pressure_plan_is_bounded_and_preserves_the_documented_reserve | runtime::matrix::tests::matrix_element_count_accepts_cudf_size_type_maximum diff --git a/.config/policies/public-api/algeon-cudf.txt b/.config/policies/public-api/algeon-cudf.txt index 4b9c88480..cf5372a91 100644 --- a/.config/policies/public-api/algeon-cudf.txt +++ b/.config/policies/public-api/algeon-cudf.txt @@ -3,6 +3,7 @@ pub mod cudf::core pub mod cudf::core::column pub struct cudf::core::column::Column impl cudf::core::Column +pub const cudf::core::Column::MAX_ELEMENTS: usize pub fn cudf::core::Column::data_type(&self) -> cudf::core::error::Result pub fn cudf::core::Column::deep_copy(&self) -> cudf::core::error::Result pub fn cudf::core::Column::device_data_ptr(&self) -> cudf::core::error::Result<*const std::os::raw::c_void> @@ -1030,6 +1031,7 @@ impl core::fmt::Display for cudf::core::types::TypeId pub fn cudf::core::types::TypeId::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct cudf::core::Column impl cudf::core::Column +pub const cudf::core::Column::MAX_ELEMENTS: usize pub fn cudf::core::Column::data_type(&self) -> cudf::core::error::Result pub fn cudf::core::Column::deep_copy(&self) -> cudf::core::error::Result pub fn cudf::core::Column::device_data_ptr(&self) -> cudf::core::error::Result<*const std::os::raw::c_void> diff --git a/.config/policies/public-api/algeon-query-engine.txt b/.config/policies/public-api/algeon-query-engine.txt index cecae9625..41cf86a7a 100644 --- a/.config/policies/public-api/algeon-query-engine.txt +++ b/.config/policies/public-api/algeon-query-engine.txt @@ -40,6 +40,7 @@ pub fn query_engine::rapids_backend::RapidsBackend::aggregate_output_names(&self pub fn query_engine::rapids_backend::RapidsBackend::binary_output_type(&self, query_planning::expr::BinaryOp, &arrow_schema::datatype::DataType, &arrow_schema::datatype::DataType) -> query_planning::error::Result pub fn query_engine::rapids_backend::RapidsBackend::check_expression(&self, &query_planning::expr::Expr, query_planning::backend::ExpressionUse) -> query_planning::error::Result<()> pub fn query_engine::rapids_backend::RapidsBackend::check_projection_output(&self, &query_planning::plan::nodes::ProjectionOutput) -> query_planning::error::Result<()> +pub fn query_engine::rapids_backend::RapidsBackend::max_column_elements(&self) -> u64 pub fn query_engine::rapids_backend::RapidsBackend::supports_type(&self, &arrow_schema::datatype::DataType) -> bool pub fn query_engine::rapids_backend::rapids_analysis_context(query_planning::policy::PlanningPolicy) -> query_planning::backend::AnalysisContext<'static> pub const fn query_engine::rapids_backend::rapids_backend() -> &'static query_engine::rapids_backend::RapidsBackend diff --git a/.config/policies/public-api/algeon-query-planning.txt b/.config/policies/public-api/algeon-query-planning.txt index fa1c16c39..9b7ed9ca1 100644 --- a/.config/policies/public-api/algeon-query-planning.txt +++ b/.config/policies/public-api/algeon-query-planning.txt @@ -40,6 +40,7 @@ pub fn query_planning::backend::BackendCapabilities::aggregate_output_names(&sel pub fn query_planning::backend::BackendCapabilities::binary_output_type(&self, query_planning::expr::BinaryOp, &arrow_schema::datatype::DataType, &arrow_schema::datatype::DataType) -> query_planning::error::Result pub fn query_planning::backend::BackendCapabilities::check_expression(&self, &query_planning::expr::Expr, query_planning::backend::ExpressionUse) -> query_planning::error::Result<()> pub fn query_planning::backend::BackendCapabilities::check_projection_output(&self, &query_planning::plan::ProjectionOutput) -> query_planning::error::Result<()> +pub fn query_planning::backend::BackendCapabilities::max_column_elements(&self) -> u64 pub fn query_planning::backend::BackendCapabilities::supports_type(&self, &arrow_schema::datatype::DataType) -> bool pub mod query_planning::capability pub enum query_planning::capability::AggregateExecutionStage @@ -381,6 +382,7 @@ pub fn query_planning::chunk::decoded::row_group_decoded_bytes_from_footer(&parq pub struct query_planning::chunk::SourceChunk pub query_planning::chunk::SourceChunk::chunk_index: usize pub query_planning::chunk::SourceChunk::estimated_chunk_bytes: core::option::Option +pub query_planning::chunk::SourceChunk::estimated_column_elements: core::option::Option pub query_planning::chunk::SourceChunk::estimated_decoded_bytes: core::option::Option pub query_planning::chunk::SourceChunk::estimated_source_read_bytes: core::option::Option pub query_planning::chunk::SourceChunk::exact_source_read_preflight_bytes: core::option::Option @@ -393,6 +395,7 @@ pub query_planning::chunk::SourceChunk::known_pure_filter_uncompressed_bytes: co pub query_planning::chunk::SourceChunk::known_row_count: core::option::Option pub query_planning::chunk::SourceChunk::known_row_group_bytes: core::option::Option pub query_planning::chunk::SourceChunk::location: query_planning::source::ParquetSourceLocation +pub query_planning::chunk::SourceChunk::max_column_elements: u64 pub query_planning::chunk::SourceChunk::partition_values: alloc::vec::Vec pub query_planning::chunk::SourceChunk::path: std::path::PathBuf pub query_planning::chunk::SourceChunk::read_columns: alloc::vec::Vec @@ -411,6 +414,7 @@ pub query_planning::chunk::SourceChunk::source_file_projected_uncompressed_bytes pub query_planning::chunk::SourceChunk::source_file_pure_filter_compressed_bytes: core::option::Option> pub query_planning::chunk::SourceChunk::source_file_pure_filter_uncompressed_bytes: core::option::Option> pub query_planning::chunk::SourceChunk::source_file_row_group_bytes: core::option::Option> +pub query_planning::chunk::SourceChunk::source_file_row_group_column_elements: core::option::Option> pub query_planning::chunk::SourceChunk::source_file_row_group_decoded_bytes: core::option::Option> pub query_planning::chunk::SourceChunk::source_file_row_group_rows: core::option::Option> pub query_planning::chunk::SourceChunk::source_file_row_groups: core::option::Option> diff --git a/components/cudf/algeon-cudf/src/core/column/owned.rs b/components/cudf/algeon-cudf/src/core/column/owned.rs index dd4716b71..3a70dcb22 100644 --- a/components/cudf/algeon-cudf/src/core/column/owned.rs +++ b/components/cudf/algeon-cudf/src/core/column/owned.rs @@ -144,6 +144,11 @@ unsafe impl Send for Column {} unsafe impl Sync for Column {} impl Column { + /// Largest element count of one column at any level: its row count and, + /// for a nested column, each child's element count. libcudf indexes + /// columns with the 32-bit `cudf::size_type`. + pub const MAX_ELEMENTS: usize = i32::MAX as usize; + pub(crate) fn backing_allocation_pointers(&self) -> Result> { let operation = "column::backing_allocation_pointers"; let mut count = 0usize; diff --git a/components/cudf/algeon-cudf/src/frame/dataframe/concat_ops.rs b/components/cudf/algeon-cudf/src/frame/dataframe/concat_ops.rs index fa94b8fd7..bc29d374e 100644 --- a/components/cudf/algeon-cudf/src/frame/dataframe/concat_ops.rs +++ b/components/cudf/algeon-cudf/src/frame/dataframe/concat_ops.rs @@ -227,10 +227,15 @@ impl DataFrame { table_refs.push(tv); } - let result_table = - ops::concatenate::concatenate_tables_with_exec(&table_refs, &output_domain, execution.exec_options())?; + let column_names = self.column_names_owned(); + let result_table = ops::concatenate::concatenate_tables_with_exec( + &table_refs, + &column_names, + &output_domain, + execution.exec_options(), + )?; - Self::from_table_with_names(result_table, self.column_names_owned()) + Self::from_table_with_names(result_table, column_names) } pub(crate) fn validate_concat_schema(&self, other: &DataFrame, operation: &'static str) -> Result<()> { diff --git a/components/cudf/algeon-cudf/src/ops/concatenate.rs b/components/cudf/algeon-cudf/src/ops/concatenate.rs index 5a6170746..180ebe917 100644 --- a/components/cudf/algeon-cudf/src/ops/concatenate.rs +++ b/components/cudf/algeon-cudf/src/ops/concatenate.rs @@ -20,6 +20,7 @@ use crate::{ column::{Column, ColumnView}, error::{self, Error, ErrorCode, FACT_COLUMN_INDEX, FACT_LEFT_TYPE_ID, FACT_RIGHT_TYPE_ID, Result}, table::{Table, TableView}, + types::TypeId, }, runtime::{execution_context::Exec, execution_domain::ExecutionDomain, execution_resolver::ResolvedExecution}, }; @@ -129,6 +130,142 @@ fn validate_table_column_types(tables: &[&TableView<'_>], operation: &'static st Ok(()) } +/// Child index of a libcudf list column's values; child 0 holds its offsets. +const LIST_VALUES_CHILD: usize = 1; + +/// Element counts of one column: its length, then the counted children +/// beneath it. A list counts its values child and a struct counts each field. +/// A string column counts no child, since libcudf stores large string +/// characters with 64-bit offsets. +#[derive(Debug, Clone, PartialEq, Eq)] +struct ColumnElementCounts { + elements: usize, + children: Vec, +} + +impl ColumnElementCounts { + /// Reads the counts from view metadata. With `exact` false, a list counts + /// its whole values child, which overstates a sliced list and needs no + /// device read. With `exact` true, a list counts only the values its rows + /// reference, which reads two offsets per list level from the device. + fn of(view: &ColumnView<'_>, exact: bool) -> Result { + let children = match view.data_type()?.id() { + TypeId::List => { + let values = if exact { view.list_sliced_child()? } else { view.child(LIST_VALUES_CHILD)? }; + vec![Self::of(&values, exact)?] + }, + TypeId::Struct => { + (0..view.num_children()?).map(|index| Self::of(&view.child(index)?, exact)).collect::>()? + }, + _ => Vec::new(), + }; + Ok(Self { elements: view.len()?, children }) + } +} + +/// A level of one concatenated column whose summed element count exceeds the +/// per-column limit. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ConcatenatedElementOverflow { + /// Rows of the concatenated column. + rows: u64, + /// Elements at the overflowing level, summed over the inputs. + elements: u64, + /// Nesting depth of that level; 0 is the row count itself. + depth: usize, +} + +impl ConcatenatedElementOverflow { + /// Largest row count whose elements at this level stay within `limit`, + /// keeping the concatenated column's elements per row. + fn row_cap(self, limit: usize) -> u64 { + let limit = limit as u64; + let row_cap = u128::from(self.rows) * u128::from(limit) / u128::from(self.elements.max(1)); + u64::try_from(row_cap).unwrap_or(u64::MAX).min(limit) + } +} + +/// The first level, depth first, at which `inputs` (one column's counts per +/// concatenation input) sum above `limit`. Inputs share one type, so their +/// count trees share one shape. +fn first_concatenated_element_overflow( + inputs: &[ColumnElementCounts], + limit: usize, +) -> Option { + fn visit( + levels: &[&ColumnElementCounts], + rows: u64, + depth: usize, + limit: usize, + ) -> Option { + let elements = levels.iter().map(|level| level.elements as u64).fold(0u64, u64::saturating_add); + if elements > limit as u64 { + return Some(ConcatenatedElementOverflow { rows, elements, depth }); + } + let child_count = levels.first().map_or(0, |level| level.children.len()); + (0..child_count).find_map(|index| { + let children = levels.iter().filter_map(|level| level.children.get(index)).collect::>(); + visit(&children, rows, depth + 1, limit) + }) + } + let roots = inputs.iter().collect::>(); + let rows = inputs.iter().map(|input| input.elements as u64).fold(0u64, u64::saturating_add); + visit(&roots, rows, 0, limit) +} + +fn concatenated_column_element_limit_error( + operation: &'static str, + column: Option<(usize, Option<&str>)>, + overflow: ConcatenatedElementOverflow, +) -> Error { + let row_cap = overflow.row_cap(Column::MAX_ELEMENTS); + let error = Error::new_with( + ErrorCode::SizeOverflow, + operation, + "concatenated column holds more elements than one libcudf column can index", + ); + let error = match column { + Some((index, Some(name))) => error.with_fact(FACT_COLUMN_INDEX, index).with_column(name.to_owned()), + Some((index, None)) => error.with_fact(FACT_COLUMN_INDEX, index), + None => error, + }; + error + .with_detail("element_count", overflow.elements.to_string()) + .with_detail("element_limit", Column::MAX_ELEMENTS.to_string()) + .with_detail("nesting_depth", overflow.depth.to_string()) + .with_detail("row_count", overflow.rows.to_string()) + .with_detail("row_cap", row_cap.to_string()) + .with_detail("reason", "concatenated_column_elements_exceed_cudf_size_type") + .with_detail("remedy", format!("keep each concatenated relation at or below {row_cap} rows")) +} + +/// Rejects a concatenation whose output column would exceed +/// [`Column::MAX_ELEMENTS`] at any level, before libcudf allocates it. +/// +/// The first pass reads view metadata only. A list input sliced from a larger +/// column can overstate its count there, so a pass that finds an overflow is +/// repeated with exact counts, which read list offsets from the device. +fn check_concatenated_column_elements( + inputs: &[&ColumnView<'_>], + operation: &'static str, + column: Option<(usize, Option<&str>)>, +) -> Result<()> { + let counts = |exact| { + inputs + .iter() + .map(|view| ColumnElementCounts::of(view, exact)) + .collect::>>() + .map_err(|error| error.reanchor(operation)) + }; + if first_concatenated_element_overflow(&counts(false)?, Column::MAX_ELEMENTS).is_none() { + return Ok(()); + } + match first_concatenated_element_overflow(&counts(true)?, Column::MAX_ELEMENTS) { + Some(overflow) => Err(concatenated_column_element_limit_error(operation, column, overflow)), + None => Ok(()), + } +} + fn compatible_column_execution(columns: &[&ColumnView<'_>]) -> Result { let Some(first) = columns.first() else { return Err(concatenate_requires_input_error("concatenate::concatenate_columns", "column")); @@ -156,6 +293,7 @@ pub(crate) fn concatenate_columns_with_exec( exec_options: &cudf_sys::cudf_exec_options_t, ) -> Result { validate_column_types(columns, "concatenate::concatenate_columns")?; + check_concatenated_column_elements(columns, "concatenate::concatenate_columns", None)?; let handles: Vec = columns.iter().map(|c| c.raw_handle()).collect(); let num_columns = checked_len_to_i32( handles.len(), @@ -177,12 +315,17 @@ pub(crate) fn concatenate_columns_with_exec( unsafe { Column::from_raw_with_domain(out, Arc::clone(output_domain)) } } +/// Concatenates `tables` row-wise. `column_names` names the shared columns in +/// order for diagnostics and may be shorter than the column count. pub(crate) fn concatenate_tables_with_exec( tables: &[&TableView<'_>], + column_names: &[String], output_domain: &Arc, exec_options: &cudf_sys::cudf_exec_options_t, ) -> Result { - validate_table_column_types(tables, "concatenate::concatenate_tables")?; + let operation = "concatenate::concatenate_tables"; + validate_table_column_types(tables, operation)?; + check_concatenated_table_elements(tables, column_names, operation)?; let handles: Vec = tables.iter().map(|t| t.raw_handle()).collect(); let num_tables = checked_len_to_i32( handles.len(), @@ -204,6 +347,35 @@ pub(crate) fn concatenate_tables_with_exec( unsafe { Table::from_raw_with_domain(out, Arc::clone(output_domain)) } } +/// Applies [`check_concatenated_column_elements`] to each column position of +/// `tables`. Tables whose column counts differ are left to libcudf, which +/// rejects them. +fn check_concatenated_table_elements( + tables: &[&TableView<'_>], + column_names: &[String], + operation: &'static str, +) -> Result<()> { + let Some(first) = tables.first() else { + return Ok(()); + }; + let column_count = first.num_columns().map_err(|error| error.reanchor(operation))?; + for table in tables { + if table.num_columns().map_err(|error| error.reanchor(operation))? != column_count { + return Ok(()); + } + } + for index in 0..column_count { + let views = tables + .iter() + .map(|table| table.column(index).map_err(|error| error.reanchor(operation))) + .collect::>>()?; + let view_refs = views.iter().collect::>(); + let name = column_names.get(index).map(String::as_str); + check_concatenated_column_elements(&view_refs, operation, Some((index, name)))?; + } + Ok(()) +} + #[cfg(any(test, feature = "test-hooks"))] pub(crate) fn trigger_concatenate_column_count_limit_error_for_test() -> Result<()> { Err(input_column_count_limit_error("concatenate::concatenate_columns", usize::MAX)) @@ -213,3 +385,65 @@ pub(crate) fn trigger_concatenate_column_count_limit_error_for_test() -> Result< pub(crate) fn trigger_concatenate_table_count_limit_error_for_test() -> Result<()> { Err(input_table_count_limit_error("concatenate::concatenate_tables", usize::MAX)) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::error::{CudfErrorClassification, FactKeyName, classify_cudf_error}; + + const VECTOR_WIDTH: usize = 2048; + + fn vector_frame(rows: usize) -> ColumnElementCounts { + ColumnElementCounts { + elements: rows, + children: vec![ColumnElementCounts { elements: rows * VECTOR_WIDTH, children: Vec::new() }], + } + } + + fn flat_frame(rows: usize) -> ColumnElementCounts { + ColumnElementCounts { elements: rows, children: Vec::new() } + } + + #[test] + fn concatenated_vector_frames_overflow_at_the_list_child_with_the_width_row_cap() { + let frames = vec![vector_frame(200_000); 15]; + + let overflow = first_concatenated_element_overflow(&frames, Column::MAX_ELEMENTS) + .expect("3,000,000 rows of 2048 values exceed the per-column limit"); + assert_eq!( + overflow, + ConcatenatedElementOverflow { rows: 3_000_000, elements: 3_000_000 * VECTOR_WIDTH as u64, depth: 1 } + ); + + let error = concatenated_column_element_limit_error("concatenate::test", Some((1, Some("fp"))), overflow); + assert_eq!(classify_cudf_error(&error), CudfErrorClassification::SizeOverflow); + assert_eq!(error.fact_value(FactKeyName::ColumnName).map(ToString::to_string).as_deref(), Some("fp")); + assert_eq!(error.detail_value("element_count"), Some("6144000000")); + assert_eq!(error.detail_value("element_limit"), Some("2147483647")); + assert_eq!(error.detail_value("row_cap"), Some("1048575")); + } + + #[test] + fn concatenation_up_to_the_width_row_cap_is_accepted_and_one_row_more_is_rejected() { + let row_cap = Column::MAX_ELEMENTS / VECTOR_WIDTH; + let at_cap = [vector_frame(row_cap - 1), vector_frame(1)]; + assert_eq!(first_concatenated_element_overflow(&at_cap, Column::MAX_ELEMENTS), None); + + let past_cap = [vector_frame(row_cap), vector_frame(1)]; + let overflow = first_concatenated_element_overflow(&past_cap, Column::MAX_ELEMENTS) + .expect("one row past the cap exceeds the per-column limit"); + assert_eq!(overflow.depth, 1); + assert_eq!(overflow.row_cap(Column::MAX_ELEMENTS), row_cap as u64); + } + + #[test] + fn concatenated_flat_rows_overflow_at_the_row_count() { + let half = Column::MAX_ELEMENTS / 2 + 1; + let overflow = first_concatenated_element_overflow(&[flat_frame(half), flat_frame(half)], Column::MAX_ELEMENTS) + .expect("rows past the limit overflow"); + + assert_eq!(overflow.depth, 0); + assert_eq!(overflow.elements, 2 * half as u64); + assert_eq!(overflow.row_cap(Column::MAX_ELEMENTS), Column::MAX_ELEMENTS as u64); + } +} diff --git a/crates/algeon-datafusion/src/error/client.rs b/crates/algeon-datafusion/src/error/client.rs index d73e3489d..17d4595d0 100644 --- a/crates/algeon-datafusion/src/error/client.rs +++ b/crates/algeon-datafusion/src/error/client.rs @@ -71,6 +71,11 @@ pub(crate) fn client_message(error: &Error) -> ClientMessage { let mut logged = client_status_message(error); let detail = client_actionable_kind(error).then(|| error.client_detail()).flatten(); let mut text = detail.map_or_else(|| logged.clone(), |detail| client_schema_detail_message(detail, &logged)); + if error.code() == ErrorCode::SizeOverflow + && let Some(column) = error.fact_value("column") + { + text = size_overflow_message(error, Some(column)); + } if error.code() == ErrorCode::FinalPlanRequirementUnsatisfied { logged.insert_str(0, NATIVE_REQUIRED_REJECTION); text.insert_str(0, NATIVE_REQUIRED_REJECTION); @@ -129,10 +134,7 @@ fn client_status_message(error: &Error) -> String { "query budget cannot fit on any eligible device: {}. Request a smaller query cap or wait for device capacity to free; do not retry the same request unchanged", error.message() ), - ErrorCode::SizeOverflow => format!( - "result size computation overflowed the platform range: {}. Change the request shape to produce a smaller result; retrying the same request unchanged cannot succeed", - error.message() - ), + ErrorCode::SizeOverflow => size_overflow_message(error, None), ErrorCode::HostMaterializationLimitExceeded => format!( "one result row exceeds the configured host materialization limit: {}. Narrow the request or raise the limit; retrying the same request unchanged cannot succeed", error.message() @@ -205,6 +207,38 @@ fn memory_failure_detail(error: &Error) -> Option { Some(format!("{}{figures}", location.unwrap_or_default())) } +/// Client text of a size overflow. When the native layers recorded the +/// per-column element limit a read or an intermediate result exceeded, the +/// text names the element count, the limit and the change that fits under it: +/// smaller Parquet row groups when one row group alone exceeds the limit, +/// otherwise a filter or a split at the row cap. `column` is the caller's +/// column name, which joins only the wire text. Counts are parsed as numbers, +/// so no foreign text enters the message. +fn size_overflow_message(error: &Error, column: Option<&str>) -> String { + const NO_RETRY: &str = "retrying the same request unchanged cannot succeed"; + let count = |key: &str| error.fact_value(key).and_then(|value| value.parse::().ok()); + let (Some(elements), Some(limit)) = (count("element_count"), count("element_limit")) else { + return format!( + "result size computation overflowed the platform range: {}. Change the request shape to produce a smaller result; {NO_RETRY}", + error.message() + ); + }; + let subject = column.map_or_else(|| "one column".to_string(), |column| format!("column `{column}`")); + if let Some(rows) = count("max_row_group_rows") { + return format!( + "one Parquet row group of {subject} decodes to {elements} elements, above the limit of {limit} elements per column. Rewrite the file with row groups of at most {rows} rows; {NO_RETRY}" + ); + } + match count("row_cap") { + Some(row_cap) => format!( + "{subject} would hold {elements} elements in one relation or intermediate result, above the limit of {limit} elements per column; at its width at most {row_cap} rows fit. Filter the input, or split the relation (for example a UNION ALL of calls over row ranges) so each part holds at most {row_cap} rows; {NO_RETRY}" + ), + None => format!( + "{subject} would hold {elements} elements in one relation or intermediate result, above the limit of {limit} elements per column. Filter the input, or split the relation (for example a UNION ALL of calls over row ranges); {NO_RETRY}" + ), + } +} + /// Client text for rejections that name the caller's own schema identifiers. /// /// Owners keep column and field names and their Arrow types out of diff --git a/crates/algeon-datafusion/src/server/error.rs b/crates/algeon-datafusion/src/server/error.rs index 6cede9014..cfc1b2039 100644 --- a/crates/algeon-datafusion/src/server/error.rs +++ b/crates/algeon-datafusion/src/server/error.rs @@ -1373,6 +1373,56 @@ mod tests { } } + /// A native size overflow that names the per-column element limit, wrapped + /// at the adapter boundary like `native_memory_failure`. + fn native_element_limit_failure(limit_facts: &[(&'static str, u64)]) -> Error { + let native = limit_facts.iter().fold( + QueryError::new( + QueryErrorCode::SizeOverflow, + "query_engine.local_parquet_chunk.read", + "secret native text naming s3://bucket/private/fingerprints.parquet", + ) + .with_fact("column", "fp") + .with_fact("element_count", 5_934_733_312_u64) + .with_fact("element_limit", 2_147_483_647_u64), + |native, (key, value)| native.with_fact(key, value), + ); + Error::from_native("algeon_datafusion.native_exec.execute", "Algeon native query execution failed", native) + } + + #[test] + fn element_limit_overflow_names_the_column_limit_and_row_cap_on_the_wire_only() { + let error = native_element_limit_failure(&[("row_cap", 1_048_575)]); + let status = adapter_error_to_status(&error); + + assert_eq!(status.code(), Code::ResourceExhausted); + assert_eq!( + status.message(), + "column `fp` would hold 5934733312 elements in one relation or intermediate result, above the limit of 2147483647 elements per column; at its width at most 1048575 rows fit. Filter the input, or split the relation (for example a UNION ALL of calls over row ranges) so each part holds at most 1048575 rows; retrying the same request unchanged cannot succeed" + ); + assert_error_metadata(&status, "size_overflow", "resource_exhausted", "permanent"); + let logged = client_message(&error).logged; + assert!(logged.starts_with("one column would hold 5934733312 elements"), "{logged}"); + assert!(!logged.contains("`fp`"), "{logged}"); + for hidden in ["secret native text", "s3://"] { + assert!(!status.message().contains(hidden), "{}", status.message()); + } + assert_no_diagnostic_fields(status.message()); + } + + #[test] + fn oversized_row_group_overflow_names_the_row_group_size_that_fits() { + let status = adapter_error_to_status(&native_element_limit_failure(&[ + ("row_group_rows", 2_897_819), + ("max_row_group_rows", 1_048_575), + ])); + + assert_eq!( + status.message(), + "one Parquet row group of column `fp` decodes to 5934733312 elements, above the limit of 2147483647 elements per column. Rewrite the file with row groups of at most 1048575 rows; retrying the same request unchanged cannot succeed" + ); + } + #[test] fn device_allocation_failures_name_the_same_facts_and_keep_the_sealed_report_pointer() { for code in [QueryErrorCode::DeviceOutOfMemory, QueryErrorCode::DeviceAllocationFailed] { diff --git a/crates/algeon-query-engine/src/chunk/tests.rs b/crates/algeon-query-engine/src/chunk/tests.rs index f007edafb..6b995acf1 100644 --- a/crates/algeon-query-engine/src/chunk/tests.rs +++ b/crates/algeon-query-engine/src/chunk/tests.rs @@ -6,12 +6,13 @@ use std::{ }; use arrow::{ - array::{Int64Array, StringArray}, - datatypes::{DataType, Field, Schema}, + array::{Array, Int32Array, Int64Array, ListArray, StringArray}, + datatypes::{DataType, Field, Int32Type, Schema}, record_batch::RecordBatch, }; use parquet::{arrow::ArrowWriter, file::properties::WriterProperties}; use query_planning::{ + backend::BackendCapabilities, chunk::{ LOCAL_PARQUET_THROUGHPUT_MAX_ROW_GROUPS_PER_CHUNK, SourceChunkPlanningOptions, SourceChunkPlanningPolicy, SourceTaskMemoryModel, decimal::validate_supplied_reader_decimal_types_match_footer, @@ -27,8 +28,9 @@ use query_planning::{ }; use super::*; -use crate::source_chunk::metadata::{ - SourceMetadataCacheKey, local_parquet_footer_metadata_path, sidecar_content_identity, +use crate::{ + rapids_backend::rapids_backend, + source_chunk::metadata::{SourceMetadataCacheKey, local_parquet_footer_metadata_path, sidecar_content_identity}, }; fn batchable_chunk(index: usize, file_index: usize, bytes: u64) -> SourceChunk { @@ -51,6 +53,8 @@ fn batchable_chunk(index: usize, file_index: usize, bytes: u64) -> SourceChunk { estimated_chunk_bytes: Some(bytes), estimated_source_read_bytes: Some(bytes), estimated_decoded_bytes: None, + estimated_column_elements: Some(1), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(bytes), source_file_row_groups: None, source_file_row_group_bytes: Some(vec![bytes]), @@ -59,6 +63,7 @@ fn batchable_chunk(index: usize, file_index: usize, bytes: u64) -> SourceChunk { source_file_pure_filter_compressed_bytes: Some(vec![0]), source_file_pure_filter_uncompressed_bytes: Some(vec![0]), source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: Some(vec![1]), reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, @@ -181,6 +186,155 @@ fn source_chunks_bound_decoded_dictionary_strings_by_the_byte_target() { fs::remove_file(path).unwrap(); } +const VECTOR_ROW_GROUP_ROWS: u64 = 200_000; +const VECTOR_WIDTH: i32 = 2048; + +/// A 2048-wide float vector source like an embedding table: each 200,000-row +/// row group decodes to 409,600,000 list values, within the backend's +/// per-column element limit, and five row groups together are too. +fn vector_source(row_group_rows: Vec) -> LocalParquetSource { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new( + "fp", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), VECTOR_WIDTH), + true, + ), + ])); + let file = LocalParquetFile::new("/tmp/missing-vector-source.parquet") + .with_row_groups((0..row_group_rows.len()).collect()) + .with_row_group_bytes(vec![1024; row_group_rows.len()]) + .with_row_group_rows(row_group_rows); + LocalParquetSource::from_files(vec![file], schema, vec!["id".to_string(), "fp".to_string()]) +} + +fn assert_within_element_limit(chunks: &[SourceChunk]) { + let limit = rapids_backend().max_column_elements(); + for chunk in chunks { + assert_eq!(chunk.max_column_elements, limit); + assert!(chunk.estimated_column_elements.is_some_and(|elements| elements <= limit), "{chunk:?}"); + } +} + +#[test] +fn vector_source_chunks_and_batches_stay_within_the_backend_column_element_limit() { + let source = vector_source(vec![VECTOR_ROW_GROUP_ROWS; 15]); + let row_group_elements = VECTOR_ROW_GROUP_ROWS * VECTOR_WIDTH as u64; + + let widened = + plan_local_parquet_source_chunks(NodeId::from_index_for_test(3), &source, SourceChunkPlanningPolicy::default()) + .unwrap(); + assert_eq!(widened.chunks.iter().map(SourceChunk::row_group_count).collect::>(), vec![5, 5, 5]); + assert_within_element_limit(&widened.chunks); + assert!(widened.chunks.iter().all(|chunk| chunk.estimated_column_elements == Some(5 * row_group_elements))); + + let single = plan_local_parquet_source_chunks( + NodeId::from_index_for_test(3), + &source, + SourceChunkPlanningPolicy::new(None, 1), + ) + .unwrap(); + assert_eq!(single.chunk_count(), 15); + let batches = plan_source_chunk_batches(&single.chunks, None); + assert_eq!(batches.iter().map(SourceChunkBatch::chunk_count).collect::>(), vec![5, 5, 5]); + for batch in &batches { + let batch_elements = batch.chunks.iter().filter_map(|chunk| chunk.estimated_column_elements).sum::(); + assert!(batch_elements <= rapids_backend().max_column_elements(), "{batch:?}"); + } + let (left, right) = batches[0].clone().split_for_retry().expect("a five-chunk batch splits"); + assert_within_element_limit(&left.chunks); + assert_within_element_limit(&right.chunks); +} + +#[test] +fn vector_row_group_above_the_column_element_limit_fails_with_the_row_group_size_that_fits() { + let source = vector_source(vec![VECTOR_ROW_GROUP_ROWS, 2_897_819]); + + let error = + plan_local_parquet_source_chunks(NodeId::from_index_for_test(3), &source, SourceChunkPlanningPolicy::default()) + .expect_err("one row group of 2,897,819 vectors cannot decode into one backend column"); + + assert_eq!(error.code(), QueryErrorCode::SizeOverflow); + let fact = |key| error.wire_fact(key).map(|value| value.into_owned()); + assert_eq!(fact("column").as_deref(), Some("fp")); + assert_eq!(fact("row_group").as_deref(), Some("1")); + assert_eq!(fact("row_group_rows").as_deref(), Some("2897819")); + assert_eq!(fact("element_count"), Some((2_897_819 * VECTOR_WIDTH as u64).to_string())); + assert_eq!(fact("element_limit"), Some(rapids_backend().max_column_elements().to_string())); + let max_row_group_rows = rapids_backend().max_column_elements() / VECTOR_WIDTH as u64; + assert_eq!(fact("max_row_group_rows"), Some(max_row_group_rows.to_string())); + assert!(fact("remedy").is_some_and(|remedy| remedy.contains(&format!("at most {max_row_group_rows} rows")))); +} + +#[test] +fn variable_list_row_groups_without_footer_value_counts_are_read_alone() { + let schema = Arc::new(Schema::new(vec![Field::new( + "values", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )])); + let file = LocalParquetFile::new("/tmp/missing-list-source.parquet") + .with_row_groups(vec![0, 1, 2]) + .with_row_group_bytes(vec![1024; 3]) + .with_row_group_rows(vec![10; 3]); + let source = LocalParquetSource::from_files(vec![file], schema, vec!["values".to_string()]); + + let plan = + plan_local_parquet_source_chunks(NodeId::from_index_for_test(3), &source, SourceChunkPlanningPolicy::default()) + .unwrap(); + + assert_eq!(plan.chunks.iter().map(SourceChunk::row_group_count).collect::>(), vec![1, 1, 1]); + assert!(plan.chunks.iter().all(|chunk| chunk.estimated_column_elements.is_none())); + assert!(plan_source_chunk_batches(&plan.chunks, None).iter().all(SourceChunkBatch::is_singleton)); +} + +#[test] +fn variable_list_elements_come_from_footer_leaf_value_counts() { + let path = std::env::temp_dir().join(format!( + "query-engine-list-elements-{}-{}.parquet", + std::process::id(), + SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos() + )); + let values = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![]), + Some(vec![Some(4), Some(5), Some(6), Some(7), Some(8)]), + None, + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("values", values.data_type().clone(), true), + ])); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])), Arc::new(values)]) + .unwrap(); + let props = WriterProperties::builder().set_max_row_group_row_count(Some(2)).build(); + let mut writer = ArrowWriter::try_new(fs::File::create(&path).unwrap(), Arc::clone(&schema), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let source = LocalParquetSource::from_files( + vec![LocalParquetFile::new(path.clone())], + schema, + vec!["key".to_string(), "values".to_string()], + ); + + let plan = plan_local_parquet_source_chunks( + NodeId::from_index_for_test(3), + &source, + SourceChunkPlanningPolicy::new(None, 1), + ) + .unwrap(); + + // Each row group counts one leaf entry per value plus one per empty or + // null list: 3 + 1 and 5 + 1. + assert_eq!( + plan.chunks.iter().map(|chunk| chunk.estimated_column_elements).collect::>(), + vec![Some(4), Some(6)] + ); + assert_eq!(plan.chunks[0].source_file_row_group_column_elements, Some(vec![4, 6])); + fs::remove_file(path).unwrap(); +} + #[test] fn adaptive_source_targets_follow_immutable_grant_and_explicit_targets() { let model = SourceTaskMemoryModel::aggregate_pipeline(true); diff --git a/crates/algeon-query-engine/src/exec/bounded/aggregate/source_adjacent.rs b/crates/algeon-query-engine/src/exec/bounded/aggregate/source_adjacent.rs index 47049de14..a8d9418ac 100644 --- a/crates/algeon-query-engine/src/exec/bounded/aggregate/source_adjacent.rs +++ b/crates/algeon-query-engine/src/exec/bounded/aggregate/source_adjacent.rs @@ -1953,6 +1953,8 @@ mod tests { estimated_chunk_bytes: Some(bytes), estimated_source_read_bytes: Some(bytes), estimated_decoded_bytes: None, + estimated_column_elements: Some(1), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(bytes), source_file_row_groups: None, source_file_row_group_bytes: None, @@ -1961,6 +1963,7 @@ mod tests { source_file_pure_filter_compressed_bytes: None, source_file_pure_filter_uncompressed_bytes: None, source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: None, reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, diff --git a/crates/algeon-query-engine/src/exec/bounded/chunks/retry.rs b/crates/algeon-query-engine/src/exec/bounded/chunks/retry.rs index 894dbc097..cefbcd903 100644 --- a/crates/algeon-query-engine/src/exec/bounded/chunks/retry.rs +++ b/crates/algeon-query-engine/src/exec/bounded/chunks/retry.rs @@ -135,6 +135,11 @@ pub(in crate::exec::bounded) fn retry_child_chunk(parent: &SourceChunk, row_grou checked_row_group_sum(parent.source_file_row_group_decoded_bytes.as_deref(), &row_group_indexes).or_else( || proportional_bytes(parent.estimated_decoded_bytes, row_group_indexes.len(), parent_row_group_count), ); + // A subset never holds more elements than its parent, so the parent's + // bound stands in when the per-row-group vector is missing. + let estimated_column_elements = + checked_row_group_sum(parent.source_file_row_group_column_elements.as_deref(), &row_group_indexes) + .or(parent.estimated_column_elements); SourceChunk { source_node_id: parent.source_node_id, @@ -155,6 +160,8 @@ pub(in crate::exec::bounded) fn retry_child_chunk(parent: &SourceChunk, row_grou estimated_chunk_bytes, estimated_source_read_bytes, estimated_decoded_bytes, + estimated_column_elements, + max_column_elements: parent.max_column_elements, exact_source_read_preflight_bytes: None, source_file_row_groups: parent.source_file_row_groups.clone(), source_file_row_group_bytes: parent.source_file_row_group_bytes.clone(), @@ -163,6 +170,7 @@ pub(in crate::exec::bounded) fn retry_child_chunk(parent: &SourceChunk, row_grou source_file_pure_filter_compressed_bytes: parent.source_file_pure_filter_compressed_bytes.clone(), source_file_pure_filter_uncompressed_bytes: parent.source_file_pure_filter_uncompressed_bytes.clone(), source_file_row_group_decoded_bytes: parent.source_file_row_group_decoded_bytes.clone(), + source_file_row_group_column_elements: parent.source_file_row_group_column_elements.clone(), source_file_row_group_rows: parent.source_file_row_group_rows.clone(), reader_decimal_types: parent.reader_decimal_types.clone(), reader_decimal_types_footer_verified: parent.reader_decimal_types_footer_verified, @@ -235,6 +243,8 @@ mod tests { estimated_chunk_bytes: Some(400), estimated_source_read_bytes: Some(400), estimated_decoded_bytes: Some(800), + estimated_column_elements: Some(40), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(400), source_file_row_groups: Some(vec![0, 1, 2, 3]), source_file_row_group_bytes: Some(vec![100, 100, 100, 100]), @@ -243,6 +253,7 @@ mod tests { source_file_pure_filter_compressed_bytes: Some(vec![0, 0, 0, 0]), source_file_pure_filter_uncompressed_bytes: Some(vec![0, 0, 0, 0]), source_file_row_group_decoded_bytes: Some(vec![200, 200, 200, 200]), + source_file_row_group_column_elements: None, source_file_row_group_rows: Some(vec![10, 10, 10, 10]), reader_decimal_types: vec![ParquetReaderDecimalType::new( "price", diff --git a/crates/algeon-query-engine/src/exec/bounded/chunks/source_read.rs b/crates/algeon-query-engine/src/exec/bounded/chunks/source_read.rs index 68783e107..dee00fd48 100644 --- a/crates/algeon-query-engine/src/exec/bounded/chunks/source_read.rs +++ b/crates/algeon-query-engine/src/exec/bounded/chunks/source_read.rs @@ -420,6 +420,8 @@ mod tests { estimated_chunk_bytes: Some(bytes), estimated_source_read_bytes: Some(bytes), estimated_decoded_bytes: None, + estimated_column_elements: Some(1), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(bytes), source_file_row_groups: None, source_file_row_group_bytes: Some(vec![bytes]), @@ -428,6 +430,7 @@ mod tests { source_file_pure_filter_compressed_bytes: Some(vec![pure_filter_bytes]), source_file_pure_filter_uncompressed_bytes: Some(vec![pure_filter_bytes]), source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: Some(vec![1]), reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, diff --git a/crates/algeon-query-engine/src/exec/bounded/facts.rs b/crates/algeon-query-engine/src/exec/bounded/facts.rs index 358847050..ca179a941 100644 --- a/crates/algeon-query-engine/src/exec/bounded/facts.rs +++ b/crates/algeon-query-engine/src/exec/bounded/facts.rs @@ -147,6 +147,8 @@ mod tests { estimated_chunk_bytes: Some(1024), estimated_source_read_bytes: Some(1024), estimated_decoded_bytes: None, + estimated_column_elements: Some(1), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(1024), source_file_row_groups: None, source_file_row_group_bytes: Some(vec![1024]), @@ -155,6 +157,7 @@ mod tests { source_file_pure_filter_compressed_bytes: Some(vec![0]), source_file_pure_filter_uncompressed_bytes: Some(vec![0]), source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: Some(vec![1]), reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, diff --git a/crates/algeon-query-engine/src/exec/bounded/join/leaf/batching.rs b/crates/algeon-query-engine/src/exec/bounded/join/leaf/batching.rs index b96603244..80a6ef45f 100644 --- a/crates/algeon-query-engine/src/exec/bounded/join/leaf/batching.rs +++ b/crates/algeon-query-engine/src/exec/bounded/join/leaf/batching.rs @@ -231,6 +231,8 @@ mod tests { estimated_chunk_bytes: Some(bytes), estimated_source_read_bytes: Some(bytes), estimated_decoded_bytes: None, + estimated_column_elements: Some(1), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(bytes), source_file_row_groups: None, source_file_row_group_bytes: Some(vec![bytes]), @@ -239,6 +241,7 @@ mod tests { source_file_pure_filter_compressed_bytes: Some(vec![0]), source_file_pure_filter_uncompressed_bytes: Some(vec![0]), source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: Some(vec![1]), reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, diff --git a/crates/algeon-query-engine/src/exec/source/cudf_contract/tests.rs b/crates/algeon-query-engine/src/exec/source/cudf_contract/tests.rs index ad5fb22c1..7efe42336 100644 --- a/crates/algeon-query-engine/src/exec/source/cudf_contract/tests.rs +++ b/crates/algeon-query-engine/src/exec/source/cudf_contract/tests.rs @@ -509,6 +509,8 @@ fn remote_direct_chunk_batch_propagates_required_footer_fetch_failure_for_unveri estimated_chunk_bytes: Some(4096), estimated_source_read_bytes: Some(4096), estimated_decoded_bytes: None, + estimated_column_elements: Some(10), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(4096), source_file_row_groups: file.row_groups.clone(), source_file_row_group_bytes: file.row_group_bytes.clone(), @@ -517,6 +519,7 @@ fn remote_direct_chunk_batch_propagates_required_footer_fetch_failure_for_unveri source_file_pure_filter_compressed_bytes: file.pure_filter_row_group_compressed_bytes.clone(), source_file_pure_filter_uncompressed_bytes: file.pure_filter_row_group_uncompressed_bytes.clone(), source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: file.row_group_rows.clone(), reader_decimal_types: file.reader_decimal_types.clone(), reader_decimal_types_footer_verified: false, @@ -583,6 +586,8 @@ fn remote_direct_chunk_batch_reuses_planner_verified_decimal_reader_type() { estimated_chunk_bytes: Some(4096), estimated_source_read_bytes: Some(4096), estimated_decoded_bytes: None, + estimated_column_elements: Some(10), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(4096), source_file_row_groups: file.row_groups.clone(), source_file_row_group_bytes: file.row_group_bytes.clone(), @@ -591,6 +596,7 @@ fn remote_direct_chunk_batch_reuses_planner_verified_decimal_reader_type() { source_file_pure_filter_compressed_bytes: file.pure_filter_row_group_compressed_bytes.clone(), source_file_pure_filter_uncompressed_bytes: file.pure_filter_row_group_uncompressed_bytes.clone(), source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: file.row_group_rows.clone(), reader_decimal_types, reader_decimal_types_footer_verified: true, @@ -635,6 +641,8 @@ fn chunk_batch_rejects_chunk_identity_mismatch_before_footer_verification() { estimated_chunk_bytes: Some(384), estimated_source_read_bytes: Some(816), estimated_decoded_bytes: None, + estimated_column_elements: Some(10), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(384), source_file_row_groups: file.row_groups.clone(), source_file_row_group_bytes: file.row_group_bytes.clone(), @@ -643,6 +651,7 @@ fn chunk_batch_rejects_chunk_identity_mismatch_before_footer_verification() { source_file_pure_filter_compressed_bytes: file.pure_filter_row_group_compressed_bytes.clone(), source_file_pure_filter_uncompressed_bytes: file.pure_filter_row_group_uncompressed_bytes.clone(), source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: file.row_group_rows.clone(), reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, diff --git a/crates/algeon-query-engine/src/exec/source/lakehouse_cache/mod.rs b/crates/algeon-query-engine/src/exec/source/lakehouse_cache/mod.rs index 537f0e270..9d314fd5c 100644 --- a/crates/algeon-query-engine/src/exec/source/lakehouse_cache/mod.rs +++ b/crates/algeon-query-engine/src/exec/source/lakehouse_cache/mod.rs @@ -219,6 +219,8 @@ mod tests { estimated_chunk_bytes: Some(1024), estimated_source_read_bytes: Some(1024), estimated_decoded_bytes: None, + estimated_column_elements: Some(10), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: Some(1024), source_file_row_groups: None, source_file_row_group_bytes: Some(vec![1024]), @@ -227,6 +229,7 @@ mod tests { source_file_pure_filter_compressed_bytes: Some(vec![0]), source_file_pure_filter_uncompressed_bytes: Some(vec![0]), source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: Some(vec![10]), reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, diff --git a/crates/algeon-query-engine/src/exec/source/metrics.rs b/crates/algeon-query-engine/src/exec/source/metrics.rs index c962c090d..dc5c576f8 100644 --- a/crates/algeon-query-engine/src/exec/source/metrics.rs +++ b/crates/algeon-query-engine/src/exec/source/metrics.rs @@ -220,6 +220,8 @@ mod tests { estimated_chunk_bytes: Some(32), estimated_source_read_bytes: Some(64), estimated_decoded_bytes: None, + estimated_column_elements: Some(10), + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: None, source_file_row_groups: Some(vec![0]), source_file_row_group_bytes: None, @@ -228,6 +230,7 @@ mod tests { source_file_pure_filter_compressed_bytes: None, source_file_pure_filter_uncompressed_bytes: None, source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: Some(vec![10]), reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, diff --git a/crates/algeon-query-engine/src/rapids_backend.rs b/crates/algeon-query-engine/src/rapids_backend.rs index 58281e1c1..bb7fa61bd 100644 --- a/crates/algeon-query-engine/src/rapids_backend.rs +++ b/crates/algeon-query-engine/src/rapids_backend.rs @@ -6,6 +6,7 @@ use arrow::datatypes::DataType as ArrowDataType; use cudf::{ + core::Column, interop::arrow_host::{arrow_type_from_cudf_dtype, cudf_dtype_from_arrow_type, is_supported_arrow_type}, ops::binary::{BinaryOperator, infer_binary_output_dtype}, }; @@ -45,6 +46,10 @@ impl BackendCapabilities for RapidsBackend { is_supported_arrow_type(data_type) } + fn max_column_elements(&self) -> u64 { + Column::MAX_ELEMENTS as u64 + } + fn binary_output_type(&self, op: BinaryOp, left: &ArrowDataType, right: &ArrowDataType) -> Result { let left_dtype = cudf_dtype_from_arrow_type(left).map_err(|error| { cudf_unsupported_type_error( diff --git a/crates/algeon-query-engine/src/source_chunk/local_parquet_plan.rs b/crates/algeon-query-engine/src/source_chunk/local_parquet_plan.rs index c4c585858..efff5a1de 100644 --- a/crates/algeon-query-engine/src/source_chunk/local_parquet_plan.rs +++ b/crates/algeon-query-engine/src/source_chunk/local_parquet_plan.rs @@ -76,7 +76,7 @@ pub fn plan_local_parquet_source_chunks_with_metadata_cache( } /// The backend inputs chunk planning consumes: the cuDF reader's minimum -/// source-read budget and its Arrow type support. +/// source-read budget, its Arrow type support and its per-column element limit. fn rapids_source_read_sizing() -> SourceReadSizing<'static> { SourceReadSizing::new(rapids_backend(), cudf::DataFrame::minimum_parquet_source_read_budget_bytes() as u64) } @@ -133,6 +133,7 @@ pub(crate) fn plan_local_parquet_source_chunks_with_metadata_cache_and_context( policy, &reader_decimal_types, reader_decimal_types_footer_verified, + sizing, )); continue; }; @@ -149,7 +150,7 @@ pub(crate) fn plan_local_parquet_source_chunks_with_metadata_cache_and_context( reader_decimal_types_footer_verified, sizing, &mut chunks, - ); + )?; } chunks = apply_source_scan_limit_to_chunks(source, chunks, sizing)?; diff --git a/crates/algeon-query-planning/src/backend.rs b/crates/algeon-query-planning/src/backend.rs index 09e4bbbd8..568ee3df8 100644 --- a/crates/algeon-query-planning/src/backend.rs +++ b/crates/algeon-query-planning/src/backend.rs @@ -60,6 +60,14 @@ pub trait BackendCapabilities: fmt::Debug + Send + Sync { /// are enforced during plan validation. fn supports_type(&self, data_type: &ArrowDataType) -> bool; + /// Largest number of elements one backend column holds, counted at every + /// level: the row count, and the child element count of each nested list. + /// + /// A column that would exceed it in one source read or one intermediate + /// result cannot be represented. Source chunk planning sizes every read + /// under this count. + fn max_column_elements(&self) -> u64; + /// The Arrow output type of the backend's binary kernel for these operands. /// /// An `Err` means the kernel has no binding for the operand pair, and diff --git a/crates/algeon-query-planning/src/chunk/assembly.rs b/crates/algeon-query-planning/src/chunk/assembly.rs index 7a09aba13..e033d7014 100644 --- a/crates/algeon-query-planning/src/chunk/assembly.rs +++ b/crates/algeon-query-planning/src/chunk/assembly.rs @@ -5,7 +5,13 @@ use std::collections::BTreeSet; use arrow::datatypes::DataType; use super::{ - SourceChunk, SourceChunkPlanningPolicy, SourceReadSizing, decoded::ReadColumnWidths, row_group::RowGroupMetadata, + SourceChunk, SourceChunkPlanningPolicy, SourceReadSizing, + decoded::ReadColumnWidths, + elements::{ + ColumnElements, read_columns_bound_elements_per_row, row_group_column_elements_from_schema, + rows_within_element_limit, + }, + row_group::RowGroupMetadata, }; use crate::{ datatypes::arrow_fixed_width_byte_size, @@ -63,6 +69,19 @@ pub(crate) fn validate_selected_row_groups( Ok(()) } +/// Groups `row_groups` of one file into chunks bounded by the policy's row +/// group cap and source-byte target, and by the backend's per-column element +/// limit ([`super::elements`]). +/// +/// A row group whose read columns have no bounded elements per row (a +/// variable-width list column without footer value counts) gets a chunk of +/// its own, the same way batch planning keeps chunks without a byte estimate +/// apart. A row group without a row count whose columns all have bounded +/// elements per row adds nothing to the element sum, as unknown bytes add +/// nothing to the byte target; its chunk has no `estimated_column_elements` +/// and stays a singleton batch. A row group whose own largest column exceeds +/// the element limit fails with [`QueryErrorCode::SizeOverflow`], since +/// row-group planning cannot split it. #[doc(hidden)] #[expect(clippy::too_many_arguments, reason = "planning inputs are distinct proof and constraint values")] pub fn append_row_group_chunks( @@ -77,10 +96,17 @@ pub fn append_row_group_chunks( reader_decimal_types_footer_verified: bool, sizing: SourceReadSizing<'_>, chunks: &mut Vec, -) { +) -> Result<()> { let max_row_groups = policy.effective_max_row_groups_per_chunk(); + let max_column_elements = sizing.max_column_elements(); let row_group_count_basis = metadata.map_or(row_groups.len(), RowGroupMetadata::row_group_count).max(1); let read_column_widths = ReadColumnWidths::of(&source.schema, &source.read_columns); + let elements_per_row_bounded = read_columns_bound_elements_per_row(&source.schema, &source.read_columns); + let file_column_elements = metadata.and_then(|metadata| { + (0..metadata.row_group_count()) + .map(|row_group| row_group_column_elements(source, metadata, row_group).map(|largest| largest.elements)) + .collect::>>() + }); let mut current = PendingChunk::default(); for row_group in row_groups { @@ -107,7 +133,32 @@ pub fn append_row_group_chunks( .and_then(|metadata| metadata.decoded_bytes(*row_group)) .or_else(|| read_column_widths.decoded_bytes(row_group_rows?, source_read_bytes?)); let planning_bytes = decoded_bytes.or(source_read_bytes); - if current.should_flush_before(planning_bytes, max_row_groups, policy.max_source_chunk_bytes) { + let column_elements = metadata.and_then(|metadata| row_group_column_elements(source, metadata, *row_group)); + if let Some(largest) = column_elements + && largest.elements > max_column_elements + { + return Err(row_group_exceeds_column_element_limit_error( + source, + file, + file_index, + *row_group, + row_group_rows, + largest, + max_column_elements, + )); + } + let column_elements = match column_elements { + Some(largest) => RowGroupElements::Known(largest.elements), + None if elements_per_row_bounded => RowGroupElements::UnknownRows, + None => RowGroupElements::Unbounded, + }; + if current.should_flush_before( + planning_bytes, + column_elements, + max_row_groups, + policy.max_source_chunk_bytes, + max_column_elements, + ) { push_pending_chunk( source_node_id, source, @@ -115,6 +166,7 @@ pub fn append_row_group_chunks( file_index, row_groups.len(), metadata, + file_column_elements.as_deref(), policy, reader_decimal_types, reader_decimal_types_footer_verified, @@ -133,6 +185,7 @@ pub fn append_row_group_chunks( pure_filter_compressed_bytes, pure_filter_uncompressed_bytes, row_group_rows, + column_elements, ); } @@ -143,6 +196,7 @@ pub fn append_row_group_chunks( file_index, row_groups.len(), metadata, + file_column_elements.as_deref(), policy, reader_decimal_types, reader_decimal_types_footer_verified, @@ -150,6 +204,65 @@ pub fn append_row_group_chunks( &mut current, chunks, ); + Ok(()) +} + +/// Element count of one row group's largest read column, as far as planning +/// can bound it. +#[derive(Debug, Clone, Copy)] +enum RowGroupElements { + Known(u64), + /// Every read column has bounded elements per row, but the row count is + /// unknown. + UnknownRows, + /// A read column has no bound on its elements per row. + Unbounded, +} + +/// The largest read column of `row_group`: the footer's leaf value counts when +/// present, otherwise the schema's fixed widths times its row count. +fn row_group_column_elements( + source: &LocalParquetSource, + metadata: &RowGroupMetadata, + row_group: usize, +) -> Option { + metadata.column_elements(row_group).or_else(|| { + row_group_column_elements_from_schema(&source.schema, &source.read_columns, metadata.row_group_rows(row_group)?) + }) +} + +fn row_group_exceeds_column_element_limit_error( + source: &LocalParquetSource, + file: &LocalParquetFile, + file_index: usize, + row_group: usize, + row_group_rows: Option, + largest: ColumnElements, + max_column_elements: u64, +) -> QueryError { + let mut error = QueryError::new( + QueryErrorCode::SizeOverflow, + "query_engine.source_chunk.plan", + "one parquet row group decodes to more elements in one column than one backend column holds", + ) + .with_fact("reason", "row_group_column_elements_exceed_backend_limit") + .with_fact("path", file.location.diagnostic_location()) + .with_fact("stable_id", parquet_file_stable_id_fact(file)) + .with_fact("file_index", file_index) + .with_fact("row_group", row_group) + .with_fact("element_count", largest.elements) + .with_fact("element_limit", max_column_elements); + if let Some(column) = largest.read_column.and_then(|position| source.read_columns.get(position)) { + error = error.with_fact("column", column); + } + if let Some(rows) = row_group_rows { + let max_row_group_rows = rows_within_element_limit(rows, largest.elements, max_column_elements); + error = error + .with_fact("row_group_rows", rows) + .with_fact("max_row_group_rows", max_row_group_rows) + .with_fact("remedy", format!("rewrite the file with row groups of at most {max_row_group_rows} rows")); + } + error } #[expect(clippy::too_many_arguments, reason = "planning inputs are distinct proof and constraint values")] pub(crate) fn push_pending_chunk( @@ -159,6 +272,7 @@ pub(crate) fn push_pending_chunk( file_index: usize, planned_row_group_count: usize, metadata: Option<&RowGroupMetadata>, + file_column_elements: Option<&[u64]>, policy: SourceChunkPlanningPolicy, reader_decimal_types: &[ParquetReaderDecimalType], reader_decimal_types_footer_verified: bool, @@ -208,6 +322,8 @@ pub(crate) fn push_pending_chunk( estimated_chunk_bytes, estimated_source_read_bytes, estimated_decoded_bytes, + estimated_column_elements: pending.known_column_elements(), + max_column_elements: sizing.max_column_elements(), exact_source_read_preflight_bytes, source_file_row_groups: file.row_groups.clone(), source_file_row_group_bytes: metadata.and_then(RowGroupMetadata::row_group_bytes_vec), @@ -217,6 +333,7 @@ pub(crate) fn push_pending_chunk( source_file_pure_filter_uncompressed_bytes: metadata .and_then(RowGroupMetadata::pure_filter_uncompressed_bytes_vec), source_file_row_group_decoded_bytes: metadata.and_then(RowGroupMetadata::decoded_bytes_vec), + source_file_row_group_column_elements: file_column_elements.map(<[u64]>::to_vec), source_file_row_group_rows: metadata.and_then(RowGroupMetadata::row_group_rows_vec), reader_decimal_types: reader_decimal_types.to_vec(), reader_decimal_types_footer_verified, @@ -325,6 +442,7 @@ pub fn single_all_row_groups_chunk( policy: SourceChunkPlanningPolicy, reader_decimal_types: &[ParquetReaderDecimalType], reader_decimal_types_footer_verified: bool, + sizing: SourceReadSizing<'_>, ) -> SourceChunk { SourceChunk { source_node_id, @@ -345,6 +463,8 @@ pub fn single_all_row_groups_chunk( estimated_chunk_bytes: file.file_size_bytes, estimated_source_read_bytes: file.file_size_bytes, estimated_decoded_bytes: None, + estimated_column_elements: None, + max_column_elements: sizing.max_column_elements(), exact_source_read_preflight_bytes: None, source_file_row_groups: file.row_groups.clone(), source_file_row_group_bytes: None, @@ -353,6 +473,7 @@ pub fn single_all_row_groups_chunk( source_file_pure_filter_compressed_bytes: None, source_file_pure_filter_uncompressed_bytes: None, source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: None, reader_decimal_types: reader_decimal_types.to_vec(), reader_decimal_types_footer_verified, @@ -396,12 +517,14 @@ pub fn zero_column_count_identity_chunk( policy, reader_decimal_types, reader_decimal_types_footer_verified, + sizing, ); chunk.reads_all_row_groups = false; chunk.known_row_count = Some(0); chunk.estimated_chunk_bytes = Some(0); chunk.estimated_source_read_bytes = Some(0); chunk.estimated_decoded_bytes = Some(0); + chunk.estimated_column_elements = Some(0); chunk.exact_source_read_preflight_bytes = exact_source_read_preflight_bytes(source, chunk.known_row_count, true, sizing); chunk.size_estimate_based = false; @@ -431,6 +554,8 @@ pub(crate) struct PendingChunk { known_pure_filter_uncompressed_bytes: u64, known_decoded_bytes: u64, known_rows: u64, + known_column_elements: u64, + has_unbounded_column_elements: bool, has_unknown_bytes: bool, has_unknown_projected_compressed_bytes: bool, has_unknown_projected_uncompressed_bytes: bool, @@ -438,6 +563,7 @@ pub(crate) struct PendingChunk { has_unknown_pure_filter_uncompressed_bytes: bool, has_unknown_decoded_bytes: bool, has_unknown_rows: bool, + has_unknown_column_elements: bool, } impl PendingChunk { @@ -453,6 +579,7 @@ impl PendingChunk { pure_filter_compressed_bytes: Option, pure_filter_uncompressed_bytes: Option, row_group_rows: Option, + column_elements: RowGroupElements, ) { self.row_groups.push(row_group); if let Some(row_group_bytes) = row_group_bytes { @@ -497,13 +624,25 @@ impl PendingChunk { } else { self.has_unknown_rows = true; } + match column_elements { + RowGroupElements::Known(elements) => { + self.known_column_elements = self.known_column_elements.saturating_add(elements); + }, + RowGroupElements::UnknownRows => self.has_unknown_column_elements = true, + RowGroupElements::Unbounded => { + self.has_unknown_column_elements = true; + self.has_unbounded_column_elements = true; + }, + } } fn should_flush_before( &self, next_row_group_bytes: Option, + next_column_elements: RowGroupElements, max_row_groups: usize, max_source_chunk_bytes: Option, + max_column_elements: u64, ) -> bool { if self.row_groups.is_empty() { return false; @@ -511,6 +650,16 @@ impl PendingChunk { if self.row_groups.len() >= max_row_groups { return true; } + match next_column_elements { + _ if self.has_unbounded_column_elements => return true, + RowGroupElements::Unbounded => return true, + RowGroupElements::Known(elements) + if self.known_column_elements.saturating_add(elements) > max_column_elements => + { + return true; + }, + RowGroupElements::Known(_) | RowGroupElements::UnknownRows => {}, + } let Some(max_source_chunk_bytes) = max_source_chunk_bytes else { return false; @@ -550,6 +699,10 @@ impl PendingChunk { (!self.has_unknown_rows).then_some(self.known_rows) } + fn known_column_elements(&self) -> Option { + (!self.has_unknown_column_elements).then_some(self.known_column_elements) + } + fn clear(&mut self) { self.row_groups.clear(); self.known_bytes = 0; @@ -560,6 +713,8 @@ impl PendingChunk { self.known_pure_filter_uncompressed_bytes = 0; self.known_decoded_bytes = 0; self.known_rows = 0; + self.known_column_elements = 0; + self.has_unbounded_column_elements = false; self.has_unknown_bytes = false; self.has_unknown_projected_compressed_bytes = false; self.has_unknown_projected_uncompressed_bytes = false; @@ -567,5 +722,6 @@ impl PendingChunk { self.has_unknown_pure_filter_uncompressed_bytes = false; self.has_unknown_decoded_bytes = false; self.has_unknown_rows = false; + self.has_unknown_column_elements = false; } } diff --git a/crates/algeon-query-planning/src/chunk/elements.rs b/crates/algeon-query-planning/src/chunk/elements.rs new file mode 100644 index 000000000..609c09eb7 --- /dev/null +++ b/crates/algeon-query-planning/src/chunk/elements.rs @@ -0,0 +1,129 @@ +//! Decoded element counts of the read columns of Parquet row groups. +//! +//! A backend column holds at most +//! [`BackendCapabilities::max_column_elements`](crate::backend::BackendCapabilities::max_column_elements) +//! elements at each level: its rows, and the children of every nested list. +//! Chunk planning keeps the largest read column of each chunk, and of each +//! multi-chunk batch, within that limit, so one planned read decodes into one +//! backend table. +//! +//! With a footer, a read column counts the larger of the row group's rows and +//! the `num_values` of each Parquet leaf beneath it. A leaf's `num_values` +//! has one entry per value plus one per null or empty list, which bounds the +//! element count of every list level on that leaf's path. Without a footer, a +//! column counts rows times the product of its `FixedSizeList` widths; a +//! variable-width list then has no count. + +use arrow::datatypes::{DataType as ArrowDataType, Schema}; +use parquet::file::metadata::ParquetMetaData; + +use crate::{parquet_stats::root_column_leaf_indices, source::LocalParquetSource}; + +/// The read column with the most decoded elements in one row group. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ColumnElements { + /// Element count of that column at its largest level; never below the + /// row group's row count. + pub(crate) elements: u64, + /// Position in the source's `read_columns` of the column holding + /// `elements`, or `None` when no read column holds more elements than the + /// row group has rows. + pub(crate) read_column: Option, +} + +impl ColumnElements { + fn rows(rows: u64) -> Self { + Self { elements: rows, read_column: None } + } + + fn widen(self, read_column: usize, elements: u64) -> Self { + if elements > self.elements { Self { elements, read_column: Some(read_column) } } else { self } + } +} + +/// The largest read column of every physical row group of `metadata`, from +/// leaf `num_values`. `None` when a read column has no column chunk. +pub(crate) fn row_group_column_elements_from_footer( + metadata: &ParquetMetaData, + source: &LocalParquetSource, +) -> Option> { + if source.read_columns.len() != source.read_column_file_indices.len() { + return None; + } + let footer_schema = metadata.file_metadata().schema_descr(); + let read_column_leaves = source + .read_column_file_indices + .iter() + .map(|root_index| { + Some(root_column_leaf_indices(footer_schema, *root_index)).filter(|leaves| !leaves.is_empty()) + }) + .collect::>>()?; + metadata + .row_groups() + .iter() + .map(|row_group| { + let rows = u64::try_from(row_group.num_rows()).ok()?; + read_column_leaves.iter().enumerate().try_fold( + ColumnElements::rows(rows), + |largest, (read_column, leaves)| { + let elements = leaves.iter().try_fold(0u64, |elements, leaf| { + Some(elements.max(u64::try_from(row_group.columns().get(*leaf)?.num_values()).ok()?)) + })?; + Some(largest.widen(read_column, elements)) + }, + ) + }) + .collect() +} + +/// The largest read column of a row group of `rows` rows, from the Arrow +/// schema alone. `None` when a read column is a variable-width list or is +/// absent from `schema`, since neither bounds its child elements per row. +pub(crate) fn row_group_column_elements_from_schema( + schema: &Schema, + read_columns: &[String], + rows: u64, +) -> Option { + read_columns.iter().enumerate().try_fold(ColumnElements::rows(rows), |largest, (read_column, column)| { + let per_row = elements_per_row(schema.field_with_name(column).ok()?.data_type())?; + Some(largest.widen(read_column, rows.checked_mul(per_row)?)) + }) +} + +/// Whether every read column has a bounded element count per row, which +/// makes a row count enough to bound a row group's elements. +pub(crate) fn read_columns_bound_elements_per_row(schema: &Schema, read_columns: &[String]) -> bool { + read_columns + .iter() + .all(|column| schema.field_with_name(column).is_ok_and(|field| elements_per_row(field.data_type()).is_some())) +} + +/// Largest per-row element count over the levels of one column: the product +/// of nested `FixedSizeList` widths, at least one for the row itself. +fn elements_per_row(data_type: &ArrowDataType) -> Option { + match data_type { + ArrowDataType::FixedSizeList(field, width) => { + Some(u64::try_from(*width).ok()?.checked_mul(elements_per_row(field.data_type())?)?.max(1)) + }, + ArrowDataType::List(_) + | ArrowDataType::LargeList(_) + | ArrowDataType::ListView(_) + | ArrowDataType::LargeListView(_) + | ArrowDataType::Map(..) => None, + ArrowDataType::Struct(fields) => { + fields.iter().try_fold(1u64, |largest, field| Some(largest.max(elements_per_row(field.data_type())?))) + }, + _ => Some(1), + } +} + +/// Largest row count at which a row group whose largest column holds +/// `elements` across `rows` rows stays within `limit`, keeping its elements +/// per row. +pub(crate) fn rows_within_element_limit(rows: u64, elements: u64, limit: u64) -> u64 { + if elements == 0 { + return limit; + } + let rows_within = u128::from(rows) * u128::from(limit) / u128::from(elements); + u64::try_from(rows_within).unwrap_or(u64::MAX).min(limit) +} diff --git a/crates/algeon-query-planning/src/chunk/mod.rs b/crates/algeon-query-planning/src/chunk/mod.rs index 8b7460e1b..b2f3604fc 100644 --- a/crates/algeon-query-planning/src/chunk/mod.rs +++ b/crates/algeon-query-planning/src/chunk/mod.rs @@ -23,6 +23,7 @@ pub mod assembly; #[doc(hidden)] pub mod decimal; pub mod decoded; +pub(crate) mod elements; mod memory; mod model; mod planner; @@ -44,8 +45,9 @@ pub use policy::{ /// What the execution backend contributes to exact source-read sizing. /// /// Chunk planning consumes these; it never derives a backend memory model of -/// its own. The minimum read budget is the backend reader's floor, and type -/// support decides whether a column has an exact fixed-width bound at all. +/// its own. The minimum read budget is the backend reader's floor, type +/// support decides whether a column has an exact fixed-width bound at all, and +/// the per-column element limit caps the rows and nested elements of one read. #[derive(Debug, Clone, Copy)] pub struct SourceReadSizing<'a> { backend: &'a dyn BackendCapabilities, @@ -63,6 +65,12 @@ impl<'a> SourceReadSizing<'a> { self.minimum_read_budget_bytes } + /// The backend's per-column element limit + /// ([`BackendCapabilities::max_column_elements`]). + pub(crate) fn max_column_elements(&self) -> u64 { + self.backend.max_column_elements() + } + /// Whether the backend can represent this Arrow type. pub(crate) fn supports_type(&self, data_type: &ArrowDataType) -> bool { self.backend.supports_type(data_type) diff --git a/crates/algeon-query-planning/src/chunk/model.rs b/crates/algeon-query-planning/src/chunk/model.rs index 3aa363b02..5b26b7bb0 100644 --- a/crates/algeon-query-planning/src/chunk/model.rs +++ b/crates/algeon-query-planning/src/chunk/model.rs @@ -112,6 +112,14 @@ pub struct SourceChunk { /// ([`crate::chunk::decoded`]), when every row group has an estimate. /// Chunk and batch planning compare it with the source-byte target. pub estimated_decoded_bytes: Option, + /// Upper bound on the decoded elements of the chunk's largest read + /// column, counted at its largest nesting level: the sum, over its row + /// groups, of each row group's largest column. `None` when a row group has no count. + pub estimated_column_elements: Option, + /// The backend's per-column element limit the chunk was planned under. + /// A multi-chunk batch keeps its summed `estimated_column_elements` + /// within it. + pub max_column_elements: u64, /// Exact cudf source-read budget in bytes for fixed-width projected columns. pub exact_source_read_preflight_bytes: Option, /// Source-level selected row groups, before chunking. @@ -128,6 +136,8 @@ pub struct SourceChunk { pub source_file_pure_filter_uncompressed_bytes: Option>, /// Per-row-group decoded read-column estimates for the full file, if available. pub source_file_row_group_decoded_bytes: Option>, + /// Per-row-group largest read-column element counts for the full file, if available. + pub source_file_row_group_column_elements: Option>, /// Per-row-group row counts for the full file, if available. pub source_file_row_group_rows: Option>, /// Per-read-column decimal storage types proven from Parquet metadata. diff --git a/crates/algeon-query-planning/src/chunk/planner.rs b/crates/algeon-query-planning/src/chunk/planner.rs index 1e3fe9283..98e5fab16 100644 --- a/crates/algeon-query-planning/src/chunk/planner.rs +++ b/crates/algeon-query-planning/src/chunk/planner.rs @@ -40,6 +40,10 @@ fn source_read_bytes_for_batch_chunks(chunks: &[SourceChunk]) -> Option { /// source-read bytes when it has none; chunks without either estimate are kept /// as singleton batches so batching cannot defeat the bounded memory policy. /// A batch's `estimated_source_read_bytes` stays the sum of encoded bytes. +/// +/// A batch's summed `estimated_column_elements` stays within the chunks' +/// `max_column_elements`, so one batch read decodes into one backend table. +/// Chunks without an element count stay singleton batches. pub fn plan_source_chunk_batches( chunks: &[SourceChunk], max_batch_source_read_bytes: Option, @@ -63,6 +67,7 @@ pub fn plan_source_chunk_batches_with_row_cap( let mut current_bytes: Option = Some(0); let mut current_planning_bytes: Option = Some(0); let mut current_rows: Option = Some(0); + let mut current_column_elements: Option = Some(0); let mut current_reason: Option<&'static str> = None; for chunk in chunks { @@ -77,6 +82,7 @@ pub fn plan_source_chunk_batches_with_row_cap( max_batch_source_read_bytes, ) .or_else(|| source_chunk_batch_row_extension_rejection(current_rows, chunk, max_batch_rows)) + .or_else(|| source_chunk_batch_column_element_extension_rejection(current_column_elements, previous, chunk)) }); let can_extend = current.last().is_some() && extension_rejection.is_none(); if !current.is_empty() && !can_extend { @@ -87,6 +93,7 @@ pub fn plan_source_chunk_batches_with_row_cap( current_bytes = Some(0); current_planning_bytes = Some(0); current_rows = Some(0); + current_column_elements = Some(0); current_reason = None; } @@ -100,6 +107,8 @@ pub fn plan_source_chunk_batches_with_row_cap( .and_then(|total| chunk_planning_bytes.and_then(|bytes| total.checked_add(bytes))); } current_rows = current_rows.and_then(|total| chunk.known_row_count.and_then(|rows| total.checked_add(rows))); + current_column_elements = current_column_elements + .and_then(|total| chunk.estimated_column_elements.and_then(|elements| total.checked_add(elements))); current.push(chunk.clone()); } @@ -155,6 +164,9 @@ fn chunk_multi_source_batch_rejection(chunk: &SourceChunk) -> Option<&'static st if batchable_chunk_source_read_bytes(chunk).is_none() { return Some("ineligible_missing_source_read_bytes"); } + if chunk.estimated_column_elements.is_none() { + return Some("ineligible_missing_column_elements"); + } None } @@ -219,6 +231,18 @@ fn source_chunk_batch_row_extension_rejection( (current_rows.saturating_add(right_rows) > max_rows).then_some("row_target_exceeded") } +fn source_chunk_batch_column_element_extension_rejection( + current_column_elements: Option, + left: &SourceChunk, + right: &SourceChunk, +) -> Option<&'static str> { + let (Some(current), Some(right_elements)) = (current_column_elements, right.estimated_column_elements) else { + return Some("ineligible_missing_column_elements"); + }; + let max_column_elements = left.max_column_elements.min(right.max_column_elements); + (current.saturating_add(right_elements) > max_column_elements).then_some("column_element_limit_exceeded") +} + fn source_chunks_have_adjacent_row_groups(left: &SourceChunk, right: &SourceChunk) -> bool { let (Some(left_last), Some(right_first)) = (left.row_group_indexes.last(), right.row_group_indexes.first()) else { return false; diff --git a/crates/algeon-query-planning/src/chunk/row_group.rs b/crates/algeon-query-planning/src/chunk/row_group.rs index 10e491f70..614922ace 100644 --- a/crates/algeon-query-planning/src/chunk/row_group.rs +++ b/crates/algeon-query-planning/src/chunk/row_group.rs @@ -2,7 +2,10 @@ use std::collections::BTreeSet; use parquet::file::metadata::ParquetMetaData; -use super::decoded::source_row_group_decoded_bytes_from_footer; +use super::{ + decoded::source_row_group_decoded_bytes_from_footer, + elements::{ColumnElements, row_group_column_elements_from_footer}, +}; use crate::{ parquet_stats::{ColumnByteKind, row_group_column_bytes}, source::{LocalParquetFile, LocalParquetSource, SourceColumnRef}, @@ -19,6 +22,8 @@ pub struct RowGroupMetadata { pure_filter_uncompressed_bytes: Option>, decoded_bytes: Option>, row_group_rows: Option>, + /// Footer-only: source-supplied facts carry no leaf value counts. + column_elements: Option>, } impl RowGroupMetadata { @@ -62,6 +67,10 @@ impl RowGroupMetadata { self.row_group_rows.as_ref().and_then(|rows| rows.get(row_group).copied()) } + pub(crate) fn column_elements(&self, row_group: usize) -> Option { + self.column_elements.as_ref().and_then(|elements| elements.get(row_group).copied()) + } + pub(crate) fn row_group_bytes_vec(&self) -> Option> { self.row_group_bytes.clone() } @@ -145,6 +154,7 @@ pub fn row_group_metadata_from_footer(metadata: &ParquetMetaData, source: &Local .iter() .map(|row_group| u64::try_from(row_group.num_rows()).ok()) .collect::>>(), + column_elements: row_group_column_elements_from_footer(metadata, source), } } @@ -190,6 +200,7 @@ pub fn source_row_group_metadata( footer_metadata.as_ref().and_then(RowGroupMetadata::row_group_rows_vec), footer_row_group_count, ); + let column_elements = footer_metadata.and_then(|footer| footer.column_elements); if row_group_bytes.is_some() || projected_compressed_bytes.is_some() @@ -198,6 +209,7 @@ pub fn source_row_group_metadata( || pure_filter_uncompressed_bytes.is_some() || decoded_bytes.is_some() || row_group_rows.is_some() + || column_elements.is_some() { Some(RowGroupMetadata { row_group_bytes, @@ -207,6 +219,7 @@ pub fn source_row_group_metadata( pure_filter_uncompressed_bytes, decoded_bytes, row_group_rows, + column_elements, }) } else { None @@ -243,6 +256,7 @@ mod tests { pure_filter_uncompressed_bytes: Some(bytes.clone()), decoded_bytes: Some(bytes), row_group_rows: Some(rows), + column_elements: None, } } diff --git a/crates/algeon-query-planning/src/chunk/scan_limit.rs b/crates/algeon-query-planning/src/chunk/scan_limit.rs index 12fe5975d..1d4e8c8a6 100644 --- a/crates/algeon-query-planning/src/chunk/scan_limit.rs +++ b/crates/algeon-query-planning/src/chunk/scan_limit.rs @@ -48,6 +48,7 @@ pub fn apply_source_scan_limit_to_chunks( chunk.scan_limit_rows = Some(0); chunk.known_row_count = Some(0); chunk.estimated_decoded_bytes = Some(0); + chunk.estimated_column_elements = Some(0); chunk.exact_source_read_preflight_bytes = exact_source_read_preflight_bytes(source, chunk.known_row_count, chunk.zero_column_count, sizing); limited.push(chunk); @@ -181,6 +182,9 @@ pub(crate) fn rewrite_scan_limited_chunk_selection( .or_else(|| { proportional_bytes(chunk.estimated_decoded_bytes, selected_row_groups.len(), chunk.row_group_count()) }); + let estimated_column_elements = + sum_selected_row_group_values(chunk.source_file_row_group_column_elements.as_ref(), &selected_row_groups) + .or(chunk.estimated_column_elements); let estimated_source_read_bytes = source_read_uncompressed_bytes( source, known_projected_uncompressed_bytes, @@ -199,6 +203,7 @@ pub(crate) fn rewrite_scan_limited_chunk_selection( chunk.estimated_chunk_bytes = estimated_chunk_bytes; chunk.estimated_source_read_bytes = estimated_source_read_bytes; chunk.estimated_decoded_bytes = estimated_decoded_bytes; + chunk.estimated_column_elements = estimated_column_elements; chunk.exact_source_read_preflight_bytes = exact_source_read_preflight_bytes(source, Some(selected_row_group_rows), chunk.zero_column_count, sizing); chunk.known_row_count = Some(output_row_count); diff --git a/crates/algeon-query-runtime/src/native_error.rs b/crates/algeon-query-runtime/src/native_error.rs index c3e61d720..468c77605 100644 --- a/crates/algeon-query-runtime/src/native_error.rs +++ b/crates/algeon-query-runtime/src/native_error.rs @@ -12,6 +12,7 @@ use std::borrow::Cow; +use cudf::core::error::FactKeyName; use query_planning::error::{MemoryFailure, MemoryUsage, QueryError, QueryErrorCode, redact_uri_tokens}; /// Wraps a `cudf::Error` at the native-engine boundary. @@ -68,6 +69,7 @@ fn cudf_query_error_with_code( .with_fact("cudf_status", source.status()) .with_fact("cudf_operation", source.operation()); error = with_cudf_ffi_detail_facts(error, source); + error = with_cudf_column_element_limit_facts(error, source); if let Some(detail) = source.budget_exhaustion() { error = error.with_memory_failure(memory_failure(detail)); } @@ -135,6 +137,23 @@ pub fn cudf_message_fact(source: &cudf::Error) -> Option { Some(redact_uri_tokens(message)) } +/// Projects the column, element count, limit and row cap of a cuDF size +/// overflow, which the client message of `size_overflow` reads. +fn with_cudf_column_element_limit_facts(mut error: QueryError, source: &cudf::Error) -> QueryError { + if source.code() != cudf::ErrorCode::SizeOverflow { + return error; + } + if let Some(column) = source.fact_value(FactKeyName::ColumnName) { + error = error.with_fact("column", column.to_string()); + } + for key in ["element_count", "element_limit", "row_cap"] { + if let Some(value) = source.detail_value(key) { + error = error.with_fact(key, value); + } + } + error +} + fn with_cudf_ffi_detail_facts(mut error: QueryError, source: &cudf::Error) -> QueryError { let Some(detail) = source.ffi_detail() else { return error; diff --git a/crates/algeon-query-runtime/src/runtime/source_identity/local_parquet.rs b/crates/algeon-query-runtime/src/runtime/source_identity/local_parquet.rs index c63e85349..fce3ccc3c 100644 --- a/crates/algeon-query-runtime/src/runtime/source_identity/local_parquet.rs +++ b/crates/algeon-query-runtime/src/runtime/source_identity/local_parquet.rs @@ -250,6 +250,8 @@ mod tests { estimated_chunk_bytes: None, estimated_source_read_bytes: None, estimated_decoded_bytes: None, + estimated_column_elements: None, + max_column_elements: u64::MAX, exact_source_read_preflight_bytes: None, source_file_row_groups: None, source_file_row_group_bytes: None, @@ -258,6 +260,7 @@ mod tests { source_file_pure_filter_compressed_bytes: None, source_file_pure_filter_uncompressed_bytes: None, source_file_row_group_decoded_bytes: None, + source_file_row_group_column_elements: None, source_file_row_group_rows: None, reader_decimal_types: Vec::new(), reader_decimal_types_footer_verified: false, diff --git a/website/docs/cuml/ml-inputs.mdx b/website/docs/cuml/ml-inputs.mdx index 8c0dd0d2d..7a16e3b5d 100644 --- a/website/docs/cuml/ml-inputs.mdx +++ b/website/docs/cuml/ml-inputs.mdx @@ -138,6 +138,19 @@ ORDER BY score DESC; Class weights are not exposed. To trade precision for recall, threshold this score instead of `prediction`, which cuts at 0.5. +## Relation size + +cuDF indexes the values of a list column with 32-bit offsets. One relation +argument, or one intermediate result inside its subquery, can hold at most +2,147,483,647 list values per column: 1,048,575 rows of a 2048-wide +`FixedSizeList` feature. Scans and filters over larger Parquet files run. The +source planner reads each file in parts under the limit, sized from the +footer's value counts or, without a footer, from fixed list widths. A larger +relation fails with `size_overflow`, naming the column, the value count and +the row count that fits. Filter or sample the training relation below that +row count, since one call trains one model. A `predict` or `input` relation +can be split into row ranges, one call each, and combined with `UNION ALL`. + ## Validation boundary SQL planning and `gpu_validate_call` check names, static Arrow types, list diff --git a/website/docs/cuvs/vector-inputs.mdx b/website/docs/cuvs/vector-inputs.mdx index 98cbf10d3..335ee10be 100644 --- a/website/docs/cuvs/vector-inputs.mdx +++ b/website/docs/cuvs/vector-inputs.mdx @@ -61,6 +61,24 @@ with `list_vector_with_additional_features`. - Variable-width `List` and `LargeList` rows must have one uniform, positive dimension after evaluation. +## Relation size + +cuDF indexes the values of a list column with 32-bit offsets. One relation +argument can therefore hold at most 2,147,483,647 list values in its vector +column, which is 1,048,575 rows of a 2048-dimension `FixedSizeList`. The same +limit applies to each intermediate result inside the subquery. Separate scalar +feature columns hold up to 2,147,483,647 rows each. + +A subquery may scan and filter a larger Parquet file. The source planner reads +it in parts that each stay under the limit, sized from the footer's value +counts or, without a footer, from fixed list widths, and only the rows the subquery +returns count toward the relation. A relation above the limit fails the query +with `size_overflow`. The message names the column, its value count and the row +count that fits at its width. Filter the subquery to fewer rows, or split the +dataset into row ranges, run one call per range and combine the results with +`UNION ALL`. For brute-force kNN over a split dataset, keep the `k` smallest +distances per query across the parts. + ## Null values Schema nullability does not matter. A nullable field type is accepted for the diff --git a/website/docs/design/data-path.mdx b/website/docs/design/data-path.mdx index 1721d3586..da5d04fe4 100644 --- a/website/docs/design/data-path.mdx +++ b/website/docs/design/data-path.mdx @@ -39,7 +39,7 @@ to distinguish metadata pruning from row filtering. | Unit | Defined by | Size describes | | --- | --- | --- | | Parquet row group | The file writer | A set of rows stored as separate column chunks, with footer metadata | -| Source chunk | Algeon's source planner | Selected row groups from one file, grouped against a decoded-byte target and a row-group cap | +| Source chunk | Algeon's source planner | Selected row groups from one file, grouped against a decoded-byte target, a row-group cap and the backend's per-column element limit | | KvikIO I/O task | The I/O layer | A range of stored bytes to read, bounded by `KVIKIO_TASK_SIZE` | The planner can combine several row groups into one source chunk. Compatible @@ -49,6 +49,12 @@ pool. One source chunk can therefore require many I/O tasks. The I/O task size controls the read subdivision; the source-chunk target still comes from the query's memory policy. +Every planned read, and every batch of chunks read together, decodes into +one cuDF table. Its row count and the child value count of each list column +stay within the 2,147,483,647 elements one cuDF column indexes, counted from +the footer's per-leaf `num_values`. A row group that alone exceeds that count +fails planning with `size_overflow` and the row group size that fits. + Source sizing uses estimates of decoded columns because dictionary-encoded strings can expand substantially during decoding. The [native settings](../integration/configuration.mdx#native-settings) define diff --git a/website/docs/integration/handling-errors.mdx b/website/docs/integration/handling-errors.mdx index 9374e8446..c0698ee72 100644 --- a/website/docs/integration/handling-errors.mdx +++ b/website/docs/integration/handling-errors.mdx @@ -146,6 +146,17 @@ the sealed attempt report. | `budget_bytes` | The query grant. | | `budget_current_bytes` | Usage sampled after the failing operation released its working set, kept out of the message. | +### Column element limit {/* #column-element-limit */} + +One cuDF column holds at most 2,147,483,647 elements at each level: its rows +and the child values of each list. A `size_overflow` caused by that limit +names the column, its element count and the limit. When one relation or +intermediate result would exceed it, the message also gives the row count that +fits at the column's width and asks for a filter or a split into row ranges +combined with `UNION ALL`. When a single Parquet row group exceeds it, the +message gives the row group size to rewrite the file with. The column name +appears only in the client message, never in the server log. + ### Non-converged algorithms {/* #algorithm-did-not-converge */} An `algorithm_did_not_converge` message names the algorithm, the iterations diff --git a/website/docs/reference/data-type-support.md b/website/docs/reference/data-type-support.md index 578a2694a..1c5ffdd76 100644 --- a/website/docs/reference/data-type-support.md +++ b/website/docs/reference/data-type-support.md @@ -51,7 +51,14 @@ reads only the top-level validity of the column. Native lists use Int32 offsets. Every offset, child length and `rows * D` must fit the nonnegative Int32 range. A `LargeList` outside that range is -rejected. Each non-null `FixedSizeList` row holds exactly `D` child +rejected. The range bounds one relation or intermediate result, such as a +UDTF relation argument, a join build side or a sort input: at most +2,147,483,647 child values per list column, or 1,048,575 rows of a 2048-wide +`FixedSizeList`. A Parquet scan over a larger file reads it in parts that each +fit. A result above the limit fails with `size_overflow`, naming the column +and the row count that fits; filter it, or split it into row ranges combined +with `UNION ALL`. A single Parquet row group above the limit fails planning, +and the message gives the row group size to rewrite the file with. Each non-null `FixedSizeList` row holds exactly `D` child positions. NaN, infinities, signed zero, root and child nulls, and empty valid lists are payload values that carry paths preserve. From 20f90fd3dc91b75e6cedd2b456b0ffe763b304b8 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Wed, 21 Oct 2026 18:33:05 +0000 Subject: [PATCH 3/8] fix(datafusion): evaluate a GPU table function call once per query execution DataFusion inlines a CTE or view at every reference, so a KMeans, Louvain or UMAP call in a CTE read twice ran twice and the two references disagreed (170 of 200,000 rows in a KMeans self-join). Island reuse was limited to cuGraph similarity with repeatable inputs, and cuVS and cuML named each physical copy separately, so no copy could be matched. Every cuGraph, cuVS and cuML call site now gets one invocation identity at logical planning (GpuFunctionInvocation), kept through plan rebuilds and part of node equality, so separately written calls stay independent. Inside a composed island, every GPU function result is reused by invocation and the function's own output-plan fingerprint; the repeatability gate and the input-scope walk are removed. Host occurrences of one invocation share one manifest entry: the first reader polled runs the function and the others replay its batches, charged to the DataFusion memory pool. References split between an island and a host boundary, and a cuGraph copy that absorbed a LIMIT, still evaluate separately; the design page documents both. Co-Authored-By: Claude Opus 5.5 --- .../src/cugraph_sql/exec/mod.rs | 4 - .../src/cugraph_sql/exec/node.rs | 49 ++--- .../src/cuml_sql/exec/node.rs | 31 ++- .../src/cuml_sql/extension_planner.rs | 12 +- .../algeon-datafusion/src/cuml_sql/logical.rs | 33 ++- .../src/cuvs_sql/exec/node.rs | 26 +-- .../src/cuvs_sql/extension_planner.rs | 11 +- .../algeon-datafusion/src/cuvs_sql/logical.rs | 25 ++- .../src/gpu_functions/invocation.rs | 41 ++++ .../src/gpu_functions/mod.rs | 2 + .../src/native_exec/admission.rs | 21 +- .../src/native_exec/fragment.rs | 22 -- .../src/native_exec/occurrence.rs | 192 ++++++++++++++++-- .../src/native_exec/producer.rs | 28 +-- .../src/native_exec/relation_function.rs | 6 +- .../src/native_exec/shared_resident.rs | 51 ++--- .../src/native_exec/shared_resident/tests.rs | 94 ++------- .../algeon-datafusion/src/planner/native.rs | 138 +++++++++---- .../contract_tests_cpu/cuml_sql_tests.rs | 35 +++- .../contract_tests_cpu/cuvs_sql_tests/mod.rs | 38 +++- .../cugraph_sql_tests/gpu.rs | 15 +- .../root_tests_gpu/cuvs_sql_gpu_tests.rs | 85 +++++++- website/docs/cugraph/graph-options.mdx | 3 +- website/docs/design/gpu-table-functions.mdx | 39 +++- 24 files changed, 674 insertions(+), 327 deletions(-) create mode 100644 crates/algeon-datafusion/src/gpu_functions/invocation.rs diff --git a/crates/algeon-datafusion/src/cugraph_sql/exec/mod.rs b/crates/algeon-datafusion/src/cugraph_sql/exec/mod.rs index 43193d1f4..7dd39b39f 100644 --- a/crates/algeon-datafusion/src/cugraph_sql/exec/mod.rs +++ b/crates/algeon-datafusion/src/cugraph_sql/exec/mod.rs @@ -1,5 +1,3 @@ -use std::sync::atomic::AtomicU64; - use admission::{ cugraph_execution_config_after_staging, edge_source_graph_vertex_estimate, edge_source_output_rows, graph_input_native_producer, lowered_native_edge_plan, plan_contains_algeon_native, @@ -110,8 +108,6 @@ use crate::{ }, }; -static NEXT_CUGRAPH_GPU_RESULT_INPUT_ID: AtomicU64 = AtomicU64::new(0); - #[derive(Clone, Debug, Default)] struct CugraphResultPlanOptions { projection: Option>, diff --git a/crates/algeon-datafusion/src/cugraph_sql/exec/node.rs b/crates/algeon-datafusion/src/cugraph_sql/exec/node.rs index 58cae472d..ead8bb66e 100644 --- a/crates/algeon-datafusion/src/cugraph_sql/exec/node.rs +++ b/crates/algeon-datafusion/src/cugraph_sql/exec/node.rs @@ -1,6 +1,6 @@ use std::{ mem, - sync::{Arc, Mutex, atomic::Ordering}, + sync::{Arc, Mutex}, time::{Duration, Instant}, }; @@ -20,24 +20,24 @@ use super::{ ExecutionPlanMetricsSet, GpuDataFrame, GpuIslandHostBoundaryConfig, GpuResidentProducer, GpuResultHandle, GraphAlgorithmRuntimeMetrics, GraphBfsFollowerOutcome, GraphBfsFollowerRoute, GraphBfsLeader, GraphBfsPublishedResult, GraphBfsSingleflightJoin, GraphBfsSingleflightKey, GraphExecutionMetrics, - GraphExecutionPlan, GraphResidencyDescriptor, MetricsSet, NEXT_CUGRAPH_GPU_RESULT_INPUT_ID, - NativeExecutionMetricsCapture, NativeExpr, Partitioning, PhysicalExpr, PhysicalSortExpr, PlanNode, PlanProperties, - RecordBatch, RecordBatchStreamAdapter, Result, Schema, SchemaRef, SendableRecordBatchStream, SortOptions, - SourceNode, StagedCugraphInputs, TaskContext, TreeNodeRecursion, TryStreamExt, UploadedCugraphInputs, - attach_graph_residency, check_host_staging_cancellation, collect_cugraph_inputs, counted_cudf_runtime, - cugraph_edge_canonical_facts, cugraph_edge_normalize_inserted, cugraph_execution_capability_evidence, - cugraph_execution_config_after_staging, cugraph_graph_residency_descriptor, cugraph_schema_query_plan, - edge_source_graph_vertex_estimate, execute_cugraph_query_to_gpu_with_options, graph_uses_logical_mapping, - materialize_gpu_output_eagerly, native_graph_spec_for_placement_hint, plan, planned_resident_graph_key, - preflight_cugraph_query_plan, prepare_cugraph_source_metadata, record_batches_to_gpu, run_gpu_execution_task, - task_context_with_cugraph_execution_ticket, validate_cugraph_graph_owner_preflight, + GraphExecutionPlan, GraphResidencyDescriptor, MetricsSet, NativeExecutionMetricsCapture, NativeExpr, Partitioning, + PhysicalExpr, PhysicalSortExpr, PlanNode, PlanProperties, RecordBatch, RecordBatchStreamAdapter, Result, Schema, + SchemaRef, SendableRecordBatchStream, SortOptions, SourceNode, StagedCugraphInputs, TaskContext, TreeNodeRecursion, + TryStreamExt, UploadedCugraphInputs, attach_graph_residency, check_host_staging_cancellation, + collect_cugraph_inputs, counted_cudf_runtime, cugraph_edge_canonical_facts, cugraph_edge_normalize_inserted, + cugraph_execution_capability_evidence, cugraph_execution_config_after_staging, cugraph_graph_residency_descriptor, + cugraph_schema_query_plan, edge_source_graph_vertex_estimate, execute_cugraph_query_to_gpu_with_options, + graph_uses_logical_mapping, materialize_gpu_output_eagerly, native_graph_spec_for_placement_hint, plan, + planned_resident_graph_key, preflight_cugraph_query_plan, prepare_cugraph_source_metadata, record_batches_to_gpu, + run_gpu_execution_task, task_context_with_cugraph_execution_ticket, validate_cugraph_graph_owner_preflight, validate_standalone_bounded_cugraph_sources, }; use crate::{ attempt::InvocationAttempt, + gpu_functions::GpuFunctionInvocation, native_exec::{ GraphClaimCalibrationCapture, SealedAttemptMetricsTarget, admission::QueryAdmissionTicket, - gpu_function_occurrence_from_task_context, gpu_resident_producer, + gpu_function_occurrence_from_task_context, }, planner::native::Invocation, report::PlanningReport, @@ -93,8 +93,7 @@ impl CugraphAlgoExec { } pub(crate) fn next_gpu_result_input_name() -> String { - let id = NEXT_CUGRAPH_GPU_RESULT_INPUT_ID.fetch_add(1, Ordering::Relaxed); - format!("__algeon_datafusion_cugraph_result_{id}") + GpuFunctionInvocation::mint("cugraph").result_input_name() } #[cfg(test)] @@ -964,24 +963,8 @@ impl GpuResidentProducer for CugraphAlgoExec { &self.gpu_result_input_name } - fn gpu_result_is_repeatable(&self) -> bool { - self.spec.algorithm.is_similarity() - && self.predicate_sources.iter().chain([&self.edge_source]).all(|input| { - gpu_resident_producer(input.as_ref()).is_some_and(GpuResidentProducer::gpu_result_is_repeatable) - }) - } - - fn reusable_gpu_result_scope(&self) -> Result>> { - if !self.gpu_result_is_repeatable() { - return Ok(None); - } - let scopes = self - .gpu_result_capability_evidence()? - .execution_scopes() - .iter() - .map(|scope| scope.execution_scope_fingerprint()) - .collect::>(); - Ok(Some(scopes)) + fn reusable_gpu_result_scope(&self) -> Result> { + Ok(Some(cugraph_schema_query_plan(&self.plan_definition())?.execution_fingerprint())) } fn gpu_result_capability_evidence(&self) -> Result { diff --git a/crates/algeon-datafusion/src/cuml_sql/exec/node.rs b/crates/algeon-datafusion/src/cuml_sql/exec/node.rs index 738258900..39fbecf67 100644 --- a/crates/algeon-datafusion/src/cuml_sql/exec/node.rs +++ b/crates/algeon-datafusion/src/cuml_sql/exec/node.rs @@ -1,9 +1,6 @@ use std::{ collections::{HashMap, HashSet}, - sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }, + sync::Arc, }; use arrow::datatypes::SchemaRef; @@ -12,7 +9,7 @@ use query_planning::plan::{MlForestModelSnapshot, QueryPlan}; use super::{ super::{ - CumlOperation, CumlRelationBinding, CumlValidatedOptions, DenseVectorBinding, + CumlLogicalNode, CumlOperation, CumlRelationBinding, CumlValidatedOptions, DenseVectorBinding, error::{feature_disabled, invalid_call}, }, lowering::lower_cuml_query_plan, @@ -29,8 +26,6 @@ use crate::{ }, }; -static NEXT_CUML_GPU_RESULT_INPUT_ID: AtomicU64 = AtomicU64::new(0); - #[derive(Clone, Debug)] /// Validated operation settings retained by the specialized physical execution wrapper. pub struct CumlRelationFunctionSpec { @@ -45,26 +40,28 @@ pub struct CumlRelationFunctionSpec { pub type CumlAlgoExec = RelationFunctionExec; impl RelationFunctionExec { + /// Plans `node`'s call over its physical `inputs`; the input and result + /// names derive from the node's invocation. pub(crate) fn try_new( - operation: CumlOperation, - options: CumlValidatedOptions, - model: Option>, + node: &CumlLogicalNode, inputs: Vec>, input_orderings: Vec>, - output_schema: SchemaRef, native_execution_config: AlgeonNativeExecutionConfig, ) -> DataFusionResult { - let id = NEXT_CUML_GPU_RESULT_INPUT_ID.fetch_add(1, Ordering::Relaxed); - let input_names = - (0..inputs.len()).map(|index| format!("__algeon_datafusion_cuml_input_{id}_{index}")).collect(); + let invocation = node.invocation(); + let input_names = (0..inputs.len()).map(|index| invocation.relation_input_name(index)).collect(); Self::try_new_with_names( - CumlRelationFunctionSpec { operation, options, model }, + CumlRelationFunctionSpec { + operation: node.operation(), + options: node.validated_options().clone(), + model: node.model().cloned(), + }, inputs, - output_schema, + node.output_schema(), native_execution_config, input_names, input_orderings, - format!("__algeon_datafusion_cuml_result_{id}"), + invocation.result_input_name(), ) } diff --git a/crates/algeon-datafusion/src/cuml_sql/extension_planner.rs b/crates/algeon-datafusion/src/cuml_sql/extension_planner.rs index 8cf8a353e..fae8ddfd4 100644 --- a/crates/algeon-datafusion/src/cuml_sql/extension_planner.rs +++ b/crates/algeon-datafusion/src/cuml_sql/extension_planner.rs @@ -43,15 +43,7 @@ impl ExtensionPlanner for CumlExtensionPlanner { .map(LexOrdering::new) }) .collect::>>()?; - CumlAlgoExec::try_new( - node.operation(), - node.validated_options().clone(), - node.model().cloned(), - physical_inputs.to_vec(), - input_orderings, - node.output_schema(), - native_execution_config, - ) - .map(|exec| Some(Arc::new(exec) as Arc)) + CumlAlgoExec::try_new(node, physical_inputs.to_vec(), input_orderings, native_execution_config) + .map(|exec| Some(Arc::new(exec) as Arc)) } } diff --git a/crates/algeon-datafusion/src/cuml_sql/logical.rs b/crates/algeon-datafusion/src/cuml_sql/logical.rs index 876a643d2..a68b5699a 100644 --- a/crates/algeon-datafusion/src/cuml_sql/logical.rs +++ b/crates/algeon-datafusion/src/cuml_sql/logical.rs @@ -23,7 +23,7 @@ use super::{ use crate::{ error::Result, gpu_functions::{ - dense_relation::known_row_count, fmt_input_orderings, input_ordering_expressions, + GpuFunctionInvocation, dense_relation::known_row_count, fmt_input_orderings, input_ordering_expressions, input_orderings_with_expressions, }, }; @@ -41,6 +41,7 @@ pub struct CumlLogicalNode { options: CumlValidatedOptions, /// Model bound at planning; present only for `cuml_forest_predict`. model: Option>, + invocation: GpuFunctionInvocation, schema: DFSchemaRef, } @@ -59,11 +60,21 @@ impl CumlLogicalNode { inputs: Vec, options_json: &str, model: Option>, + ) -> DataFusionResult { + Self::try_new_for_invocation(operation, inputs, options_json, model, GpuFunctionInvocation::mint("cuml")) + } + + fn try_new_for_invocation( + operation: CumlOperation, + inputs: Vec, + options_json: &str, + model: Option>, + invocation: GpuFunctionInvocation, ) -> DataFusionResult { let (options, output) = Self::validate_call(operation, &inputs, options_json, model.as_deref())?; let schema = Arc::new(DFSchema::try_from(output.as_ref().clone())?); let input_orderings = vec![Vec::new(); inputs.len()]; - Ok(Self { operation, inputs, input_orderings, options, model, schema }) + Ok(Self { operation, inputs, input_orderings, options, model, invocation, schema }) } fn validate_call( @@ -318,6 +329,10 @@ impl CumlLogicalNode { pub fn output_schema(&self) -> SchemaRef { Arc::clone(self.schema.inner()) } + + pub(crate) fn invocation(&self) -> GpuFunctionInvocation { + self.invocation + } } #[cfg(feature = "server")] @@ -370,6 +385,7 @@ impl PartialEq for CumlLogicalNode { && self.input_orderings == other.input_orderings && self.options == other.options && self.model == other.model + && self.invocation == other.invocation } } impl Eq for CumlLogicalNode {} @@ -380,17 +396,26 @@ impl Hash for CumlLogicalNode { self.input_orderings.hash(state); self.options.hash(state); self.model.hash(state); + self.invocation.hash(state); } } impl PartialOrd for CumlLogicalNode { fn partial_cmp(&self, other: &Self) -> Option { - (self.operation.descriptor().sql_function_name, &self.inputs, &self.input_orderings, &self.options, &self.model) + ( + self.operation.descriptor().sql_function_name, + &self.inputs, + &self.input_orderings, + &self.options, + &self.model, + &self.invocation, + ) .partial_cmp(&( other.operation.descriptor().sql_function_name, &other.inputs, &other.input_orderings, &other.options, &other.model, + &other.invocation, )) } } @@ -437,7 +462,7 @@ impl UserDefinedLogicalNodeCore for CumlLogicalNode { }; let options_json = dense_relation_function(self.operation).options_json(&inputs, self.options.value_options().clone())?; - Self::try_new_with_model(self.operation, inputs, &options_json, self.model.clone())? + Self::try_new_for_invocation(self.operation, inputs, &options_json, self.model.clone(), self.invocation)? .with_input_orderings(input_orderings) .map_err(DataFusionError::from) } diff --git a/crates/algeon-datafusion/src/cuvs_sql/exec/node.rs b/crates/algeon-datafusion/src/cuvs_sql/exec/node.rs index 7d989a01a..b2700effe 100644 --- a/crates/algeon-datafusion/src/cuvs_sql/exec/node.rs +++ b/crates/algeon-datafusion/src/cuvs_sql/exec/node.rs @@ -1,9 +1,6 @@ use std::{ collections::{HashMap, HashSet}, - sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }, + sync::Arc, }; use arrow::datatypes::SchemaRef; @@ -12,7 +9,7 @@ use query_planning::plan::QueryPlan; use super::{ super::{ - CuvsOperation, CuvsValidatedOptions, CuvsVectorRelationBinding, DenseVectorBinding, + CuvsLogicalNode, CuvsOperation, CuvsValidatedOptions, CuvsVectorRelationBinding, DenseVectorBinding, error::{feature_disabled, invalid_call}, }, lowering::lower_cuvs_query_plan, @@ -29,8 +26,6 @@ use crate::{ }, }; -static NEXT_CUVS_GPU_RESULT_INPUT_ID: AtomicU64 = AtomicU64::new(0); - #[derive(Clone, Debug)] /// Validated operation settings retained by the specialized physical execution wrapper. pub struct CuvsRelationFunctionSpec { @@ -44,25 +39,24 @@ pub struct CuvsRelationFunctionSpec { pub type CuvsAlgoExec = RelationFunctionExec; impl RelationFunctionExec { + /// Plans `node`'s call over its physical `inputs`; the input and result + /// names derive from the node's invocation. pub(crate) fn try_new( - operation: CuvsOperation, - options: CuvsValidatedOptions, + node: &CuvsLogicalNode, inputs: Vec>, input_orderings: Vec>, - output_schema: SchemaRef, native_execution_config: AlgeonNativeExecutionConfig, ) -> DataFusionResult { - let id = NEXT_CUVS_GPU_RESULT_INPUT_ID.fetch_add(1, Ordering::Relaxed); - let input_names = - (0..inputs.len()).map(|index| format!("__algeon_datafusion_cuvs_input_{id}_{index}")).collect(); + let invocation = node.invocation(); + let input_names = (0..inputs.len()).map(|index| invocation.relation_input_name(index)).collect(); Self::try_new_with_names( - CuvsRelationFunctionSpec { operation, options }, + CuvsRelationFunctionSpec { operation: node.operation(), options: node.validated_options().clone() }, inputs, - output_schema, + node.output_schema(), native_execution_config, input_names, input_orderings, - format!("__algeon_datafusion_cuvs_result_{id}"), + invocation.result_input_name(), ) } diff --git a/crates/algeon-datafusion/src/cuvs_sql/extension_planner.rs b/crates/algeon-datafusion/src/cuvs_sql/extension_planner.rs index fd59a3cb6..b06e200b4 100644 --- a/crates/algeon-datafusion/src/cuvs_sql/extension_planner.rs +++ b/crates/algeon-datafusion/src/cuvs_sql/extension_planner.rs @@ -43,14 +43,7 @@ impl ExtensionPlanner for CuvsExtensionPlanner { .map(LexOrdering::new) }) .collect::>>()?; - CuvsAlgoExec::try_new( - node.operation(), - node.validated_options().clone(), - physical_inputs.to_vec(), - input_orderings, - node.output_schema(), - native_execution_config, - ) - .map(|exec| Some(Arc::new(exec) as Arc)) + CuvsAlgoExec::try_new(node, physical_inputs.to_vec(), input_orderings, native_execution_config) + .map(|exec| Some(Arc::new(exec) as Arc)) } } diff --git a/crates/algeon-datafusion/src/cuvs_sql/logical.rs b/crates/algeon-datafusion/src/cuvs_sql/logical.rs index 71cf41ae9..16e639e24 100644 --- a/crates/algeon-datafusion/src/cuvs_sql/logical.rs +++ b/crates/algeon-datafusion/src/cuvs_sql/logical.rs @@ -20,7 +20,7 @@ use super::{ use crate::{ error::{Error, Result}, gpu_functions::{ - dense_relation::known_row_count, fmt_input_orderings, input_ordering_expressions, + GpuFunctionInvocation, dense_relation::known_row_count, fmt_input_orderings, input_ordering_expressions, input_orderings_with_expressions, }, }; @@ -36,6 +36,7 @@ pub struct CuvsLogicalNode { /// empty when the input declares no order. input_orderings: Vec>, options: CuvsValidatedOptions, + invocation: GpuFunctionInvocation, schema: DFSchemaRef, } @@ -46,6 +47,15 @@ impl CuvsLogicalNode { /// missing columns, unsupported types, or incompatible cardinalities return typed /// Algeon errors through the DataFusion error boundary. pub fn try_new(operation: CuvsOperation, inputs: Vec, options_json: &str) -> DataFusionResult { + Self::try_new_for_invocation(operation, inputs, options_json, GpuFunctionInvocation::mint("cuvs")) + } + + fn try_new_for_invocation( + operation: CuvsOperation, + inputs: Vec, + options_json: &str, + invocation: GpuFunctionInvocation, + ) -> DataFusionResult { let descriptor = operation.descriptor(); if inputs.len() != descriptor.relation_roles.len() { return Err(invalid_call( @@ -65,7 +75,7 @@ impl CuvsLogicalNode { validate_known_cardinalities(operation, &inputs, &options)?; let schema = derive_output_schema(operation, &inputs, &options, &dimensions, element)?; let input_orderings = vec![Vec::new(); inputs.len()]; - Ok(Self { operation, inputs, input_orderings, options, schema }) + Ok(Self { operation, inputs, input_orderings, options, invocation, schema }) } /// Declares each relation argument's order, which fixes its `row_ordinal` @@ -100,6 +110,10 @@ impl CuvsLogicalNode { pub fn output_schema(&self) -> SchemaRef { Arc::clone(self.schema.inner()) } + + pub(crate) fn invocation(&self) -> GpuFunctionInvocation { + self.invocation + } } impl PartialEq for CuvsLogicalNode { @@ -108,6 +122,7 @@ impl PartialEq for CuvsLogicalNode { && self.inputs == other.inputs && self.input_orderings == other.input_orderings && self.options == other.options + && self.invocation == other.invocation } } @@ -119,16 +134,18 @@ impl Hash for CuvsLogicalNode { self.inputs.hash(state); self.input_orderings.hash(state); self.options.hash(state); + self.invocation.hash(state); } } impl PartialOrd for CuvsLogicalNode { fn partial_cmp(&self, other: &Self) -> Option { - (&self.operation, &self.inputs, &self.input_orderings, &self.options).partial_cmp(&( + (&self.operation, &self.inputs, &self.input_orderings, &self.options, &self.invocation).partial_cmp(&( &other.operation, &other.inputs, &other.input_orderings, &other.options, + &other.invocation, )) } } @@ -177,7 +194,7 @@ impl UserDefinedLogicalNodeCore for CuvsLogicalNode { }; let options_json = dense_relation_function(self.operation).options_json(&inputs, self.options.value_options().clone())?; - Self::try_new(self.operation, inputs, &options_json)? + Self::try_new_for_invocation(self.operation, inputs, &options_json, self.invocation)? .with_input_orderings(input_orderings) .map_err(DataFusionError::from) } diff --git a/crates/algeon-datafusion/src/gpu_functions/invocation.rs b/crates/algeon-datafusion/src/gpu_functions/invocation.rs new file mode 100644 index 000000000..3ac3224d4 --- /dev/null +++ b/crates/algeon-datafusion/src/gpu_functions/invocation.rs @@ -0,0 +1,41 @@ +//! Identity of one SQL call site of a GPU table function. + +use std::sync::atomic::{AtomicU64, Ordering}; + +static NEXT_GPU_FUNCTION_INVOCATION_ID: AtomicU64 = AtomicU64::new(0); + +/// One SQL invocation of a GPU table function, minted when the relation +/// planner builds the call's logical node. +/// +/// DataFusion inlines a CTE or view at each reference by copying its logical +/// plan, so every copy of one call site carries the same identity and every +/// physical occurrence planned from those copies reads one evaluation per +/// query execution. Two call sites with identical text receive different +/// identities and evaluate independently. The identity also enters the +/// logical node's equality, which keeps DataFusion from treating two call +/// sites as one plan. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(crate) struct GpuFunctionInvocation { + family: &'static str, + id: u64, +} + +impl GpuFunctionInvocation { + /// Mints an identity unique within this process for one call site of a + /// function in `family`. + pub(crate) fn mint(family: &'static str) -> Self { + Self { family, id: NEXT_GPU_FUNCTION_INVOCATION_ID.fetch_add(1, Ordering::Relaxed) } + } + + /// Name that binds the invocation's result in a consumer's GPU-resident + /// input map and identifies its host occurrences in the attempt manifest. + pub(crate) fn result_input_name(self) -> String { + format!("__algeon_datafusion_{}_result_{}", self.family, self.id) + } + + /// Name that binds the relation argument at `index` in the invocation's + /// own native plan. + pub(crate) fn relation_input_name(self, index: usize) -> String { + format!("__algeon_datafusion_{}_input_{}_{index}", self.family, self.id) + } +} diff --git a/crates/algeon-datafusion/src/gpu_functions/mod.rs b/crates/algeon-datafusion/src/gpu_functions/mod.rs index ed2c0e4d2..f2a881db9 100644 --- a/crates/algeon-datafusion/src/gpu_functions/mod.rs +++ b/crates/algeon-datafusion/src/gpu_functions/mod.rs @@ -11,6 +11,7 @@ pub(crate) mod dense_relation_errors; mod descriptor; mod family; mod input_ordering; +mod invocation; mod named_args; mod registration; pub(crate) mod scalar_arguments; @@ -24,6 +25,7 @@ pub(crate) use family::GpuFunctionFamily; pub(crate) use input_ordering::{ fmt_input_orderings, input_ordering_expressions, input_orderings_with_expressions, resolve_input_ordering, }; +pub(crate) use invocation::GpuFunctionInvocation; pub(crate) use named_args::{NamedArgumentDefinition, NamedArgumentKind, bind_named_arguments, unknown_name_message}; pub(crate) use registration::install_gpu_function_family; pub(crate) use table_functions::gpu_metadata_table_functions; diff --git a/crates/algeon-datafusion/src/native_exec/admission.rs b/crates/algeon-datafusion/src/native_exec/admission.rs index 5906b6d1c..2f10efd50 100644 --- a/crates/algeon-datafusion/src/native_exec/admission.rs +++ b/crates/algeon-datafusion/src/native_exec/admission.rs @@ -49,7 +49,7 @@ use tokio::sync::Notify; use crate::{ attempt::execute_in_internal_attempt, error::{Error, ErrorCode, Result}, - native_exec::metrics::NativeExecutionMetricsCapture, + native_exec::{metrics::NativeExecutionMetricsCapture, occurrence::SharedHostOccurrenceResults}, planner::native::Invocation, }; @@ -190,11 +190,13 @@ struct TicketAdmission { } /// The ticket of one DataFusion execution of a native admission root, with -/// the occurrence round of that execution. +/// the occurrence round and the shared host occurrence results of that +/// execution. A later execution of the same plan starts with neither. #[derive(Debug)] struct ActiveQueryAdmissionTicket { ticket: Arc, execution_round: AttemptExecutionRound, + host_occurrence_results: Arc, } pub(crate) fn query_admission_ticket_from_task_context( @@ -206,6 +208,15 @@ pub(crate) fn query_admission_ticket_from_task_context( .map(|extension| Arc::clone(&extension.ticket)) } +/// Host occurrence evaluations shared by the readers of one DataFusion +/// execution, when `task_context` executes under an admission root. +pub(crate) fn shared_host_occurrence_results(task_context: &TaskContext) -> Option> { + task_context + .session_config() + .get_extension::() + .map(|extension| Arc::clone(&extension.host_occurrence_results)) +} + /// The ticket that native work in `task_context` admits under: the active /// execution's ticket, else the ticket its invocation minted. /// @@ -251,7 +262,11 @@ pub(crate) fn task_context_with_query_admission_ticket( ticket.current_capability_evidence().and_then(|evidence| evidence.execution_manifest().cloned()), ); let mut session_config = task_context.session_config().clone(); - session_config.set_extension(Arc::new(ActiveQueryAdmissionTicket { ticket, execution_round })); + session_config.set_extension(Arc::new(ActiveQueryAdmissionTicket { + ticket, + execution_round, + host_occurrence_results: Arc::default(), + })); Arc::new(TaskContext::new( task_context.task_id(), task_context.session_id(), diff --git a/crates/algeon-datafusion/src/native_exec/fragment.rs b/crates/algeon-datafusion/src/native_exec/fragment.rs index 2e7391ff8..7ea0da399 100644 --- a/crates/algeon-datafusion/src/native_exec/fragment.rs +++ b/crates/algeon-datafusion/src/native_exec/fragment.rs @@ -22,8 +22,6 @@ use futures::future::BoxFuture; #[cfg(test)] use query_engine::exec::admit_query_plan; use query_engine::exec::{QueryExecutionRequest, execute_query_plan, output::QueryOutput}; -#[cfg(feature = "cugraph")] -use query_planning::plan::PlanNode; use query_planning::{ admission::{AdmittedPlan, AttemptCapabilityEvidence}, plan::{QueryPlan, ResultSink, row_bounds::query_plan_hard_row_upper_bound}, @@ -679,26 +677,6 @@ fn validate_gpu_inputs(plan: &QueryPlan, gpu_inputs: &[Arc], } impl GpuResidentProducer for AlgeonNativeExec { - #[cfg(feature = "cugraph")] - fn gpu_result_is_repeatable(&self) -> bool { - self.native_plan().nodes().iter().all(|node| match &node.kind { - PlanNode::Source(_) - | PlanNode::Filter(_) - | PlanNode::Projection(_) - | PlanNode::Aggregate(_) - | PlanNode::Window(_) - | PlanNode::Sort(_) - | PlanNode::Limit(_) - | PlanNode::Join(_) - | PlanNode::Union(_) - | PlanNode::EdgeNormalize(_) => true, - PlanNode::GraphAlgorithm(_) | PlanNode::VectorAlgorithm(_) | PlanNode::MlAlgorithm(_) => false, - }) && self.gpu_inputs.iter().all(|input| { - super::producer::gpu_resident_producer(input.as_ref()) - .is_some_and(GpuResidentProducer::gpu_result_is_repeatable) - }) - } - /// The name of a GPU-resident unit's output. Only GPU-resident units are /// classified as producers, so a host unit never answers. fn gpu_result_input_name(&self) -> &str { diff --git a/crates/algeon-datafusion/src/native_exec/occurrence.rs b/crates/algeon-datafusion/src/native_exec/occurrence.rs index 7cb92c58a..82ee45e5e 100644 --- a/crates/algeon-datafusion/src/native_exec/occurrence.rs +++ b/crates/algeon-datafusion/src/native_exec/occurrence.rs @@ -1,16 +1,34 @@ -//! Final-plan binding for one host-materializing GPU UDTF occurrence. +//! Final-plan binding for host-materializing GPU UDTF occurrences. -use std::sync::Arc; +use std::{ + collections::HashMap, + sync::{Arc, Mutex, PoisonError}, +}; +use arrow::record_batch::RecordBatch; use datafusion::{ - common::{DataFusionError, Result as DataFusionResult, tree_node::TreeNodeRecursion}, - execution::context::TaskContext, + common::{ + DataFusionError, Result as DataFusionResult, tree_node::TreeNodeRecursion, + utils::memory::get_record_batch_memory_size, + }, + execution::{ + context::TaskContext, + memory_pool::{MemoryConsumer, MemoryReservation}, + }, physical_expr::PhysicalExpr, - physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream}, + physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, + stream::RecordBatchStreamAdapter, + }, }; +use futures::{StreamExt, stream}; +use tokio::sync::OnceCell; +use super::admission::shared_host_occurrence_results; use crate::error::{Error, ErrorCode, Result}; +const SHARED_HOST_RESULT_MEMORY_CONSUMER: &str = "GpuFunctionHostOccurrenceExec[shared-host-result]"; + #[derive(Debug, Clone)] struct ActiveGpuFunctionOccurrence { occurrence_id: Arc, @@ -18,13 +36,25 @@ struct ActiveGpuFunctionOccurrence { } /// A transparent final-plan node that binds a shared physical producer to one -/// concrete invocation occurrence. Its task-context lineage lets runtime -/// dependency validation use the physical path selected by DataFusion. +/// manifest occurrence. Its task-context lineage lets runtime dependency +/// validation use the physical path selected by DataFusion. +/// +/// DataFusion copies a CTE or view at each reference, so one GPU function +/// invocation can reach several host boundaries. Binding gives every copy with +/// the same producer identity and host plan one occurrence id and records the +/// number of these readers. When there are several, the first reader to poll +/// drives one execution of the input, which begins and completes the single +/// manifest entry, and every reader replays the retained Arrow batches. A +/// reader that a consumer polls before another reader has drained still +/// drives the shared execution, so no reader waits on the order in which the +/// consumer polls its inputs. #[derive(Debug)] pub(crate) struct GpuFunctionHostOccurrenceExec { input: Arc, occurrence_id: Arc, producer_identity: Arc, + readers: usize, + account_memory_pool: bool, properties: Arc, } @@ -39,9 +69,21 @@ impl GpuFunctionHostOccurrenceExec { input, occurrence_id: occurrence_id.into(), producer_identity: producer_identity.into(), + readers: 1, + account_memory_pool: false, } } + /// Marks this node as one of `readers` final-plan nodes bound to the same + /// occurrence. With more than one reader, the retained batches are + /// charged to the DataFusion memory pool when `account_memory_pool` is + /// set, the policy native host output follows. + pub(crate) fn with_readers(mut self, readers: usize, account_memory_pool: bool) -> Self { + self.readers = readers; + self.account_memory_pool = account_memory_pool; + self + } + pub(crate) fn input(&self) -> &Arc { &self.input } @@ -53,6 +95,111 @@ impl GpuFunctionHostOccurrenceExec { pub(crate) fn producer_identity(&self) -> &str { &self.producer_identity } + + pub(crate) fn readers(&self) -> usize { + self.readers + } + + fn execute_shared( + &self, + task_context: Arc, + results: &SharedHostOccurrenceResults, + ) -> SendableRecordBatchStream { + let evaluation = results.reader_evaluation(&self.occurrence_id, self.readers); + let input = Arc::clone(&self.input); + let account_memory_pool = self.account_memory_pool; + // The reader that polls first runs the input. When it is dropped + // before the result is ready, the next reader to poll starts again. + let output = async move { + evaluation.get_or_init(|| evaluate_host_result(input, task_context, account_memory_pool)).await.clone() + }; + let batches = stream::once(output).flat_map(|result| match result { + Ok(retained) => { + stream::iter((0..retained.batches.len()).map(move |index| Ok(retained.batches[index].clone()))) + .left_stream() + }, + Err(error) => stream::iter([Err(DataFusionError::Shared(error))]).right_stream(), + }); + Box::pin(RecordBatchStreamAdapter::new(self.input.schema(), batches)) + } +} + +async fn evaluate_host_result( + input: Arc, + task_context: Arc, + account_memory_pool: bool, +) -> SharedHostOutput { + let reservation = account_memory_pool.then(|| { + MemoryConsumer::new(SHARED_HOST_RESULT_MEMORY_CONSUMER) + .with_can_spill(false) + .register(task_context.memory_pool()) + }); + let mut batches = Vec::new(); + let collected = async { + let mut input = input.execute(0, task_context)?; + while let Some(batch) = input.next().await { + let batch = batch?; + if let Some(reservation) = &reservation { + reservation.try_grow(get_record_batch_memory_size(&batch))?; + } + batches.push(batch); + } + Ok::<_, DataFusionError>(()) + } + .await; + match collected { + Ok(()) => Ok(Arc::new(RetainedHostResult { batches, _reservation: reservation })), + Err(error) => Err(Arc::new(error)), + } +} + +/// The Arrow result of one shared occurrence, retained until the last reader +/// stream releases it. +struct RetainedHostResult { + batches: Vec, + _reservation: Option, +} + +type SharedHostOutput = std::result::Result, Arc>; + +struct PendingReaders { + evaluation: Arc>, + remaining: usize, +} + +/// Shared host occurrence evaluations of one DataFusion execution of a native +/// admission root, by occurrence id. +/// +/// An entry leaves the map once its last reader has taken the evaluation, so +/// the retained batches live as long as a reader stream holds them. Entries of +/// readers that never execute drop with the execution. An entry holds no task +/// context, which keeps it from owning the execution that owns this map. +#[derive(Default)] +pub(crate) struct SharedHostOccurrenceResults { + pending: Mutex, PendingReaders>>, +} + +impl std::fmt::Debug for SharedHostOccurrenceResults { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner); + f.debug_struct("SharedHostOccurrenceResults").field("pending_occurrences", &pending.len()).finish() + } +} + +impl SharedHostOccurrenceResults { + /// The evaluation every one of the `readers` of `occurrence_id` awaits. + fn reader_evaluation(&self, occurrence_id: &Arc, readers: usize) -> Arc> { + let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner); + let entry = pending + .entry(Arc::clone(occurrence_id)) + .or_insert_with(|| PendingReaders { evaluation: Arc::default(), remaining: readers }); + let evaluation = Arc::clone(&entry.evaluation); + entry.remaining = entry.remaining.saturating_sub(1); + if entry.remaining == 0 { + pending.remove(occurrence_id); + } + evaluation + } } pub(crate) fn gpu_function_occurrence_from_task_context( @@ -96,7 +243,11 @@ fn task_context_with_gpu_function_occurrence( impl DisplayAs for GpuFunctionHostOccurrenceExec { fn fmt_as(&self, _: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "GpuFunctionHostOccurrenceExec: occurrence_id={}", self.occurrence_id) + write!(f, "GpuFunctionHostOccurrenceExec: occurrence_id={}", self.occurrence_id)?; + if self.readers > 1 { + write!(f, ", readers={}", self.readers)?; + } + Ok(()) } } @@ -127,17 +278,24 @@ impl ExecutionPlan for GpuFunctionHostOccurrenceExec { let [input] = children.try_into().map_err(|children: Vec>| { DataFusionError::Internal(format!("GPU function occurrence expected one child, got {}", children.len())) })?; - Ok(Arc::new(Self::new(input, Arc::clone(&self.occurrence_id), Arc::clone(&self.producer_identity)))) + Ok(Arc::new( + Self::new(input, Arc::clone(&self.occurrence_id), Arc::clone(&self.producer_identity)) + .with_readers(self.readers, self.account_memory_pool), + )) } fn execute(&self, partition: usize, task_context: Arc) -> DataFusionResult { - self.input.execute( - partition, - task_context_with_gpu_function_occurrence( - task_context, - Arc::clone(&self.occurrence_id), - Arc::clone(&self.producer_identity), - ), - ) + let shared_results = (self.readers > 1 && partition == 0) + .then(|| shared_host_occurrence_results(task_context.as_ref())) + .flatten(); + let task_context = task_context_with_gpu_function_occurrence( + task_context, + Arc::clone(&self.occurrence_id), + Arc::clone(&self.producer_identity), + ); + match shared_results { + Some(results) => Ok(self.execute_shared(task_context, &results)), + None => self.input.execute(partition, task_context), + } } } diff --git a/crates/algeon-datafusion/src/native_exec/producer.rs b/crates/algeon-datafusion/src/native_exec/producer.rs index 9b54e431f..149300cba 100644 --- a/crates/algeon-datafusion/src/native_exec/producer.rs +++ b/crates/algeon-datafusion/src/native_exec/producer.rs @@ -33,19 +33,21 @@ pub(crate) trait GpuResidentProducer: Send + Sync { /// matches the consumer plan's GPU-resident source name. fn gpu_result_input_name(&self) -> &str; - /// Whether this producer's inputs and operation permit reuse of one - /// logical invocation's result. Host-staged inputs and GPU functions - /// without a repeatability contract keep the default. - #[cfg(feature = "cugraph")] - fn gpu_result_is_repeatable(&self) -> bool { - false - } - - /// Native output identities eligible for reuse within one composed island. - /// The invocation name remains part of the key, keeping independently - /// authored calls separate. Output pushdown must change this identity; - /// the result scope also binds the complete input producer subtree. - fn reusable_gpu_result_scope(&self) -> Result>> { + /// Execution scope of the relation a GPU table function invocation + /// produces, or `None` for a producer that is not a GPU function call. + /// + /// Every physical occurrence of one SQL invocation shares + /// [`gpu_result_input_name`](Self::gpu_result_input_name), and occurrences + /// with equal names and scopes read one evaluation within a composed + /// island. The scope covers the function's validated options and any + /// output pushdown, which give two occurrences of one invocation + /// different relations. It excludes the input subtrees' plans: each + /// occurrence is planned from the same logical call, and a physical + /// rewrite of one copy's inputs preserves the relation they compute. + /// Facts a producer derives from its inputs, such as cuGraph's row and + /// vertex bounds, stay in the scope, so copies whose inputs report + /// different bounds evaluate separately. + fn reusable_gpu_result_scope(&self) -> Result> { Ok(None) } diff --git a/crates/algeon-datafusion/src/native_exec/relation_function.rs b/crates/algeon-datafusion/src/native_exec/relation_function.rs index 372333fa3..0b1a19fe5 100644 --- a/crates/algeon-datafusion/src/native_exec/relation_function.rs +++ b/crates/algeon-datafusion/src/native_exec/relation_function.rs @@ -25,7 +25,7 @@ use query_planning::{ admission::AttemptCapabilityEvidence, error::QueryError, expr::Expr as NativeExpr, - plan::{NodeId, ProjectionOutput, QueryPlan, QueryPlanBuilder}, + plan::{NodeId, PlanFingerprint, ProjectionOutput, QueryPlan, QueryPlanBuilder}, source::{GpuResidentSource, SourceNode}, }; use query_runtime::runtime::{ExecutionContext, resident::GpuResultHandle}; @@ -336,6 +336,10 @@ impl GpuResidentProducer for RelationFunctionExec { self.capability_evidence_for_sink(RelationFunctionSink::HostArrow) } + fn reusable_gpu_result_scope(&self) -> Result> { + Ok(Some(plan_for_sink(&self.native_host_plan, RelationFunctionSink::GpuResident).execution_fingerprint())) + } + fn native_relation_input_names(&self) -> Option<&[String]> { Some(&self.input_names) } diff --git a/crates/algeon-datafusion/src/native_exec/shared_resident.rs b/crates/algeon-datafusion/src/native_exec/shared_resident.rs index 53c3b58d2..d843d4589 100644 --- a/crates/algeon-datafusion/src/native_exec/shared_resident.rs +++ b/crates/algeon-datafusion/src/native_exec/shared_resident.rs @@ -9,10 +9,14 @@ //! unit takes the result; the device buffers then live as long as a reading //! unit holds its input map. //! -//! Repeatable GPU functions also register their reading occurrences here. -//! Their logical invocation name and native output scope distinguish separate -//! calls and pushed result options. Counts exclude the children of reused -//! parents because those children execute only for the first reader. +//! GPU table function occurrences register here as well. DataFusion copies a +//! CTE or view at each reference, so one SQL invocation can appear several +//! times in an island. Occurrences with one invocation name and one output +//! scope evaluate once, and every later occurrence reads that result, which +//! keeps the references of a function with nondeterministic results +//! consistent. Separate call sites have separate names and pushed result +//! options have separate scopes. Counts exclude the children of reused +//! occurrences because those children execute only for the first reader. use std::{ collections::BTreeMap, @@ -43,7 +47,7 @@ pub(crate) struct SharedResidentResults { #[derive(Clone, PartialEq, Eq, PartialOrd, Ord)] pub(crate) struct ReusableProducerIdentity { invocation: String, - scope: Vec, + scope: PlanFingerprint, } #[derive(PartialEq, Eq, PartialOrd, Ord)] @@ -157,43 +161,12 @@ fn count_consumers( } fn reusable_producer_identity(plan: &dyn ExecutionPlan) -> Result> { - if let Some(cooperative) = plan.downcast_ref::() { - return reusable_producer_identity(cooperative.input().as_ref()); - } let Some(producer) = gpu_resident_producer(plan) else { return Ok(None); }; - let Some(mut scope) = producer.reusable_gpu_result_scope()? else { - return Ok(None); - }; - for child in plan.children() { - if !append_input_execution_scopes(child.as_ref(), &mut scope)? { - return Ok(None); - } - } - Ok(Some(ReusableProducerIdentity { invocation: producer.gpu_result_input_name().to_string(), scope })) -} - -fn append_input_execution_scopes(plan: &dyn ExecutionPlan, scopes: &mut Vec) -> Result { - if let Some(cooperative) = plan.downcast_ref::() { - return append_input_execution_scopes(cooperative.input().as_ref(), scopes); - } - let Some(producer) = gpu_resident_producer(plan) else { - return Ok(false); - }; - scopes.extend( - producer - .gpu_result_capability_evidence()? - .execution_scopes() - .iter() - .map(|scope| scope.execution_scope_fingerprint()), - ); - for child in plan.children() { - if !append_input_execution_scopes(child.as_ref(), scopes)? { - return Ok(false); - } - } - Ok(true) + Ok(producer + .reusable_gpu_result_scope()? + .map(|scope| ReusableProducerIdentity { invocation: producer.gpu_result_input_name().to_string(), scope })) } /// `task_context` carrying a fresh result scope for one composed island. diff --git a/crates/algeon-datafusion/src/native_exec/shared_resident/tests.rs b/crates/algeon-datafusion/src/native_exec/shared_resident/tests.rs index 32e3b1ae1..1ccf24296 100644 --- a/crates/algeon-datafusion/src/native_exec/shared_resident/tests.rs +++ b/crates/algeon-datafusion/src/native_exec/shared_resident/tests.rs @@ -5,9 +5,8 @@ use datafusion::physical_plan::{ ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, coop::CooperativeExec, empty::EmptyExec, }; use query_planning::{ - expr::Expr, - plan::{ProjectionOutput, QueryPlanBuilder}, - source::{GpuResidentSource, LocalParquetFile, LocalParquetSource, SourceNode}, + plan::QueryPlanBuilder, + source::{LocalParquetFile, LocalParquetSource, SourceNode}, }; use super::SharedResidentResults; @@ -16,7 +15,7 @@ use crate::{ CugraphAlgoExec, CugraphAlgorithm, CugraphAlgorithmParams, CugraphAlgorithmSpec, CugraphGraphOptions, CugraphLeidenParams, CugraphSimilarityParams, GpuIslandHostBoundaryConfig, result_schema_for_params, }, - native_exec::{AlgeonNativeExec, AlgeonNativeExecutionConfig, NativeOutput, admit_for_test, gpu_resident_producer}, + native_exec::{AlgeonNativeExec, AlgeonNativeExecutionConfig, NativeOutput, admit_for_test}, }; fn native_input(columns: &[&str], name: &str) -> Arc { @@ -106,86 +105,20 @@ fn repeated_invocations_preserve_output_pushdown_and_independent_calls() { } #[test] -fn replaced_relation_inputs_do_not_share_a_result() { +fn rewritten_relation_inputs_read_one_result_of_their_invocation() { let original = algorithm("similar", CugraphAlgorithmParams::Similarity(CugraphSimilarityParams::default()), None, false); - let changed = Arc::clone(&original) + let rewritten = Arc::clone(&original) .replace_children( - vec![native_input(&["src", "dst"], "changed_edges"), native_input(&["first", "second"], "pairs")], + vec![native_input(&["src", "dst"], "rewritten_edges"), native_input(&["first", "second"], "pairs")], ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), ) .unwrap(); - let results = SharedResidentResults::for_gpu_inputs(&[ - Arc::clone(&original), - Arc::clone(&original), - Arc::clone(&changed), - Arc::clone(&changed), - ]) - .unwrap(); - let original = results.reusable_identity(original.as_ref()).unwrap().unwrap(); - let changed = results.reusable_identity(changed.as_ref()).unwrap().unwrap(); - assert!(original != changed, "same invocation and schema with different input contents must remain separate"); -} - -fn similarity_edges(input: Arc) -> Arc { - let input_name = gpu_resident_producer(input.as_ref()).unwrap().gpu_result_input_name().to_string(); - let mut builder = QueryPlanBuilder::new(); - let source = builder.add_source(SourceNode::GpuResident(GpuResidentSource::new(input_name, input.schema()))); - let projection = builder - .add_projection( - source, - vec![ - ProjectionOutput::new(Expr::column("first"), Field::new("src", DataType::Int64, false)), - ProjectionOutput::new(Expr::column("second"), Field::new("dst", DataType::Int64, false)), - ], - ) - .unwrap(); - let plan = builder.finish_with_host_sink(projection).unwrap().clone_with_gpu_resident_sink(); - let config = AlgeonNativeExecutionConfig::default(); - Arc::new( - AlgeonNativeExec::try_new( - admit_for_test(&plan, &config).unwrap(), - vec![input], - NativeOutput::GpuResident { input_name: "nested_edges".to_string() }, - config, - None, - ) - .unwrap(), - ) -} - -#[test] -fn nested_function_input_contents_are_part_of_reuse_identity() { - let inner = algorithm("inner", CugraphAlgorithmParams::Similarity(CugraphSimilarityParams::default()), None, false); - let changed_inner = Arc::clone(&inner) - .replace_children( - vec![native_input(&["src", "dst"], "changed_edges"), native_input(&["first", "second"], "pairs")], - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - .unwrap(); - let outer = algorithm("outer", CugraphAlgorithmParams::Similarity(CugraphSimilarityParams::default()), None, false); - let original = Arc::clone(&outer) - .replace_children( - vec![similarity_edges(inner), native_input(&["first", "second"], "pairs")], - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - .unwrap(); - let changed = outer - .replace_children( - vec![similarity_edges(changed_inner), native_input(&["first", "second"], "pairs")], - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - .unwrap(); - let results = SharedResidentResults::for_gpu_inputs(&[ - Arc::clone(&original), - Arc::clone(&original), - Arc::clone(&changed), - Arc::clone(&changed), - ]) - .unwrap(); + let results = SharedResidentResults::for_gpu_inputs(&[Arc::clone(&original), Arc::clone(&rewritten)]).unwrap(); let original = results.reusable_identity(original.as_ref()).unwrap().unwrap(); - let changed = results.reusable_identity(changed.as_ref()).unwrap().unwrap(); - assert!(original != changed, "a nested function's changed input must distinguish the outer result"); + let rewritten = results.reusable_identity(rewritten.as_ref()).unwrap().unwrap(); + assert!(original == rewritten, "a physical rewrite of one copy's inputs keeps the invocation's one result"); + assert_eq!(results.reusable_consumers[&original], 2); } #[test] @@ -207,7 +140,7 @@ fn reused_functions_register_only_executed_reads_of_shared_inputs() { } #[test] -fn nondeterministic_functions_and_host_staging_exclude_reuse() { +fn nondeterministic_functions_and_host_staged_inputs_reuse_one_result() { let leiden = algorithm("worlds", CugraphAlgorithmParams::Leiden(CugraphLeidenParams::default()), None, false); let host = algorithm("host", CugraphAlgorithmParams::Similarity(CugraphSimilarityParams::default()), None, true); let results = SharedResidentResults::for_gpu_inputs(&[ @@ -218,8 +151,7 @@ fn nondeterministic_functions_and_host_staging_exclude_reuse() { ]) .unwrap(); for plan in [leiden, host] { - let producer = gpu_resident_producer(plan.as_ref()).unwrap(); - assert!(!producer.gpu_result_is_repeatable()); - assert!(results.reusable_identity(plan.as_ref()).unwrap().is_none()); + let identity = results.reusable_identity(plan.as_ref()).unwrap().expect("every GPU function call is reusable"); + assert_eq!(results.reusable_consumers[&identity], 2); } } diff --git a/crates/algeon-datafusion/src/planner/native.rs b/crates/algeon-datafusion/src/planner/native.rs index d97f525de..5765f373e 100644 --- a/crates/algeon-datafusion/src/planner/native.rs +++ b/crates/algeon-datafusion/src/planner/native.rs @@ -15,12 +15,16 @@ //! The selection contract lives in `website/docs/design/native-lowering.mdx` //! ("Select an execution path" and "Require GPU execution"). -use std::sync::{Arc, Mutex, PoisonError}; +use std::{ + collections::HashMap, + sync::{Arc, Mutex, PoisonError}, +}; use datafusion::{ common::{ DataFusionError, Result as DataFusionResult, config::{ConfigOptions, OptimizerOptions}, + tree_node::{Transformed, TreeNode}, }, execution::SessionState, logical_expr::LogicalPlan, @@ -31,7 +35,7 @@ use query_planning::{ admission::{AttemptCapabilityEvidence, AttemptExecutionManifest, AttemptExecutionOccurrence}, correlation::QueryId, optimizer::NativeOptimizerDiagnostics, - plan::QueryPlan, + plan::{PlanFingerprint, QueryPlan}, source::SourceDiagnostics, }; @@ -455,8 +459,10 @@ impl AlgeonNativePhysicalOptimizerRule { // disposition before the ticket is minted so the admission that // follows has a planning predecessor in the lifecycle. if binds_final_plan { - let mut next_occurrence = 0_u64; - optimized.plan = bind_gpu_function_host_occurrences(optimized.plan, true, &mut next_occurrence)?; + optimized.plan = bind_final_gpu_function_host_occurrences( + optimized.plan, + native_execution_config.host_output_memory_pool_accounting(), + )?; } if binds_final_plan && plan_contains_gpu_execution(optimized.plan.as_ref()) @@ -589,27 +595,71 @@ fn gpu_function_execution_manifest(plan: &dyn ExecutionPlan) -> Result#` in pre-order, so the same function reached twice -/// (DataFusion inlines CTE references) gets two occurrences. A function read on -/// the device by a native unit or by another producer is a GPU-resident input -/// and is left unwrapped. Execution later validates each host occurrence -/// against the manifest built from these ids. +/// Ids are `#` in pre-order. DataFusion copies a CTE or view at +/// each reference, so one invocation can reach several host boundaries; the +/// copies with one producer identity and one host plan share an id and read +/// one evaluation. A pushed result option changes the host plan and keeps its +/// own id. A function read on the device by a native unit or by another +/// producer is a GPU-resident input and is left unwrapped. Execution later +/// validates each host occurrence against the manifest built from these ids. +fn bind_final_gpu_function_host_occurrences( + plan: Arc, + account_memory_pool: bool, +) -> DataFusionResult> { + let mut binding = HostOccurrenceBinding::default(); + let plan = bind_gpu_function_host_occurrences(plan, true, &mut binding)?; + if binding.readers.values().all(|readers| *readers == 1) { + return Ok(plan); + } + plan.transform_up(|node| { + let Some(bound) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let readers = binding.readers[bound.occurrence_id()]; + if readers == bound.readers() { + return Ok(Transformed::no(node)); + } + Ok(Transformed::yes(Arc::new( + GpuFunctionHostOccurrenceExec::new( + Arc::clone(bound.input()), + bound.occurrence_id(), + bound.producer_identity(), + ) + .with_readers(readers, account_memory_pool), + ) as Arc)) + }) + .map(|transformed| transformed.data) +} + +/// Occurrence ids assigned while binding one final plan. +#[derive(Default)] +struct HostOccurrenceBinding { + next_occurrence: u64, + /// Occurrence id of each producer identity and host plan fingerprint. + occurrence_ids: HashMap<(String, PlanFingerprint), String>, + /// Number of bound nodes that read each occurrence id. + readers: HashMap, +} + fn bind_gpu_function_host_occurrences( plan: Arc, materializes_host_result: bool, - next_occurrence: &mut u64, + binding: &mut HostOccurrenceBinding, ) -> DataFusionResult> { if plan.downcast_ref::().is_some() { return Ok(plan); @@ -619,7 +669,7 @@ fn bind_gpu_function_host_occurrences( let mut rewritten_children = Vec::with_capacity(children.len()); let mut children_changed = false; for child in children { - let rewritten = bind_gpu_function_host_occurrences(Arc::clone(&child), false, next_occurrence)?; + let rewritten = bind_gpu_function_host_occurrences(Arc::clone(&child), false, binding)?; children_changed |= !Arc::ptr_eq(&child, &rewritten); rewritten_children.push(rewritten); } @@ -636,7 +686,7 @@ fn bind_gpu_function_host_occurrences( for child in children { let child_materializes_host_result = producer.is_none() || gpu_resident_producer(child.as_ref()).is_none(); let rewritten = - bind_gpu_function_host_occurrences(Arc::clone(&child), child_materializes_host_result, next_occurrence)?; + bind_gpu_function_host_occurrences(Arc::clone(&child), child_materializes_host_result, binding)?; children_changed |= !Arc::ptr_eq(&child, &rewritten); rewritten_children.push(rewritten); } @@ -648,12 +698,23 @@ fn bind_gpu_function_host_occurrences( let Some(producer) = gpu_resident_producer(plan.as_ref()) else { return Ok(plan); }; - if !materializes_host_result || producer.host_result_query_plan()?.is_none() { + if !materializes_host_result { return Ok(plan); } + let Some(host_plan) = producer.host_result_query_plan()? else { + return Ok(plan); + }; let producer_identity = producer.gpu_result_input_name().to_string(); - let occurrence_id = format!("{producer_identity}#{}", *next_occurrence); - *next_occurrence = next_occurrence.saturating_add(1); + let HostOccurrenceBinding { next_occurrence, occurrence_ids, readers } = binding; + let occurrence_id = occurrence_ids + .entry((producer_identity.clone(), host_plan.execution_fingerprint())) + .or_insert_with(|| { + let occurrence_id = format!("{producer_identity}#{next_occurrence}"); + *next_occurrence = next_occurrence.saturating_add(1); + occurrence_id + }) + .clone(); + *readers.entry(occurrence_id.clone()).or_default() += 1; Ok(Arc::new(GpuFunctionHostOccurrenceExec::new(plan, occurrence_id, producer_identity))) } @@ -661,13 +722,15 @@ fn bind_gpu_function_host_occurrences( /// returns the occurrence ids that `plan`'s parent reads through it. /// /// Each entry records the host occurrences below it as predecessors, in the -/// order DataFusion executes them. A host-materializing producer without a -/// bound occurrence is an adapter defect: binding and collection must have -/// walked the same final plan. +/// order DataFusion executes them. Readers that share an occurrence id add +/// one entry, from the first reader in pre-order: their input subtrees are +/// copies of one call, whose nested occurrences share ids in the same way. A +/// host-materializing producer without a bound occurrence is an adapter +/// defect: binding and collection must have walked the same final plan. fn collect_gpu_function_occurrences( plan: &dyn ExecutionPlan, materializes_host_result: bool, - occurrences: &mut Vec, + occurrences: &mut Vec<(String, AttemptExecutionOccurrence)>, ) -> Result> { if plan.downcast_ref::().is_some_and(AlgeonNativeExec::materializes_host) { let mut predecessors = Vec::new(); @@ -682,6 +745,9 @@ fn collect_gpu_function_occurrences( } if let Some(bound) = plan.downcast_ref::() { let predecessors = collect_gpu_function_occurrences(bound.input().as_ref(), false, occurrences)?; + if occurrences.iter().any(|(occurrence_id, _)| occurrence_id == bound.occurrence_id()) { + return Ok(vec![bound.occurrence_id().to_string()]); + } let producer = gpu_resident_producer(bound.input().as_ref()).ok_or_else(|| { Error::new( ErrorCode::InvalidExecutionPlan, @@ -700,12 +766,15 @@ fn collect_gpu_function_occurrences( .with_log_safe_fact("reason", "bound_gpu_occurrence_host_plan_missing") .with_log_safe_fact("occurrence_id", bound.occurrence_id()) })?; - occurrences.push(AttemptExecutionOccurrence::host_arrow_for_producer( - bound.occurrence_id(), - bound.producer_identity(), - &host_plan, - predecessors, - producer.host_result_has_deferred_graph_facts(), + occurrences.push(( + bound.occurrence_id().to_string(), + AttemptExecutionOccurrence::host_arrow_for_producer( + bound.occurrence_id(), + bound.producer_identity(), + &host_plan, + predecessors, + producer.host_result_has_deferred_graph_facts(), + ), )); return Ok(vec![bound.occurrence_id().to_string()]); } @@ -1012,8 +1081,7 @@ mod tests { let (_backend, config) = backend_admission_config(); let graph = deferred_cugraph_exec(&config); let plan = Arc::new(CooperativeExec::new(graph)) as Arc; - let mut next_occurrence = 0; - let plan = bind_gpu_function_host_occurrences(plan, true, &mut next_occurrence) + let plan = bind_final_gpu_function_host_occurrences(plan, false) .expect("final-plan binding should assign graph occurrence lineage"); let (_, evidence) = plan_native_admission(plan.as_ref(), &config, None) diff --git a/crates/algeon-datafusion/tests/contract_tests_cpu/cuml_sql_tests.rs b/crates/algeon-datafusion/tests/contract_tests_cpu/cuml_sql_tests.rs index 7b8693f87..924e0bbf3 100644 --- a/crates/algeon-datafusion/tests/contract_tests_cpu/cuml_sql_tests.rs +++ b/crates/algeon-datafusion/tests/contract_tests_cpu/cuml_sql_tests.rs @@ -36,7 +36,7 @@ use datafusion::{ datasource::MemTable, execution::{SessionStateBuilder, context::QueryPlanner}, logical_expr::{LogicalPlan, UserDefinedLogicalNodeCore}, - physical_plan::ExecutionPlan, + physical_plan::{ExecutionPlan, displayable}, physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}, prelude::SessionContext, }; @@ -585,6 +585,39 @@ async fn sql_schemas_match_ir_and_physical_nodes() { } } +/// DataFusion copies a CTE at each reference. Both physical copies keep the +/// call site's invocation and bind one result name, while two call sites with +/// identical text keep separate names. +#[tokio::test] +async fn cte_references_share_one_invocation_and_separate_call_sites_do_not() { + let ctx = context(); + let umap = format!("{CUML_UMAP_FUNCTION}(input => (SELECT id, d0, d1 FROM vectors))"); + let shared = gpu_result_inputs( + &ctx, + &format!("WITH emb AS (SELECT * FROM {umap}) SELECT a.id FROM emb a JOIN emb b ON a.id = b.id"), + ) + .await; + assert_eq!(shared.len(), 2, "{shared:?}"); + assert_eq!(shared[0], shared[1], "inlined references of one call must read one result"); + let separate = gpu_result_inputs(&ctx, &format!("SELECT a.id FROM {umap} a JOIN {umap} b ON a.id = b.id")).await; + assert_eq!(separate.len(), 2, "{separate:?}"); + assert_ne!(separate[0], separate[1], "separately authored calls must stay independent"); +} + +/// The `gpu_result_input` of every cuML exec in the optimized physical plan. +async fn gpu_result_inputs(ctx: &SessionContext, sql: &str) -> Vec { + let logical = ctx.state().create_logical_plan(sql).await.unwrap(); + let physical = ctx.state().create_physical_plan(&logical).await.unwrap(); + displayable(physical.as_ref()) + .indent(false) + .to_string() + .lines() + .filter(|line| line.contains("CumlAlgoExec")) + .filter_map(|line| line.split("gpu_result_input=").nth(1)) + .map(|name| name.split(',').next().unwrap_or(name).trim().to_string()) + .collect() +} + #[tokio::test] async fn ordered_input_relations_keep_their_sort_through_physical_optimization() { let ctx = context(); diff --git a/crates/algeon-datafusion/tests/contract_tests_cpu/cuvs_sql_tests/mod.rs b/crates/algeon-datafusion/tests/contract_tests_cpu/cuvs_sql_tests/mod.rs index ed1c2e68f..13d87ebad 100644 --- a/crates/algeon-datafusion/tests/contract_tests_cpu/cuvs_sql_tests/mod.rs +++ b/crates/algeon-datafusion/tests/contract_tests_cpu/cuvs_sql_tests/mod.rs @@ -46,7 +46,7 @@ use datafusion::{ physical_expr::PhysicalExpr, physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, ReplaceChildrenOptions, - SendableRecordBatchStream, + SendableRecordBatchStream, displayable, }, physical_planner::DefaultPhysicalPlanner, prelude::SessionContext, @@ -741,6 +741,42 @@ async fn cte_relations_are_preserved_as_two_logical_children() { assert!(!queries.contains("EmptyRelation"), "{queries}"); } +/// DataFusion copies a CTE at each reference. Both physical copies keep the +/// call site's invocation and bind one result name, while two call sites with +/// identical text keep separate names. +#[tokio::test] +async fn cte_references_share_one_invocation_and_separate_call_sites_do_not() { + let ctx = cuvs_context(); + let kmeans = format!( + "{CUVS_KMEANS_FUNCTION}(input => (SELECT item_id AS id, d0, d1 FROM dataset_vectors), n_clusters => 1)" + ); + let shared = gpu_result_inputs( + &ctx, + &format!("WITH km AS (SELECT id, cluster_id FROM {kmeans}) SELECT a.id FROM km a JOIN km b ON a.id = b.id"), + ) + .await; + assert_eq!(shared.len(), 2, "{shared:?}"); + assert_eq!(shared[0], shared[1], "inlined references of one call must read one result"); + let separate = + gpu_result_inputs(&ctx, &format!("SELECT a.id FROM {kmeans} a JOIN {kmeans} b ON a.id = b.id")).await; + assert_eq!(separate.len(), 2, "{separate:?}"); + assert_ne!(separate[0], separate[1], "separately authored calls must stay independent"); +} + +/// The `gpu_result_input` of every cuVS exec in the optimized physical plan. +async fn gpu_result_inputs(ctx: &SessionContext, sql: &str) -> Vec { + let logical = logical_plan(ctx, sql).await; + let physical = ctx.state().create_physical_plan(&logical).await.expect("cuVS physical planning"); + displayable(physical.as_ref()) + .indent(false) + .to_string() + .lines() + .filter(|line| line.contains("CuvsAlgoExec")) + .filter_map(|line| line.split("gpu_result_input=").nth(1)) + .map(|name| name.split(',').next().unwrap_or(name).trim().to_string()) + .collect() +} + #[tokio::test] async fn direct_join_filter_projection_subqueries_remain_real_children() { let ctx = cuvs_context(); diff --git a/crates/algeon-datafusion/tests/cugraph_tests_gpu/cugraph_sql_tests/gpu.rs b/crates/algeon-datafusion/tests/cugraph_tests_gpu/cugraph_sql_tests/gpu.rs index 08c68f677..2594b6614 100644 --- a/crates/algeon-datafusion/tests/cugraph_tests_gpu/cugraph_sql_tests/gpu.rs +++ b/crates/algeon-datafusion/tests/cugraph_tests_gpu/cugraph_sql_tests/gpu.rs @@ -258,7 +258,10 @@ async fn functions_only_keeps_sql_limit_outside_the_cugraph_function() { assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); } -async fn assert_repeated_cugraph_cte_executes(mode: AlgeonGpuExecutionMode, expected_gpu_executions: usize) { +/// DataFusion copies the CTE at both references. Functions-only execution +/// keeps both copies as host occurrences and a native join reads one producer; +/// either way the two references read one evaluation. +async fn assert_repeated_cugraph_cte_executes(mode: AlgeonGpuExecutionMode, displayed_occurrences: usize) { let (_backend, ctx) = backend_cugraph_sql_context_with_mode(SessionStateBuilder::new_with_default_features(), mode); register_edge_memtable(&ctx, "shared_cte_edges", DataType::Int64, DataType::Int64, DataType::Float64); let plan = ctx @@ -279,19 +282,15 @@ async fn assert_repeated_cugraph_cte_executes(mode: AlgeonGpuExecutionMode, expe let rendered = displayable(plan.as_ref()).indent(false).to_string(); assert_eq!( rendered.matches("CugraphAlgoExec: algorithm=DegreesAll").count(), - expected_gpu_executions, - "the execution policy must preserve its declared CTE sharing behavior:\n{rendered}" + displayed_occurrences, + "the execution policy must preserve its declared CTE placement:\n{rendered}" ); let batches = collect(Arc::clone(&plan), ctx.task_ctx()) .await .expect("both references to one cuGraph table provider must execute"); assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); let metrics = NativeDataFusionExecutionMetrics::from_metrics(&plan_metrics_recursive(plan.as_ref())); - assert_eq!( - metrics.graph_runtime_rows.len(), - expected_gpu_executions, - "each unshared graph occurrence must execute once" - ); + assert_eq!(metrics.graph_runtime_rows.len(), 1, "both CTE references must read one graph evaluation"); } #[tokio::test] diff --git a/crates/algeon-datafusion/tests/root_tests_gpu/cuvs_sql_gpu_tests.rs b/crates/algeon-datafusion/tests/root_tests_gpu/cuvs_sql_gpu_tests.rs index 913a0d3bd..ce6806d72 100644 --- a/crates/algeon-datafusion/tests/root_tests_gpu/cuvs_sql_gpu_tests.rs +++ b/crates/algeon-datafusion/tests/root_tests_gpu/cuvs_sql_gpu_tests.rs @@ -725,7 +725,7 @@ async fn assert_shared_cugraph_cte_lineage_to_two_cuvs_consumers(mode: AlgeonGpu .expect("both shared-producer branches must satisfy their own downstream occurrence"); assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 8); let metrics = NativeDataFusionExecutionMetrics::from_metrics(&recursive_plan_metrics(plan.as_ref())); - assert_eq!(metrics.graph_runtime_rows.len(), 2, "each unshared CTE occurrence executes once"); + assert_eq!(metrics.graph_runtime_rows.len(), 1, "both references of the ForceAtlas2 CTE read one evaluation"); close_cuvs_test_backend(backend, ctx); } @@ -1513,6 +1513,89 @@ async fn cte_aggregate_read_by_gpu_function_inputs_and_joins_is_computed_once() assert_eq!(run_shared_cte_functions_only(NESTED_PERSONAS_SQL, temp_dir.path()).await, rows); } +/// Both references of the `km` CTE read one K-means call site. +const KMEANS_SELF_JOIN_SQL: &str = "WITH km AS (\ + SELECT id, cluster_id FROM cuvs_kmeans(\ + input => (SELECT id, d0, d1 FROM cluster_points), n_clusters => 2, max_iter => 5, init => 'random')) \ + SELECT COUNT(*) AS rows_joined, SUM(CASE WHEN a.cluster_id <> b.cluster_id THEN 1 ELSE 0 END) AS disagreeing \ + FROM km a JOIN km b ON a.id = b.id"; + +/// Two call sites with identical text. +const KMEANS_TWO_CALLS_SQL: &str = "SELECT COUNT(*) AS rows_joined, \ + SUM(CASE WHEN a.cluster_id <> b.cluster_id THEN 1 ELSE 0 END) AS disagreeing \ + FROM cuvs_kmeans(\ + input => (SELECT id, d0, d1 FROM cluster_points), n_clusters => 2, max_iter => 5, init => 'random') a \ + JOIN cuvs_kmeans(\ + input => (SELECT id, d0, d1 FROM cluster_points), n_clusters => 2, max_iter => 5, init => 'random') b \ + ON a.id = b.id"; + +/// Registers 64 points with unique ids, half at (0, 1) and half at (10, 11). +async fn register_cluster_points(ctx: &SessionContext, dir: &std::path::Path) { + let path = dir.join("cluster_points.parquet"); + let schema = wide_vector_schema("id"); + let ids = (0..64).collect::>(); + let d0 = ids.iter().map(|id| if *id < 32 { 0.0 } else { 10.0 }).collect::>(); + let d1 = ids.iter().map(|id| if *id < 32 { 1.0 } else { 11.0 }).collect::>(); + let batch = wide_vector_batch(Arc::clone(&schema), "id", &ids, &d0, &d1); + let mut writer = ArrowWriter::try_new(File::create(&path).expect("create Parquet fixture"), schema, None) + .expect("create Parquet writer"); + writer.write(&batch).expect("write Parquet fixture"); + writer.close().expect("finish Parquet fixture"); + ctx.register_parquet("cluster_points", path.to_str().expect("UTF-8 Parquet path"), ParquetReadOptions::default()) + .await + .expect("register cluster points"); +} + +/// Runs `sql` under `mode` in one attempt and returns its joined and +/// disagreeing row counts with the number of vector algorithm tasks executed. +async fn run_kmeans_join(mode: AlgeonGpuExecutionMode, sql: &str, dir: &std::path::Path) -> (i64, i64, usize) { + let (backend, ctx) = cuvs_bounded_context_with_mode(mode); + register_cluster_points(&ctx, dir).await; + let attempt = backend.begin_attempt(AttemptOptions::default()).expect("begin cuVS attempt"); + let completion = attempt.completion(); + let plan = attempt.plan_sql(&ctx, sql).await.expect("plan the K-means join"); + let mut stream = attempt.execute_physical_plan(plan, &ctx).expect("execute the K-means join"); + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch.expect("K-means join batch")); + } + drop(stream); + let _report = completion.await; + let metrics = attempt.sealed_native_execution_snapshot().expect("sealed native execution snapshot").metrics; + let vector_tasks = metrics + .task_operator_rows + .iter() + .filter(|row| row.operator_kind == "vector_algorithm") + .map(|row| row.task_count) + .sum::(); + drop(attempt); + close_cuvs_test_backend(backend, ctx); + let [batch] = batches.as_slice() else { + panic!("the aggregate returns one batch: {batches:?}"); + }; + let count = |column: usize| batch.column(column).as_any().downcast_ref::().expect("count").value(0); + (count(0), count(1), vector_tasks) +} + +#[tokio::test] +async fn kmeans_cte_references_read_one_evaluation_and_separate_calls_do_not() { + let temp_dir = tempfile::tempdir().expect("temporary Parquet directory"); + // A native join reads one K-means producer for both references. + assert_eq!( + run_kmeans_join(AlgeonGpuExecutionMode::NativePreferred, KMEANS_SELF_JOIN_SQL, temp_dir.path()).await, + (64, 0, 1) + ); + // A CPU join reads two host occurrences that replay one evaluation. + assert_eq!( + run_kmeans_join(AlgeonGpuExecutionMode::FunctionsOnly, KMEANS_SELF_JOIN_SQL, temp_dir.path()).await, + (64, 0, 1) + ); + for mode in [AlgeonGpuExecutionMode::NativePreferred, AlgeonGpuExecutionMode::FunctionsOnly] { + let (rows_joined, _, vector_tasks) = run_kmeans_join(mode, KMEANS_TWO_CALLS_SQL, temp_dir.path()).await; + assert_eq!((rows_joined, vector_tasks), (64, 2), "identical call sites evaluate independently under {mode:?}"); + } +} + fn cuvs_test_backend_with_mode(mode: AlgeonGpuExecutionMode) -> AlgeonGpuBackend { let profile = AlgeonGpuDeviceProfile::new(0) .with_memory_limit(GpuDeviceMemoryLimit::Explicit { limit_bytes: 512 * 1024 * 1024 }); diff --git a/website/docs/cugraph/graph-options.mdx b/website/docs/cugraph/graph-options.mdx index ed2f899f2..790335bce 100644 --- a/website/docs/cugraph/graph-options.mdx +++ b/website/docs/cugraph/graph-options.mdx @@ -142,7 +142,8 @@ Look for both fields in the physical plan: ## Reuse algorithm results -A view stores a query definition. Later reads can execute the algorithm again. +A view stores a query definition. Each later query that reads it executes the +algorithm again. Community partitions and layout coordinates can change between executions. The [reproducibility section](../design/gpu-table-functions.mdx#reproducibility) lists which functions repeat. Use `CREATE OR REPLACE TABLE ... AS` when later diff --git a/website/docs/design/gpu-table-functions.mdx b/website/docs/design/gpu-table-functions.mdx index a2546772d..64438259f 100644 --- a/website/docs/design/gpu-table-functions.mdx +++ b/website/docs/design/gpu-table-functions.mdx @@ -182,12 +182,35 @@ Shared relational subplans across GPU function inputs and query roots Separate graph invocations that read an identical edge relation build separate graph instances. -Within a composed island execution, repeated readers of a cuGraph similarity -invocation share its GPU-resident result when relational edge and side-input -producers have deterministic GPU dependencies. The shared invocation identity -incorporates the logical function call and execution-scope fingerprints of its -input subtrees. Registered reader occurrences retain the result handle under the -attempt's grant until the final reader completes. +### One evaluation for each call site {/* #one-evaluation-per-invocation */} + +The relation planner gives each SQL call of a GPU function an invocation +identity. DataFusion copies a CTE or view at every reference, and each copy +keeps that identity, while two calls with identical text receive different +ones. Within one query execution, the references of one call read a single +evaluation. A function in the [reproducibility table](#reproducibility) then +returns one relation to every reader, apart from the separate evaluations +listed below. + +Inside a composed island, the first occurrence to run computes the +GPU-resident result. Later occurrences read that result and skip their own +input subplans. The handle stays under the attempt's grant until the final +reader completes. A result option pushed into one occurrence, such as a +`LIMIT` absorbed by a cuGraph producer, changes the relation that occurrence +computes and gives it a separate evaluation. + +References that cross into DataFusion operators as host Arrow share one +evaluation as well. The first reference that DataFusion polls runs the +function, and every reference replays the retained batches; the attempt +manifest records one host occurrence for them. The batches stay in host +memory until the last reference finishes and count against the DataFusion +memory pool when host output accounting is enabled. + +A composed island and the host boundaries keep separate results. When one +reference is read inside an island and another crosses into DataFusion, or +two islands each read a reference, the call runs once in each of them. +Separate queries, including repeated executions of one prepared statement, +evaluate it again. ## Reproducibility @@ -215,7 +238,9 @@ Two runs of the same query on the same data, options and GPU return the same rel Connected-component labels (the minimum vertex ID of each component), degrees, triangle counts and core numbers involve no floating-point accumulation and repeat. The similarity functions `cugraph_jaccard`, `cugraph_sorensen`, `cugraph_overlap` and `cugraph_cosine` contain no atomic accumulation. None of the remaining functions has a `seed` option (the `seeds` argument of the betweenness functions lists source vertices and draws no random numbers). -Floating-point results from atomic reductions require tolerance comparisons rather than bitwise equality. Community and cluster identifiers are arbitrary integers; validation checks partition membership rather than label values. To preserve non-deterministic results across queries, materialize the relation using `CREATE TABLE ... AS` ([Reuse algorithm results](../cugraph/graph-options.mdx#reuse-algorithm-results)). +Floating-point results from atomic reductions require tolerance comparisons rather than bitwise equality. Community and cluster identifiers are arbitrary integers; validation checks partition membership rather than label values. + +The table compares separate executions. Within one query, the references of one CTE or view read [one evaluation](#one-evaluation-per-invocation) and agree with each other, except where that section lists separate evaluations. Two calls written separately with the same arguments are two evaluations and can disagree. To preserve non-deterministic results across queries, materialize the relation using `CREATE TABLE ... AS` ([Reuse algorithm results](../cugraph/graph-options.mdx#reuse-algorithm-results)). A seed fixes random draws over one row order. Without a top-level `ORDER BY`, the row order of an input subquery can change between runs, thereby shifting seeded results as well. See [Input order](../cuml/ml-inputs.mdx#input-order) and [Relation identity](../cuvs/vector-inputs.mdx#relation-identity). From b581e0b5a604259751b968fe29e7c611ea1a8983 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Thu, 22 Oct 2026 00:27:13 +0000 Subject: [PATCH 4/8] feat(cudf): build zero-row frames from an Arrow schema on the device A read known to return no rows (scan LIMIT 0, every file or row group pruned) built its zero-row frame by uploading an empty Arrow batch, which counts as a host Arrow import and breaks device-resident contracts once a dynamic filter prunes a whole source. Runtime::empty_dataframe_from_arrow_schema maps each field through the host-import type mapping and builds the column from a cudf type tree in the new shim call cudf_column_create_empty_from_type_tree_ex: a LIST gets an empty INT32 offsets child and its element child, a STRUCT its field children, as cudf::empty_like does. No device memory is allocated and the columns carry the runtime's provenance. The engine's empty local Parquet frame uses it. Co-Authored-By: Claude Opus 5.5 --- .config/policies/public-api/algeon-cudf.txt | 4 + .../native/include/cudf_rust/column.h | 26 ++++ .../algeon-cudf-sys/native/src/abi_checks.cpp | 7 + .../native/src/c_header_smoke.c | 1 + .../algeon-cudf-sys/native/src/column.cpp | 70 ++++++++++ .../cudf/algeon-cudf-sys/src/bindings.rs | 15 +++ .../cudf/algeon-cudf/src/core/column/owned.rs | 24 ++++ .../src/interop/arrow_host/empty.rs | 100 ++++++++++++++ .../algeon-cudf/src/interop/arrow_host/mod.rs | 1 + .../frame/dataframe/arrow_tests.rs | 124 +++++++++++++++++- .../src/exec/source/scan_limit.rs | 33 +---- 11 files changed, 378 insertions(+), 27 deletions(-) create mode 100644 components/cudf/algeon-cudf/src/interop/arrow_host/empty.rs diff --git a/.config/policies/public-api/algeon-cudf.txt b/.config/policies/public-api/algeon-cudf.txt index cf5372a91..d9835352b 100644 --- a/.config/policies/public-api/algeon-cudf.txt +++ b/.config/policies/public-api/algeon-cudf.txt @@ -4178,6 +4178,8 @@ impl cudf::Runtime pub fn cudf::Runtime::dataframe_from_record_batch(&self, &arrow_array::record_batch::RecordBatch) -> cudf::core::error::Result pub fn cudf::Runtime::dataframe_from_record_batches(&self, &[arrow_array::record_batch::RecordBatch]) -> cudf::core::error::Result impl cudf::Runtime +pub fn cudf::Runtime::empty_dataframe_from_arrow_schema(&self, &arrow_schema::schema::Schema) -> cudf::core::error::Result +impl cudf::Runtime pub fn cudf::Runtime::narrow_packed_matrix_to_float32(&self, &cudf::PackedMatrix) -> cudf::core::error::Result pub fn cudf::Runtime::pack_columns(&self, &[&cudf::core::Column], cudf::DenseMatrixLayout) -> cudf::core::error::Result pub fn cudf::Runtime::pack_flat_column(&self, &cudf::core::Column, usize, usize, cudf::DenseMatrixLayout) -> cudf::core::error::Result @@ -5709,6 +5711,8 @@ impl cudf::Runtime pub fn cudf::Runtime::dataframe_from_record_batch(&self, &arrow_array::record_batch::RecordBatch) -> cudf::core::error::Result pub fn cudf::Runtime::dataframe_from_record_batches(&self, &[arrow_array::record_batch::RecordBatch]) -> cudf::core::error::Result impl cudf::Runtime +pub fn cudf::Runtime::empty_dataframe_from_arrow_schema(&self, &arrow_schema::schema::Schema) -> cudf::core::error::Result +impl cudf::Runtime pub fn cudf::Runtime::narrow_packed_matrix_to_float32(&self, &cudf::PackedMatrix) -> cudf::core::error::Result pub fn cudf::Runtime::pack_columns(&self, &[&cudf::core::Column], cudf::DenseMatrixLayout) -> cudf::core::error::Result pub fn cudf::Runtime::pack_flat_column(&self, &cudf::core::Column, usize, usize, cudf::DenseMatrixLayout) -> cudf::core::error::Result diff --git a/components/cudf/algeon-cudf-sys/native/include/cudf_rust/column.h b/components/cudf/algeon-cudf-sys/native/include/cudf_rust/column.h index 106975299..b918ac43d 100644 --- a/components/cudf/algeon-cudf-sys/native/include/cudf_rust/column.h +++ b/components/cudf/algeon-cudf-sys/native/include/cudf_rust/column.h @@ -74,6 +74,32 @@ cudf_status_t cudf_column_create_empty_ex(cudf_type_id_t type_id, cudf_exec_options_t const* options, cudf_column_handle* out); +/// One node of a column type tree laid out in depth-first pre-order. The +/// `num_children` nodes after a parent describe its children, each followed +/// by its own subtree. `scale` is read for DECIMAL types only. +typedef struct cudf_column_type_node_t { + cudf_type_id_t type_id; + int32_t scale; + int32_t num_children; +} cudf_column_type_node_t; + +/// Create a zero-row column whose type, including every nested child type, +/// is the pre-order tree `nodes[0..num_nodes)`. +/// +/// A LIST node takes exactly one child, its element type. A STRUCT node takes +/// one or more children, its field types. Every other node takes none. +/// EMPTY and DICTIONARY32 are rejected with `CUDF_ERROR_TYPE`; a malformed +/// tree, including nodes left over after the root subtree, is rejected with +/// `CUDF_ERROR_INVALID_ARGUMENT`. No device memory is allocated. +/// +/// `options` must be non-NULL and must name an explicit valid device, a +/// non-NULL stream, and a non-NULL memory resource. The returned column +/// records that execution provenance. +cudf_status_t cudf_column_create_empty_from_type_tree_ex(cudf_column_type_node_t const* nodes, + int32_t num_nodes, + cudf_exec_options_t const* options, + cudf_column_handle* out); + /// Create a numeric column from host data with explicit execution options. /// /// `options` must be non-NULL and must name an explicit valid device, a diff --git a/components/cudf/algeon-cudf-sys/native/src/abi_checks.cpp b/components/cudf/algeon-cudf-sys/native/src/abi_checks.cpp index f3f5ad239..c27de9509 100644 --- a/components/cudf/algeon-cudf-sys/native/src/abi_checks.cpp +++ b/components/cudf/algeon-cudf-sys/native/src/abi_checks.cpp @@ -4,6 +4,7 @@ #include "cudf_rust/ast.h" #include "cudf_rust/binary.h" +#include "cudf_rust/column.h" #include "cudf_rust/copying.h" #include "cudf_rust/datetime.h" #include "cudf_rust/device.h" @@ -202,6 +203,12 @@ static_assert(CUDF_AST_IS_NULL == static_cast(cudf::ast::ast_operator::IS_N static_assert(CUDF_AST_NOT == static_cast(cudf::ast::ast_operator::NOT)); static_assert(CUDF_AST_CAST_TO_FLOAT64 == static_cast(cudf::ast::ast_operator::CAST_TO_FLOAT64)); +static_assert(std::is_standard_layout_v, "cudf_column_type_node_t must be standard layout"); +static_assert(offsetof(cudf_column_type_node_t, type_id) == 0, "cudf_column_type_node_t::type_id must be the first field"); +static_assert(offsetof(cudf_column_type_node_t, scale) == sizeof(cudf_type_id_t), "cudf_column_type_node_t::scale must follow type_id"); +static_assert(offsetof(cudf_column_type_node_t, num_children) == offsetof(cudf_column_type_node_t, scale) + sizeof(int32_t), "cudf_column_type_node_t::num_children must follow scale"); +static_assert(sizeof(cudf_column_type_node_t) == 3 * sizeof(int32_t), "cudf_column_type_node_t must be three 32-bit fields"); + static_assert(std::is_standard_layout_v, "cudf_row_predicate_node_t must be standard layout"); static_assert(offsetof(cudf_row_predicate_node_t, kind) == 0, "cudf_row_predicate_node_t::kind must be the first field"); static_assert(sizeof(static_cast(nullptr)->lhs_index) == sizeof(int32_t), "cudf_row_predicate_node_t::lhs_index must be 32-bit"); diff --git a/components/cudf/algeon-cudf-sys/native/src/c_header_smoke.c b/components/cudf/algeon-cudf-sys/native/src/c_header_smoke.c index 8024e0596..bed23b0f3 100644 --- a/components/cudf/algeon-cudf-sys/native/src/c_header_smoke.c +++ b/components/cudf/algeon-cudf-sys/native/src/c_header_smoke.c @@ -72,6 +72,7 @@ static int cudf_rust_c_header_smoke(void) { SMOKE_CHECK_FN(cudf_status_t, cudf_column_adopt_device_buffer_nullable, cudf_type_id_t, cudf_device_buffer_handle, cudf_device_buffer_handle, int32_t, int32_t, cudf_column_handle*); // Column and table creation, copying, and inspection + SMOKE_CHECK_FN(cudf_status_t, cudf_column_create_empty_from_type_tree_ex, cudf_column_type_node_t const*, int32_t, cudf_exec_options_t const*, cudf_column_handle*); SMOKE_CHECK_FN(cudf_status_t, cudf_column_create_numeric_ex, cudf_type_id_t, void const*, int32_t, uint8_t const*, int32_t, cudf_exec_options_t const*, cudf_column_handle*); SMOKE_CHECK_FN(cudf_status_t, cudf_column_create_fixed_width_from_host_ex, cudf_type_id_t, int32_t, void const*, int32_t, uint8_t const*, int32_t, cudf_exec_options_t const*, cudf_column_handle*); SMOKE_CHECK_FN(cudf_status_t, cudf_column_create_sequence_ex, cudf_scalar_handle, cudf_scalar_handle, int32_t, cudf_exec_options_t const*, cudf_column_handle*); diff --git a/components/cudf/algeon-cudf-sys/native/src/column.cpp b/components/cudf/algeon-cudf-sys/native/src/column.cpp index 5fab16d0f..3dc14605d 100644 --- a/components/cudf/algeon-cudf-sys/native/src/column.cpp +++ b/components/cudf/algeon-cudf-sys/native/src/column.cpp @@ -77,6 +77,47 @@ cudf_status_t resolve_column_view_null_count(cudf::column_view const& view, return CUDF_SUCCESS; } +// Mirrors cudf::empty_like for nested columns: a zero-row LIST holds an empty +// INT32 offsets child before its element child, and no level owns a buffer. +std::unique_ptr make_empty_column_from_type_tree(cudf_column_type_node_t const* nodes, + int32_t num_nodes, + int32_t& cursor) +{ + CUDF_EXPECTS(cursor < num_nodes, + "type tree ends before every declared child is described", + std::invalid_argument); + auto const& node = nodes[cursor++]; + CUDF_EXPECTS(node.num_children >= 0, "type tree num_children must be >= 0", std::invalid_argument); + CUDF_EXPECTS(node.type_id > CUDF_TYPE_EMPTY && node.type_id <= CUDF_TYPE_STRUCT && + node.type_id != CUDF_TYPE_DICTIONARY32, + "type tree node has no empty column construction", + cudf::data_type_error); + + std::vector> children; + if (node.type_id == CUDF_TYPE_LIST) { + CUDF_EXPECTS( + node.num_children == 1, "type tree LIST node requires exactly one child", std::invalid_argument); + children.push_back(cudf::make_empty_column(cudf::data_type{cudf::type_id::INT32})); + children.push_back(make_empty_column_from_type_tree(nodes, num_nodes, cursor)); + } else if (node.type_id == CUDF_TYPE_STRUCT) { + CUDF_EXPECTS( + node.num_children > 0, "type tree STRUCT node requires at least one child", std::invalid_argument); + children.reserve(static_cast(node.num_children)); + for (int32_t child = 0; child < node.num_children; ++child) { + children.push_back(make_empty_column_from_type_tree(nodes, num_nodes, cursor)); + } + } else { + CUDF_EXPECTS(node.num_children == 0, "type tree leaf node must have no children", std::invalid_argument); + return cudf::make_empty_column(host_fixed_width_dtype(node.type_id, node.scale)); + } + return std::make_unique(cudf::data_type{to_cudf_type_id(node.type_id)}, + 0, + rmm::device_buffer{}, + rmm::device_buffer{}, + 0, + std::move(children)); +} + } // namespace static cudf_status_t create_column_from_device_data_impl(cudf_type_id_t type_id, @@ -149,6 +190,35 @@ cudf_status_t cudf_column_create_empty_ex(cudf_type_id_t type_id, CUDF_RUST_CATCH } +cudf_status_t cudf_column_create_empty_from_type_tree_ex(cudf_column_type_node_t const* nodes, + int32_t num_nodes, + cudf_exec_options_t const* options, + cudf_column_handle* out) +{ + CUDF_RUST_NULL_CHECK(out, "out is null"); + *out = nullptr; + CUDF_RUST_NULL_CHECK(nodes, "nodes is null"); + if (num_nodes <= 0) { + set_error(CUDF_ERROR_INVALID_ARGUMENT, "num_nodes must be > 0"); + return CUDF_ERROR_INVALID_ARGUMENT; + } + CUDF_RUST_TRY + { + auto exec = std::unique_ptr{}; + cudf_status_t const exec_status = cudf_resolve_required_exec_options(options, &exec); + if (exec_status != CUDF_SUCCESS) { return exec_status; } + + int32_t cursor = 0; + auto col = make_empty_column_from_type_tree(nodes, num_nodes, cursor); + CUDF_EXPECTS(cursor == num_nodes, + "type tree has nodes after the root subtree", + std::invalid_argument); + *out = cudf_release_column_handle(std::move(col), cudf_capture_exec_provenance(*exec)); + return CUDF_SUCCESS; + } + CUDF_RUST_CATCH +} + cudf_status_t cudf_column_create_numeric_ex(cudf_type_id_t type_id, void const* data, int32_t size, diff --git a/components/cudf/algeon-cudf-sys/src/bindings.rs b/components/cudf/algeon-cudf-sys/src/bindings.rs index 1a48d65ae..d55c13906 100644 --- a/components/cudf/algeon-cudf-sys/src/bindings.rs +++ b/components/cudf/algeon-cudf-sys/src/bindings.rs @@ -274,6 +274,14 @@ pub struct cudf_interop_semantic_column_t { pub num_child_column_ids: i32, pub keys_column_id: u32, } +#[doc = " One node of a column type tree laid out in depth-first pre-order. The\n `num_children` nodes after a parent describe its children, each followed\n by its own subtree. `scale` is read for DECIMAL types only."] +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cudf_column_type_node_t { + pub type_id: cudf_type_id_t, + pub scale: i32, + pub num_children: i32, +} #[repr(C)] #[derive(Debug)] pub struct cudf_ast_expression_handle_state { @@ -1041,6 +1049,13 @@ unsafe extern "C" { options: *const cudf_exec_options_t, out: *mut cudf_column_handle, ) -> cudf_status_t; + #[doc = " Create a zero-row column whose type, including every nested child type,\n is the pre-order tree `nodes[0..num_nodes)`.\n\n A LIST node takes exactly one child, its element type. A STRUCT node takes\n one or more children, its field types. Every other node takes none.\n EMPTY and DICTIONARY32 are rejected with `CUDF_ERROR_TYPE`; a malformed\n tree, including nodes left over after the root subtree, is rejected with\n `CUDF_ERROR_INVALID_ARGUMENT`. No device memory is allocated.\n\n `options` must be non-NULL and must name an explicit valid device, a\n non-NULL stream, and a non-NULL memory resource. The returned column\n records that execution provenance."] + pub fn cudf_column_create_empty_from_type_tree_ex( + nodes: *const cudf_column_type_node_t, + num_nodes: i32, + options: *const cudf_exec_options_t, + out: *mut cudf_column_handle, + ) -> cudf_status_t; #[doc = " Create a numeric column from host data with explicit execution options.\n\n `options` must be non-NULL and must name an explicit valid device, a\n non-NULL stream, and a non-NULL memory resource. `options->mr` must outlive\n the returned column.\n\n If host-source copies are queued, the selected stream is synchronized\n before return on success or failure so the borrowed host buffers are no\n longer in use."] pub fn cudf_column_create_numeric_ex( type_id: cudf_type_id_t, diff --git a/components/cudf/algeon-cudf/src/core/column/owned.rs b/components/cudf/algeon-cudf/src/core/column/owned.rs index 3a70dcb22..9873351c4 100644 --- a/components/cudf/algeon-cudf/src/core/column/owned.rs +++ b/components/cudf/algeon-cudf/src/core/column/owned.rs @@ -193,6 +193,30 @@ impl Column { Self::from_handle_with_domain(handle, domain) } + /// Creates a zero-row column whose type, children included, is the pre-order type tree `nodes`. + pub(crate) fn empty_from_type_tree_in_domain( + nodes: &[cudf_sys::cudf_column_type_node_t], + domain: Arc, + operation: &'static str, + ) -> Result { + let num_nodes = usize_to_i32_len(nodes.len(), "type_tree_node_count", operation)?; + let mut handle = ptr::null_mut(); + let exec_options = domain.enter_execution()?; + error::check_status_in( + // SAFETY: `nodes` stays borrowed for the call and holds `num_nodes` initialized entries. + unsafe { + cudf_sys::cudf_column_create_empty_from_type_tree_ex( + nodes.as_ptr(), + num_nodes, + exec_options.options(), + &mut handle, + ) + }, + operation, + )?; + Self::from_handle_with_domain(handle, domain) + } + fn create_empty_with_dtype_in_domain( dtype: DataType, domain: Arc, diff --git a/components/cudf/algeon-cudf/src/interop/arrow_host/empty.rs b/components/cudf/algeon-cudf/src/interop/arrow_host/empty.rs new file mode 100644 index 000000000..f7a49696e --- /dev/null +++ b/components/cudf/algeon-cudf/src/interop/arrow_host/empty.rs @@ -0,0 +1,100 @@ +//! Zero-row device frames typed by an Arrow schema, built without host data. + +use std::sync::Arc; + +use arrow::datatypes::{DataType as ArrowDataType, Schema}; + +use super::schema::cudf_dtype_from_arrow_type_in; +use crate::{ + core::{ + column::Column, + error::{self, Error, ErrorCode, Result}, + types::{DataType, TypeId}, + }, + frame::dataframe::DataFrame, + runtime::Runtime, +}; + +const OPERATION: &str = "arrow_host::empty_dataframe_from_arrow_schema"; + +impl Runtime { + /// Create a zero-row `DataFrame` typed by an Arrow schema inside this + /// runtime, without a host copy. + /// + /// Field names become column names in schema order. Each field gets the + /// cudf type that [`DataFrame::from_record_batch`] would create for it, + /// including the element type of a list, so the frame concatenates with + /// non-empty frames of the same schema. A `Struct` field becomes a cudf + /// `STRUCT` column whose children follow the same rules. Columns carry + /// this runtime's retained provenance and allocate no device memory. Arrow + /// nullability is not encoded in cudf types and is ignored. A schema with + /// no fields gives the same frame as [`Runtime::empty_dataframe`]. + /// + /// # Errors + /// + /// Returns [`ErrorCode::UnsupportedDtype`] for a field type outside that + /// mapping, such as `Dictionary` or a `Struct` without fields, and a + /// duplicate column name error when two fields share a name. + pub fn empty_dataframe_from_arrow_schema(&self, schema: &Schema) -> Result { + if schema.fields().is_empty() { + return Ok(self.empty_dataframe()); + } + let mut names = Vec::with_capacity(schema.fields().len()); + let mut columns = Vec::with_capacity(schema.fields().len()); + let mut nodes = Vec::new(); + for field in schema.fields() { + nodes.clear(); + push_cudf_type_tree(field.data_type(), &mut nodes) + .map_err(|error| error.with_detail("field_name", field.name().clone()))?; + columns.push(Column::empty_from_type_tree_in_domain( + &nodes, + Arc::clone(self.execution_domain()), + OPERATION, + )?); + names.push(field.name().clone()); + } + DataFrame::new(names, columns).map_err(|error| error.reanchor(OPERATION)) + } +} + +/// Appends the cudf type tree of `data_type` in pre-order. Leaf and list types +/// come from the host import mapping, which also decides the element types a +/// list may carry; only `Struct` is expanded here. +fn push_cudf_type_tree(data_type: &ArrowDataType, nodes: &mut Vec) -> Result<()> { + if let ArrowDataType::Struct(fields) = data_type { + if fields.is_empty() { + return Err(Error::new_with( + ErrorCode::UnsupportedDtype, + OPERATION, + "a cudf STRUCT column requires at least one child field", + ) + .with_detail("arrow_type", format!("{data_type:?}")) + .with_detail("reason", "struct_without_fields")); + } + nodes.push(type_node(DataType::new(TypeId::Struct), fields.len())?); + for field in fields { + push_cudf_type_tree(field.data_type(), nodes) + .map_err(|error| error.with_detail("child_field_name", field.name().clone()))?; + } + return Ok(()); + } + + let dtype = cudf_dtype_from_arrow_type_in(data_type, OPERATION)?; + let element = match data_type { + ArrowDataType::List(child) | ArrowDataType::LargeList(child) | ArrowDataType::FixedSizeList(child, _) => { + Some(child) + }, + _ => None, + }; + nodes.push(type_node(dtype, usize::from(element.is_some()))?); + match element { + Some(child) => push_cudf_type_tree(child.data_type(), nodes), + None => Ok(()), + } +} + +fn type_node(dtype: DataType, num_children: usize) -> Result { + let num_children = i32::try_from(num_children) + .map_err(|_| error::range_limit_error(OPERATION, "num_children", "i32::MAX", num_children))?; + Ok(cudf_sys::cudf_column_type_node_t { type_id: dtype.id().to_raw(), scale: dtype.scale(), num_children }) +} diff --git a/components/cudf/algeon-cudf/src/interop/arrow_host/mod.rs b/components/cudf/algeon-cudf/src/interop/arrow_host/mod.rs index 4bc81371f..034c02ffc 100644 --- a/components/cudf/algeon-cudf/src/interop/arrow_host/mod.rs +++ b/components/cudf/algeon-cudf/src/interop/arrow_host/mod.rs @@ -27,6 +27,7 @@ //! assert_eq!(df2.num_columns(), 2); //! ``` mod batch; +mod empty; mod estimate; mod export; mod import; diff --git a/components/cudf/algeon-cudf/tests/semantic_contracts_gpu/frame/dataframe/arrow_tests.rs b/components/cudf/algeon-cudf/tests/semantic_contracts_gpu/frame/dataframe/arrow_tests.rs index 3d3b4f887..cbcb2a3e4 100644 --- a/components/cudf/algeon-cudf/tests/semantic_contracts_gpu/frame/dataframe/arrow_tests.rs +++ b/components/cudf/algeon-cudf/tests/semantic_contracts_gpu/frame/dataframe/arrow_tests.rs @@ -14,6 +14,7 @@ use arrow::{ }; use cudf::{ DataFrame, DataType, ErrorCode, ErrorKind, Exec, TypeId, + core::column::ColumnView, frame::dataframe::GpuRecordBatch, interop::arrow_host::{all_null_list_series, cudf_dtype_from_arrow_type, estimate_list_import_device_memory_bytes}, testing::{ @@ -29,7 +30,8 @@ use cudf::{ use crate::common::{ column_from_slice, column_from_slice_nullable, column_from_strs, column_from_strs_nullable, dataframe_from_columns, - error_context_value, nullable_series, sample_runtime, series, + error_context_value, int64_array, nested_dataframe, nullable_series, sample_runtime, series, struct_array, + with_test_runtime, }; fn import_record_batch(batch: &RecordBatch) -> cudf::Result { @@ -2037,6 +2039,118 @@ fn test_from_record_batches_rejects_schema_mismatch() { assert_eq!(error.code(), ErrorCode::TypeMismatch); } +/// Pre-order cudf types of a column, skipping list offsets children so a +/// zero-row column and a non-empty one of the same type compare equal. +fn logical_type_tree(view: &ColumnView<'_>, out: &mut Vec) { + let dtype = view.data_type().unwrap(); + out.push(dtype); + match dtype.id() { + TypeId::List => logical_type_tree(&view.child(1).unwrap(), out), + TypeId::Struct => { + for index in 0..view.num_children().unwrap() { + logical_type_tree(&view.child(index).unwrap(), out); + } + }, + _ => {}, + } +} + +fn frame_type_trees(frame: &DataFrame) -> Vec<(String, Vec)> { + frame + .column_names() + .into_iter() + .map(|name| { + let mut tree = Vec::new(); + logical_type_tree(&frame.series(name).unwrap().testing_view().unwrap(), &mut tree); + (name.to_string(), tree) + }) + .collect() +} + +fn test_empty_dataframe_from_arrow_schema_keeps_nested_types_without_host_import() { + let weights = ListArray::from_iter_primitive::([ + Some(vec![Some(0.5)]), + Some(vec![]), + Some(vec![Some(1.5), Some(2.5)]), + ]); + let columns: Vec<(&str, ArrayRef)> = vec![ + ("id", Arc::new(Int32Array::from(vec![1, 2, 3]))), + ("name", Arc::new(StringArray::from(vec!["a", "b", "c"]))), + ("amount", Arc::new(Decimal128Array::from(vec![100, 250, -5]).with_precision_and_scale(38, 2).unwrap())), + ("seen_at", Arc::new(TimestampMicrosecondArray::from(vec![1, 2, 3]).with_timezone("UTC"))), + ( + "scores", + Arc::new(ListArray::from_iter_primitive::([ + Some(vec![Some(1.0)]), + Some(vec![]), + Some(vec![Some(2.0), Some(3.0)]), + ])), + ), + ( + "embedding", + Arc::new(FixedSizeListArray::from_iter_primitive::( + [Some(vec![Some(1.0), Some(2.0)]), Some(vec![Some(3.0), Some(4.0)]), Some(vec![Some(5.0), Some(6.0)])], + 2, + )), + ), + ("info", struct_array(vec![("count", int64_array(&[7, 8, 9])), ("weights", Arc::new(weights))])), + ]; + let schema = RecordBatch::try_from_iter(columns.clone()).unwrap().schema(); + + let (empty, rows) = with_test_runtime(|runtime| { + reset_arrow_host_interop_observation_for_test(); + let empty = runtime.empty_dataframe_from_arrow_schema(&schema).unwrap(); + assert_eq!(arrow_host_import_call_count_for_test(), 0); + (empty, nested_dataframe(runtime, columns)) + }); + + let decimal = DataType::fixed_point(TypeId::Decimal128, -2).unwrap(); + let expected = [ + ("id", vec![DataType::new(TypeId::Int32)]), + ("name", vec![DataType::new(TypeId::String)]), + ("amount", vec![decimal]), + ("seen_at", vec![DataType::new(TypeId::TimestampMicroseconds)]), + ("scores", vec![DataType::new(TypeId::List), DataType::new(TypeId::Float64)]), + ("embedding", vec![DataType::new(TypeId::List), DataType::new(TypeId::Float32)]), + ( + "info", + vec![ + DataType::new(TypeId::Struct), + DataType::new(TypeId::Int64), + DataType::new(TypeId::List), + DataType::new(TypeId::Float64), + ], + ), + ] + .map(|(name, tree)| (name.to_string(), tree)) + .to_vec(); + assert_eq!(empty.num_rows(), 0); + assert_eq!(frame_type_trees(&empty), expected); + assert_eq!(frame_type_trees(&rows), expected); + + let combined = empty.on(Exec::RETAINED).concat(&rows).unwrap(); + assert_eq!(combined.num_rows(), 3); + assert_eq!(frame_type_trees(&combined), expected); + assert_eq!(combined.series("id").unwrap().to_vec_strict::().unwrap(), vec![1, 2, 3]); +} + +fn test_empty_dataframe_from_arrow_schema_rejects_unmapped_field_type() { + let schema = Schema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new( + "code", + ArrowDataType::Dictionary(Box::new(ArrowDataType::Int32), Box::new(ArrowDataType::Utf8)), + true, + ), + ]); + + let err = sample_runtime().empty_dataframe_from_arrow_schema(&schema).unwrap_err(); + + assert_eq!(err.code(), ErrorCode::UnsupportedDtype); + assert_eq!(err.operation(), "arrow_host::empty_dataframe_from_arrow_schema"); + assert_eq!(error_context_value(&err, "field_name").as_deref(), Some("code")); +} + pub(super) const CASES: &[(&str, fn())] = &[ ( "arrow_tests::test_arrow_datatype_try_from_type_id_supported_mappings", @@ -2291,4 +2405,12 @@ pub(super) const CASES: &[(&str, fn())] = &[ ), ("arrow_tests::test_from_record_batches_rejects_empty_input", test_from_record_batches_rejects_empty_input), ("arrow_tests::test_from_record_batches_rejects_schema_mismatch", test_from_record_batches_rejects_schema_mismatch), + ( + "arrow_tests::test_empty_dataframe_from_arrow_schema_keeps_nested_types_without_host_import", + test_empty_dataframe_from_arrow_schema_keeps_nested_types_without_host_import, + ), + ( + "arrow_tests::test_empty_dataframe_from_arrow_schema_rejects_unmapped_field_type", + test_empty_dataframe_from_arrow_schema_rejects_unmapped_field_type, + ), ]; diff --git a/crates/algeon-query-engine/src/exec/source/scan_limit.rs b/crates/algeon-query-engine/src/exec/source/scan_limit.rs index 5a89ce034..dae3ebc36 100644 --- a/crates/algeon-query-engine/src/exec/source/scan_limit.rs +++ b/crates/algeon-query-engine/src/exec/source/scan_limit.rs @@ -1,41 +1,22 @@ use std::time::Instant; -use arrow::{ - array::new_empty_array, - record_batch::{RecordBatch, RecordBatchOptions}, -}; - use super::{ - DataFrame, ExecutionContext, ExecutionMetrics, LocalParquetSource, QueryError, QueryErrorCode, Result, SourceChunk, - cudf_query_error, local_parquet_execution_frame_schema, parquet_chunk_path_fact, + DataFrame, ExecutionContext, ExecutionMetrics, LocalParquetSource, Result, SourceChunk, cudf_query_error, + local_parquet_execution_frame_schema, parquet_chunk_path_fact, }; -/// Builds a zero-row device frame with the execution-frame schema, by uploading an empty Arrow -/// batch. It stands in for a read that is known to return no rows, such as a scan limit of 0 or a -/// source whose files were all pruned. +/// Builds a zero-row device frame with the execution-frame schema from the schema alone, with no +/// host Arrow import. It stands in for a read that is known to return no rows, such as a scan limit +/// of 0 or a source whose files or row groups were all pruned. pub(super) fn empty_local_parquet_execution_frame( context: &ExecutionContext, source: &LocalParquetSource, ) -> Result { let schema = local_parquet_execution_frame_schema(source)?; - let columns = schema.fields().iter().map(|field| new_empty_array(field.data_type())).collect(); - let batch = RecordBatch::try_new_with_options( - std::sync::Arc::clone(&schema), - columns, - &RecordBatchOptions::new().with_row_count(Some(0)), - ) - .map_err(|error| { - QueryError::new( - QueryErrorCode::Internal, - "query_engine.local_parquet_source.empty_scan_limit", - "failed to construct zero-row local parquet source batch", - ) - .with_source(error) - })?; - DataFrame::from_record_batch(&context.runtime()?, &batch).map_err(|error| { + context.runtime()?.empty_dataframe_from_arrow_schema(&schema).map_err(|error| { cudf_query_error( "query_engine.local_parquet_source.empty_scan_limit", - "failed to upload zero-row local parquet source batch", + "failed to build zero-row local parquet source frame", error, ) }) From 58dbe35c76501b8d39bef36f783ae89fffaeb035 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Thu, 22 Oct 2026 00:27:13 +0000 Subject: [PATCH 5/8] fix(planning): route dynamic filters into every join input that drops unmatched rows Route discovery accepted only inner joins, so `x IN (SELECT ...)`, which plans as a left semi join, read every row group of the probe source: on 2.61 billion rows a one-column count took 0.59 s and 8.5 GiB against 0.06 s and 1.2 GiB for the literal list. join_type_discards_unmatched_rows now decides which input a key payload may reduce: both inputs of inner and semi joins, the right input of left outer, non-null-aware anti and mark joins, neither input of null-aware anti, full and cross joins. Key routes, keyset-constraint proofs and residual-equality candidates all use it. The TPC-H rule golden gains routes on q04, q17, q18 and q20. The same count now takes 0.05 s and 1.2 GiB. Co-Authored-By: Claude Opus 5.5 --- .../goldens/tpch_optimizer_rule_rows.txt | 10 +- .../native_report/report_execution_tests.rs | 2 +- .../dynamic_filter/fixtures.rs | 65 +++++++---- .../dynamic_filter/shared_sources.rs | 39 ++++++- .../dynamic_filter/source_routes.rs | 106 ++++++++++++++++++ .../aggregate_join_pushdown.rs | 11 +- .../root_joins/hash/dynamic_filter.rs | 92 +++++++++++++++ .../src/dynamic_filter.rs | 3 +- .../src/optimizer/dynamic_filter.rs | 71 ++++++++---- .../dynamic_filter/keyset_constraints.rs | 24 ++-- .../dynamic_filter/residual_equality.rs | 58 +++++----- website/docs/design/gpu-execution.mdx | 13 +++ 12 files changed, 403 insertions(+), 91 deletions(-) diff --git a/crates/algeon-datafusion/src/native/lowering_tests/goldens/tpch_optimizer_rule_rows.txt b/crates/algeon-datafusion/src/native/lowering_tests/goldens/tpch_optimizer_rule_rows.txt index bc65a44a3..5635806f7 100644 --- a/crates/algeon-datafusion/src/native/lowering_tests/goldens/tpch_optimizer_rule_rows.txt +++ b/crates/algeon-datafusion/src/native/lowering_tests/goldens/tpch_optimizer_rule_rows.txt @@ -22,7 +22,7 @@ q04 0 rule=projection_simplification,considered_count=8,applied_count=0,candidat q04 0 rule=semijoin_reduction,considered_count=1,applied_count=0,candidates=rule=semijoin_reduction,candidate_node_id=n6,before_root_node_id=n6,after_root_node_id=none,crossed_join_node_ids=none,target_subtree_id=n0,key_mappings=o_orderkey->o_orderkey@n0,cost_threshold_multiplier=8,cost_gate=none,applied=false,rejected_reason=UnsupportedLeftShape q04 0 rule=count_non_null_column_pruning,considered_count=1,applied_count=0,candidates=rule=count_non_null_column_pruning,aggregate_node_id=n8,column=__algeon_native_aggregate_arg_0,source_node_id=none,file_column_index=none,applied=false,rejected_reason=UnsupportedInputLineage q04 0 rule=shared_subplan_materialization,considered_count=1,candidate_count=0,applied_count=0,candidates=none -q04 0 rule=dynamic_filter_route_discovery,considered_count=1,route_count=0,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n6,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n5,consumer_input_node_id=n2,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape +q04 0 rule=dynamic_filter_route_discovery,considered_count=1,route_count=1,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n6,dynamic_filter_id=df0,producer_build_side=right,producer_input_node_id=n5,consumer_input_node_id=n2,consumer_kind=local_parquet_source,consumer_node_id=n0,consumer_source_node_id=n0,producer_key_column=l_orderkey,consumer_key_column=o_orderkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=172,total_source_bytes=172,routed=true,rejected_reason=none q04 0 rule=aggregate_join_pushdown,considered_count=1,applied_count=0,candidates=rule=aggregate_join_pushdown,aggregate_node_id=n8,join_node_id=n7,applied=false,rejected_reason=UnsupportedAggregateExpression q05 0 rule=projection_simplification,considered_count=11,applied_count=3,candidates=rule=projection_simplification,candidate_node_id=n21,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n18,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n2,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n4,replacement_node_id=n3,applied=true,rejected_reason=none|rule=projection_simplification,candidate_node_id=n8,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n11,replacement_node_id=n9,applied=true,rejected_reason=none|rule=projection_simplification,candidate_node_id=n15,replacement_node_id=n12,applied=true,rejected_reason=none|rule=projection_simplification,candidate_node_id=n19,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n14,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n4,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n12,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch q05 0 rule=join_build_side_selection,considered_count=10,applied_count=3,candidates=rule=join_build_side_selection,candidate_node_id=n4,before_build_side=right,after_build_side=none,left_estimated_rows=5,right_estimated_rows=3,left_estimated_bytes=223,right_estimated_bytes=125,left_advisory_rows=none,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=conservative_bytes,applied=false,rejected_reason=NoBetterBuildSide|rule=join_build_side_selection,candidate_node_id=n10,before_build_side=right,after_build_side=left,left_estimated_rows=3,right_estimated_rows=3,left_estimated_bytes=118,right_estimated_bytes=165,left_advisory_rows=none,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=conservative_bytes,applied=true,rejected_reason=none|rule=join_build_side_selection,candidate_node_id=n12,before_build_side=right,after_build_side=none,left_estimated_rows=none,right_estimated_rows=3,left_estimated_bytes=none,right_estimated_bytes=236,left_advisory_rows=3,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=small_current_side,applied=false,rejected_reason=NoBetterBuildSide|rule=join_build_side_selection,candidate_node_id=n13,before_build_side=right,after_build_side=left,left_estimated_rows=3,right_estimated_rows=none,left_estimated_bytes=118,right_estimated_bytes=none,left_advisory_rows=none,right_advisory_rows=3,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=small_other_side,applied=true,rejected_reason=none|rule=join_build_side_selection,candidate_node_id=n14,before_build_side=right,after_build_side=none,left_estimated_rows=none,right_estimated_rows=none,left_estimated_bytes=none,right_estimated_bytes=none,left_advisory_rows=3,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=none,applied=false,rejected_reason=NoBetterBuildSide|rule=join_build_side_selection,candidate_node_id=n5,before_build_side=left,after_build_side=none,left_estimated_rows=3,right_estimated_rows=3,left_estimated_bytes=118,right_estimated_bytes=165,left_advisory_rows=none,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=conservative_bytes,applied=false,rejected_reason=NoBetterBuildSide|rule=join_build_side_selection,candidate_node_id=n7,before_build_side=right,after_build_side=none,left_estimated_rows=none,right_estimated_rows=3,left_estimated_bytes=none,right_estimated_bytes=236,left_advisory_rows=3,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=small_current_side,applied=false,rejected_reason=NoBetterBuildSide|rule=join_build_side_selection,candidate_node_id=n8,before_build_side=left,after_build_side=none,left_estimated_rows=3,right_estimated_rows=none,left_estimated_bytes=118,right_estimated_bytes=none,left_advisory_rows=none,right_advisory_rows=3,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=small_current_side,applied=false,rejected_reason=NoBetterBuildSide|rule=join_build_side_selection,candidate_node_id=n13,before_build_side=right,after_build_side=none,left_estimated_rows=5,right_estimated_rows=3,left_estimated_bytes=223,right_estimated_bytes=125,left_advisory_rows=none,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=conservative_bytes,applied=false,rejected_reason=NoBetterBuildSide|rule=join_build_side_selection,candidate_node_id=n16,before_build_side=right,after_build_side=left,left_estimated_rows=none,right_estimated_rows=none,left_estimated_bytes=none,right_estimated_bytes=none,left_advisory_rows=none,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=grouped_aggregate_keyset,applied=true,rejected_reason=none @@ -106,7 +106,7 @@ q17 0 rule=join_build_side_selection,considered_count=2,applied_count=0,candidat q17 0 rule=dimension_derived_row_group_pruning,considered_count=2,applied_count=0,candidates=rule=dimension_derived_row_group_pruning,candidate_node_id=n4,dimension_source_node_id=n0,target_source_node_id=n3,dimension_key_column=p_partkey,target_key_column=l_partkey,derived_min_key=none,derived_max_key=none,selected_row_groups_before=none,selected_row_groups_after=none,max_dimension_rows=1000000,applied=false,rejected_reason=DimensionPredicateAbsent|rule=dimension_derived_row_group_pruning,candidate_node_id=n4,dimension_source_node_id=n3,target_source_node_id=n0,dimension_key_column=l_partkey,target_key_column=p_partkey,derived_min_key=none,derived_max_key=none,selected_row_groups_before=none,selected_row_groups_after=none,max_dimension_rows=1000000,applied=false,rejected_reason=DimensionPredicateAbsent q17 0 rule=local_parquet_source_reuse,considered_count=1,applied_count=0,candidates=rule=local_parquet_source_reuse,source_node_ids=n3+n6,reuse_count=2,stable_file_ids=local:/lineitem.parquet:5023:mtime=,superset_read_columns=l_partkey+l_quantity+l_extendedprice,branch_filter_count=0,max_variant_selected_bytes=524288000,applied=false,rejected_reason=UnsupportedSourceContract q17 0 rule=shared_subplan_materialization,considered_count=2,candidate_count=0,applied_count=0,candidates=none -q17 0 rule=dynamic_filter_route_discovery,considered_count=2,route_count=1,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n4,dynamic_filter_id=df0,producer_build_side=right,producer_input_node_id=n3,consumer_input_node_id=n2,consumer_kind=local_parquet_source,consumer_node_id=n0,consumer_source_node_id=n0,producer_key_column=l_partkey,consumer_key_column=p_partkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=198,total_source_bytes=198,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n9,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n8,consumer_input_node_id=n5,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape +q17 0 rule=dynamic_filter_route_discovery,considered_count=2,route_count=2,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n4,dynamic_filter_id=df0,producer_build_side=right,producer_input_node_id=n3,consumer_input_node_id=n2,consumer_kind=local_parquet_source,consumer_node_id=n0,consumer_source_node_id=n0,producer_key_column=l_partkey,consumer_key_column=p_partkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=198,total_source_bytes=198,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n9,dynamic_filter_id=df1,producer_build_side=right,producer_input_node_id=n8,consumer_input_node_id=n5,consumer_kind=local_parquet_source,consumer_node_id=n0,consumer_source_node_id=n0,producer_key_column=l_partkey,consumer_key_column=p_partkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=198,total_source_bytes=198,routed=true,rejected_reason=none q17 0 rule=aggregate_join_pushdown,considered_count=2,applied_count=0,candidates=rule=aggregate_join_pushdown,aggregate_node_id=n7,join_node_id=none,applied=false,rejected_reason=InputNotJoin|rule=aggregate_join_pushdown,aggregate_node_id=n10,join_node_id=n9,applied=false,rejected_reason=JoinResidualUnsupported q18 0 rule=projection_simplification,considered_count=5,applied_count=1,candidates=rule=projection_simplification,candidate_node_id=n14,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n6,replacement_node_id=n1,applied=true,rejected_reason=none|rule=projection_simplification,candidate_node_id=n3,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n13,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n5,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch q18 0 rule=semijoin_reduction,considered_count=1,applied_count=1,candidates=rule=semijoin_reduction,candidate_node_id=n9,before_root_node_id=n9,after_root_node_id=n9,crossed_join_node_ids=n4+n2,target_subtree_id=n1,key_mappings=o_orderkey->o_orderkey@n1,cost_threshold_multiplier=8,cost_gate=signal=aggregate_keyset_shape,threshold_multiplier=8,accepted=true,applied=true,rejected_reason=none @@ -114,7 +114,7 @@ q18 0 rule=join_build_side_selection,considered_count=4,applied_count=1,candidat q18 0 rule=dimension_derived_row_group_pruning,considered_count=4,applied_count=0,candidates=rule=dimension_derived_row_group_pruning,candidate_node_id=n7,dimension_source_node_id=n0,target_source_node_id=n1,dimension_key_column=c_custkey,target_key_column=o_custkey,derived_min_key=none,derived_max_key=none,selected_row_groups_before=none,selected_row_groups_after=none,max_dimension_rows=1000000,applied=false,rejected_reason=DimensionPredicateAbsent|rule=dimension_derived_row_group_pruning,candidate_node_id=n7,dimension_source_node_id=n1,target_source_node_id=n0,dimension_key_column=o_custkey,target_key_column=c_custkey,derived_min_key=none,derived_max_key=none,selected_row_groups_before=none,selected_row_groups_after=none,max_dimension_rows=1000000,applied=false,rejected_reason=DimensionPredicateAbsent|rule=dimension_derived_row_group_pruning,candidate_node_id=n10,dimension_source_node_id=n1,target_source_node_id=none,dimension_key_column=o_orderkey,target_key_column=none,derived_min_key=none,derived_max_key=none,selected_row_groups_before=none,selected_row_groups_after=none,max_dimension_rows=1000000,applied=false,rejected_reason=TargetKeyLineageUnavailable|rule=dimension_derived_row_group_pruning,candidate_node_id=n10,dimension_source_node_id=none,target_source_node_id=none,dimension_key_column=none,target_key_column=none,derived_min_key=none,derived_max_key=none,selected_row_groups_before=none,selected_row_groups_after=none,max_dimension_rows=1000000,applied=false,rejected_reason=DimensionKeyLineageUnavailable q18 0 rule=local_parquet_source_reuse,considered_count=1,applied_count=1,candidates=rule=local_parquet_source_reuse,source_node_ids=n2+n8,reuse_count=2,stable_file_ids=local:/lineitem.parquet:5023:mtime=,superset_read_columns=l_orderkey+l_quantity,branch_filter_count=0,max_variant_selected_bytes=524288000,applied=true,rejected_reason=none q18 0 rule=shared_subplan_materialization,considered_count=3,candidate_count=0,applied_count=0,candidates=none -q18 0 rule=dynamic_filter_route_discovery,considered_count=3,route_count=2,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n6,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n5,consumer_input_node_id=n1,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape|rule=dynamic_filter_route_discovery,candidate_node_id=n7,dynamic_filter_id=df0,producer_build_side=left,producer_input_node_id=n0,consumer_input_node_id=n6,consumer_kind=local_parquet_source,consumer_node_id=n1,consumer_source_node_id=n1,producer_key_column=c_custkey,consumer_key_column=o_custkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=224,total_source_bytes=224,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n9,dynamic_filter_id=df1,producer_build_side=right,producer_input_node_id=n8,consumer_input_node_id=n7,consumer_kind=local_parquet_source,consumer_node_id=n1,consumer_source_node_id=n1,producer_key_column=l_orderkey,consumer_key_column=o_orderkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=224,total_source_bytes=224,routed=true,rejected_reason=none +q18 0 rule=dynamic_filter_route_discovery,considered_count=3,route_count=3,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n6,dynamic_filter_id=df0,producer_build_side=right,producer_input_node_id=n5,consumer_input_node_id=n1,consumer_kind=local_parquet_source,consumer_node_id=n1,consumer_source_node_id=n1,producer_key_column=l_orderkey,consumer_key_column=o_orderkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=224,total_source_bytes=224,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n7,dynamic_filter_id=df1,producer_build_side=left,producer_input_node_id=n0,consumer_input_node_id=n6,consumer_kind=local_parquet_source,consumer_node_id=n1,consumer_source_node_id=n1,producer_key_column=c_custkey,consumer_key_column=o_custkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=224,total_source_bytes=224,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n9,dynamic_filter_id=df2,producer_build_side=right,producer_input_node_id=n8,consumer_input_node_id=n7,consumer_kind=local_parquet_source,consumer_node_id=n1,consumer_source_node_id=n1,producer_key_column=l_orderkey,consumer_key_column=o_orderkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=224,total_source_bytes=224,routed=true,rejected_reason=none q18 0 rule=aggregate_join_pushdown,considered_count=2,applied_count=1,candidates=rule=aggregate_join_pushdown,aggregate_node_id=n3,join_node_id=none,applied=false,rejected_reason=InputNotJoin|rule=aggregate_join_pushdown,aggregate_node_id=n10,join_node_id=n9,applied=true,rejected_reason=none q19 0 rule=projection_simplification,considered_count=6,applied_count=0,candidates=rule=projection_simplification,candidate_node_id=n8,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n6,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n2,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n10,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n8,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n3,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch q19 0 rule=join_build_side_selection,considered_count=1,applied_count=0,candidates=rule=join_build_side_selection,candidate_node_id=n7,before_build_side=right,after_build_side=none,left_estimated_rows=3,right_estimated_rows=3,left_estimated_bytes=379,right_estimated_bytes=245,left_advisory_rows=none,right_advisory_rows=none,left_advisory_bytes=none,right_advisory_bytes=none,small_build_side_bytes=67108864,small_build_side_rows=1000000,selection_basis=conservative_bytes,applied=false,rejected_reason=NoBetterBuildSide @@ -128,7 +128,7 @@ q20 0 rule=join_build_side_selection,considered_count=2,applied_count=2,candidat q20 0 rule=join_filter_reassociation,considered_count=1,applied_count=0,candidates=rule=join_filter_reassociation,candidate_node_id=n4,crossed_join_node_id=none,filtered_child_node_id=none,filtering_node_id=none,applied=false,rejected_reason=UnsupportedPreservedInputShape q20 0 rule=dimension_derived_row_group_pruning,considered_count=2,applied_count=0,candidates=rule=dimension_derived_row_group_pruning,candidate_node_id=n4,dimension_source_node_id=n0,target_source_node_id=n3,dimension_key_column=n_nationkey,target_key_column=s_nationkey,derived_min_key=none,derived_max_key=none,selected_row_groups_before=none,selected_row_groups_after=none,max_dimension_rows=1000000,applied=false,rejected_reason=DimensionPredicateAbsent|rule=dimension_derived_row_group_pruning,candidate_node_id=n4,dimension_source_node_id=n3,target_source_node_id=n0,dimension_key_column=s_nationkey,target_key_column=n_nationkey,derived_min_key=none,derived_max_key=none,selected_row_groups_before=none,selected_row_groups_after=none,max_dimension_rows=1000000,applied=false,rejected_reason=DimensionPredicateAbsent q20 0 rule=shared_subplan_materialization,considered_count=1,candidate_count=0,applied_count=0,candidates=none -q20 0 rule=dynamic_filter_route_discovery,considered_count=4,route_count=1,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n4,dynamic_filter_id=df0,producer_build_side=left,producer_input_node_id=n2,consumer_input_node_id=n3,consumer_kind=local_parquet_source,consumer_node_id=n3,consumer_source_node_id=n3,producer_key_column=n_nationkey,consumer_key_column=s_nationkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=324,total_source_bytes=324,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n9,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n8,consumer_input_node_id=n5,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape|rule=dynamic_filter_route_discovery,candidate_node_id=n14,dynamic_filter_id=none,producer_build_side=left,producer_input_node_id=n9,consumer_input_node_id=n13,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape|rule=dynamic_filter_route_discovery,candidate_node_id=n15,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n14,consumer_input_node_id=n4,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape +q20 0 rule=dynamic_filter_route_discovery,considered_count=4,route_count=3,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n4,dynamic_filter_id=df0,producer_build_side=left,producer_input_node_id=n2,consumer_input_node_id=n3,consumer_kind=local_parquet_source,consumer_node_id=n3,consumer_source_node_id=n3,producer_key_column=n_nationkey,consumer_key_column=s_nationkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=324,total_source_bytes=324,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n9,dynamic_filter_id=df1,producer_build_side=right,producer_input_node_id=n8,consumer_input_node_id=n5,consumer_kind=local_parquet_source,consumer_node_id=n5,consumer_source_node_id=n5,producer_key_column=p_partkey,consumer_key_column=ps_partkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=165,total_source_bytes=165,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n14,dynamic_filter_id=none,producer_build_side=left,producer_input_node_id=n9,consumer_input_node_id=n13,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape|rule=dynamic_filter_route_discovery,candidate_node_id=n15,dynamic_filter_id=df2,producer_build_side=right,producer_input_node_id=n14,consumer_input_node_id=n4,consumer_kind=local_parquet_source,consumer_node_id=n3,consumer_source_node_id=n3,producer_key_column=ps_suppkey,consumer_key_column=s_suppkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=324,total_source_bytes=324,routed=true,rejected_reason=none q20 0 rule=aggregate_join_pushdown,considered_count=1,applied_count=0,candidates=rule=aggregate_join_pushdown,aggregate_node_id=n12,join_node_id=none,applied=false,rejected_reason=InputNotJoin q21 0 rule=projection_simplification,considered_count=12,applied_count=1,candidates=rule=projection_simplification,candidate_node_id=n21,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n17,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n5,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n7,replacement_node_id=n3,applied=true,rejected_reason=none|rule=projection_simplification,candidate_node_id=n2,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n12,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n21,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n17,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n13,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n6,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n2,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n11,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch q21 0 rule=semijoin_reduction,considered_count=1,applied_count=0,candidates=rule=semijoin_reduction,candidate_node_id=n14,before_root_node_id=n14,after_root_node_id=none,crossed_join_node_ids=none,target_subtree_id=none,key_mappings=none,cost_threshold_multiplier=8,cost_gate=none,applied=false,rejected_reason=UnsupportedLeftShape @@ -138,7 +138,7 @@ q21 0 rule=dimension_derived_row_group_pruning,considered_count=6,applied_count= q21 0 rule=local_parquet_source_reuse,considered_count=1,applied_count=0,candidates=rule=local_parquet_source_reuse,source_node_ids=n4+n14,reuse_count=2,stable_file_ids=local:/lineitem.parquet:5023:mtime=,superset_read_columns=l_orderkey+l_suppkey+l_commitdate+l_receiptdate,branch_filter_count=0,max_variant_selected_bytes=524288000,applied=false,rejected_reason=UnsupportedSourceContract q21 0 rule=count_non_null_column_pruning,considered_count=1,applied_count=0,candidates=rule=count_non_null_column_pruning,aggregate_node_id=n18,column=__algeon_native_aggregate_arg_0,source_node_id=none,file_column_index=none,applied=false,rejected_reason=UnsupportedInputLineage q21 0 rule=shared_subplan_materialization,considered_count=1,candidate_count=0,applied_count=0,candidates=none -q21 0 rule=dynamic_filter_route_discovery,considered_count=5,route_count=2,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n7,dynamic_filter_id=df0,producer_build_side=right,producer_input_node_id=n6,consumer_input_node_id=n3,consumer_kind=local_parquet_source,consumer_node_id=n3,consumer_source_node_id=n3,producer_key_column=l_suppkey,consumer_key_column=s_suppkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=221,total_source_bytes=221,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n8,dynamic_filter_id=df1,producer_build_side=right,producer_input_node_id=n2,consumer_input_node_id=n7,consumer_kind=local_parquet_source,consumer_node_id=n3,consumer_source_node_id=n3,producer_key_column=n_nationkey,consumer_key_column=s_nationkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=221,total_source_bytes=221,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n12,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n11,consumer_input_node_id=n8,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=n4,producer_key_column=o_orderkey,consumer_key_column=l_orderkey,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=212,total_source_bytes=212,routed=false,rejected_reason=ConsumerSourceReusedInPlan|rule=dynamic_filter_route_discovery,candidate_node_id=n15,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n14,consumer_input_node_id=n13,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape|rule=dynamic_filter_route_discovery,candidate_node_id=n16,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n6,consumer_input_node_id=n15,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape +q21 0 rule=dynamic_filter_route_discovery,considered_count=5,route_count=2,candidates=rule=dynamic_filter_route_discovery,candidate_node_id=n7,dynamic_filter_id=df0,producer_build_side=right,producer_input_node_id=n6,consumer_input_node_id=n3,consumer_kind=local_parquet_source,consumer_node_id=n3,consumer_source_node_id=n3,producer_key_column=l_suppkey,consumer_key_column=s_suppkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=221,total_source_bytes=221,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n8,dynamic_filter_id=df1,producer_build_side=right,producer_input_node_id=n2,consumer_input_node_id=n7,consumer_kind=local_parquet_source,consumer_node_id=n3,consumer_source_node_id=n3,producer_key_column=n_nationkey,consumer_key_column=s_nationkey,filter_kinds=range+small_in_list+hash_set_contains,consumer_lowering=reader_predicate,consumer_reduction_stage=source_reader,selected_source_bytes=221,total_source_bytes=221,routed=true,rejected_reason=none|rule=dynamic_filter_route_discovery,candidate_node_id=n12,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n11,consumer_input_node_id=n8,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=n4,producer_key_column=o_orderkey,consumer_key_column=l_orderkey,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=212,total_source_bytes=212,routed=false,rejected_reason=ConsumerSourceReusedInPlan|rule=dynamic_filter_route_discovery,candidate_node_id=n15,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n14,consumer_input_node_id=n13,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=n4,producer_key_column=l_orderkey,consumer_key_column=l_orderkey,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=212,total_source_bytes=212,routed=false,rejected_reason=ConsumerSourceReusedInPlan|rule=dynamic_filter_route_discovery,candidate_node_id=n16,dynamic_filter_id=none,producer_build_side=right,producer_input_node_id=n6,consumer_input_node_id=n15,consumer_kind=none,consumer_node_id=none,consumer_source_node_id=none,producer_key_column=none,consumer_key_column=none,filter_kinds=none,consumer_lowering=none,consumer_reduction_stage=none,selected_source_bytes=none,total_source_bytes=none,routed=false,rejected_reason=UnsupportedJoinShape q21 0 rule=aggregate_join_pushdown,considered_count=1,applied_count=0,candidates=rule=aggregate_join_pushdown,aggregate_node_id=n18,join_node_id=n17,applied=false,rejected_reason=JoinResidualUnsupported q22 0 rule=projection_simplification,considered_count=8,applied_count=0,candidates=rule=projection_simplification,candidate_node_id=n13,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n10,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n9,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n4,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n13,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n10,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n9,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch|rule=projection_simplification,candidate_node_id=n4,replacement_node_id=none,applied=false,rejected_reason=SchemaMismatch q22 0 rule=join_filter_reassociation,considered_count=1,applied_count=0,candidates=rule=join_filter_reassociation,candidate_node_id=n8,crossed_join_node_id=none,filtered_child_node_id=none,filtering_node_id=none,applied=false,rejected_reason=NoPushableResidualPredicate diff --git a/crates/algeon-datafusion/src/native_report/report_execution_tests.rs b/crates/algeon-datafusion/src/native_report/report_execution_tests.rs index dafafe758..eb5c1a8b8 100644 --- a/crates/algeon-datafusion/src/native_report/report_execution_tests.rs +++ b/crates/algeon-datafusion/src/native_report/report_execution_tests.rs @@ -827,7 +827,7 @@ fn native_datafusion_planning_report_projects_native_optimizer_diagnostics() { ";rule=shared_resident_extraction,considered_count=6,extracted_count=0,rejected_count=0,candidates=" )); assert!(optimizer_cell.contains("decision=Skipped,reason=NotRepeated,resident_input_name=none")); - assert!(optimizer_cell.contains(";rule=dynamic_filter_route_discovery,considered_count=2,route_count=1")); + assert!(optimizer_cell.contains(";rule=dynamic_filter_route_discovery,considered_count=2,route_count=2")); assert!(optimizer_cell.contains(";rule=aggregate_join_pushdown,considered_count=1,applied_count=0")); assert!(text.contains("\tnative_optimizer\t")); assert!(text.contains("rule=semijoin_reduction,considered_count=1,applied_count=1")); diff --git a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/fixtures.rs b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/fixtures.rs index e33ced668..66b390e9b 100644 --- a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/fixtures.rs +++ b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/fixtures.rs @@ -243,6 +243,9 @@ pub(super) struct SharedPartsuppDynamicFilterPlan { pub(super) struct SharedPartsuppDynamicFilterPlanOptions { pub(super) extra_unconstrained_partsupp_consumer: bool, pub(super) final_join_type: JoinType, + /// Puts the aggregate on the final join's left input and the part side on + /// its right input. + pub(super) final_join_swaps_inputs: bool, pub(super) final_join_null_equality: JoinNullEquality, pub(super) derived_aggregate_group_key: bool, pub(super) filter_part_build_keyset: bool, @@ -255,6 +258,7 @@ impl Default for SharedPartsuppDynamicFilterPlanOptions { Self { extra_unconstrained_partsupp_consumer: false, final_join_type: JoinType::Inner, + final_join_swaps_inputs: false, final_join_null_equality: JoinNullEquality::NullsDoNotMatch, derived_aggregate_group_key: false, filter_part_build_keyset: false, @@ -360,29 +364,50 @@ pub(super) fn q02_like_shared_partsupp_dynamic_filter_plan( } else { part_partsupp_join }; + let final_join_spec = if options.final_join_swaps_inputs { + JoinSpec::new( + options.final_join_type, + vec![ + JoinKey::new( + ColumnRef::input(0, aggregate_key_name, aggregate_key_name), + ColumnRef::input(1, "p_partkey", "p_partkey"), + ), + JoinKey::new( + ColumnRef::input(0, "min_supplycost", "min_supplycost"), + ColumnRef::input(1, "ps_supplycost", "ps_supplycost"), + ), + ], + ) + .with_output_projection(vec![ + JoinOutputColumn::left(ColumnRef::input(0, aggregate_key_name, aggregate_key_name), "p_partkey"), + JoinOutputColumn::left(ColumnRef::input(0, "min_supplycost", "min_supplycost"), "ps_supplycost"), + ]) + } else { + JoinSpec::new( + options.final_join_type, + vec![ + JoinKey::new( + ColumnRef::input(0, "p_partkey", "p_partkey"), + ColumnRef::input(1, aggregate_key_name, aggregate_key_name), + ), + JoinKey::new( + ColumnRef::input(0, "ps_supplycost", "ps_supplycost"), + ColumnRef::input(1, "min_supplycost", "min_supplycost"), + ), + ], + ) + .with_output_projection(vec![ + JoinOutputColumn::left(ColumnRef::input(0, "p_partkey", "p_partkey"), "p_partkey"), + JoinOutputColumn::left(ColumnRef::input(0, "ps_supplycost", "ps_supplycost"), "ps_supplycost"), + ]) + }; + let (final_join_left, final_join_right) = + if options.final_join_swaps_inputs { (partsupp_min, final_join_left) } else { (final_join_left, partsupp_min) }; let final_join = builder .add_join( final_join_left, - partsupp_min, - JoinSpec::new( - options.final_join_type, - vec![ - JoinKey::new( - ColumnRef::input(0, "p_partkey", "p_partkey"), - ColumnRef::input(1, aggregate_key_name, aggregate_key_name), - ), - JoinKey::new( - ColumnRef::input(0, "ps_supplycost", "ps_supplycost"), - ColumnRef::input(1, "min_supplycost", "min_supplycost"), - ), - ], - ) - .with_null_equality(options.final_join_null_equality) - .with_output_projection(vec![ - JoinOutputColumn::left(ColumnRef::input(0, "p_partkey", "p_partkey"), "p_partkey"), - JoinOutputColumn::left(ColumnRef::input(0, "ps_supplycost", "ps_supplycost"), "ps_supplycost"), - ]) - .with_build_side(JoinSide::Right), + final_join_right, + final_join_spec.with_null_equality(options.final_join_null_equality).with_build_side(JoinSide::Right), Arc::clone(&result_schema), ) .unwrap(); diff --git a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/shared_sources.rs b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/shared_sources.rs index 36ac1ab0c..97bc95318 100644 --- a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/shared_sources.rs +++ b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/shared_sources.rs @@ -318,9 +318,20 @@ fn native_optimizer_rejects_join_probe_dynamic_filter_for_reused_local_source() assert!(routes.is_empty()); } +// The final join drops aggregate rows whose key matches no part key, whether it +// is an inner join or a semi join, so both constrain the shared partsupp path. #[test] fn native_optimizer_routes_q02_like_shared_partsupp_source_when_all_key_consumers_are_constrained() { - let fixture = q02_like_shared_partsupp_dynamic_filter_plan(SharedPartsuppDynamicFilterPlanOptions::default()); + for final_join_type in [JoinType::Inner, JoinType::LeftSemi] { + assert_q02_like_shared_partsupp_source_routed(final_join_type); + } +} + +fn assert_q02_like_shared_partsupp_source_routed(final_join_type: JoinType) { + let fixture = q02_like_shared_partsupp_dynamic_filter_plan(SharedPartsuppDynamicFilterPlanOptions { + final_join_type, + ..SharedPartsuppDynamicFilterPlanOptions::default() + }); let options = NativeOptimizerOptions { dynamic_filter_route_discovery: DynamicFilterRouteDiscoveryOptions::default(), join_build_side_selection: JoinBuildSideSelectionOptions::disabled(), @@ -412,8 +423,12 @@ fn t_q02_like_cases() { native_optimizer_rejects_q02_like_shared_source_filter_with_unconstrained_extra_consumer, ), ( - "optimizer_contracts::dynamic_filter::shared_sources::native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_outer_joined", - native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_outer_joined, + "optimizer_contracts::dynamic_filter::shared_sources::native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_full_outer_joined", + native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_full_outer_joined, + ), + ( + "optimizer_contracts::dynamic_filter::shared_sources::native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_left_outer_preserved", + native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_left_outer_preserved, ), ( "optimizer_contracts::dynamic_filter::shared_sources::native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_uses_null_equal_join", @@ -434,11 +449,25 @@ fn t_q02_like_cases() { ]); } -fn native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_outer_joined() { - let fixture = q02_like_shared_partsupp_dynamic_filter_plan(SharedPartsuppDynamicFilterPlanOptions { +fn native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_full_outer_joined() { + assert_q02_like_shared_partsupp_source_not_routed(SharedPartsuppDynamicFilterPlanOptions { + final_join_type: JoinType::Full, + ..SharedPartsuppDynamicFilterPlanOptions::default() + }); +} + +/// A left outer join keeps every left row, so the aggregate path on its left +/// input is not constrained by the part keys. +fn native_optimizer_rejects_q02_like_shared_source_filter_when_aggregate_key_is_left_outer_preserved() { + assert_q02_like_shared_partsupp_source_not_routed(SharedPartsuppDynamicFilterPlanOptions { final_join_type: JoinType::Left, + final_join_swaps_inputs: true, ..SharedPartsuppDynamicFilterPlanOptions::default() }); +} + +fn assert_q02_like_shared_partsupp_source_not_routed(plan_options: SharedPartsuppDynamicFilterPlanOptions) { + let fixture = q02_like_shared_partsupp_dynamic_filter_plan(plan_options); let options = NativeOptimizerOptions { dynamic_filter_route_discovery: DynamicFilterRouteDiscoveryOptions::default(), join_build_side_selection: JoinBuildSideSelectionOptions::disabled(), diff --git a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/source_routes.rs b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/source_routes.rs index bae43e92d..dea21a8d7 100644 --- a/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/source_routes.rs +++ b/crates/algeon-query-engine/tests/native_contract_tests_cpu/optimizer_contracts/dynamic_filter/source_routes.rs @@ -1,3 +1,5 @@ +use query_planning::optimizer::NativeOptimizerResult; + use super::*; #[test] @@ -446,3 +448,107 @@ fn native_optimizer_discovers_dynamic_filter_route_for_inner_keyed_residual_join assert_eq!(routes[0].producer_columns, vec!["right_key".to_string()]); assert_eq!(routes[0].consumer_columns, vec!["left_key".to_string()]); } + +/// A keyed `left_key = right_key` join whose build side is the small input, so +/// the size heuristics never decline the probe source. Returns the plan with the +/// left and right source ids. +fn build_small_left_only_output_join_plan(join_type: JoinType, build_side: JoinSide) -> (QueryPlan, NodeId, NodeId) { + let (small, large) = ((16 * 1024 * 1024, 100_000), (512 * 1024 * 1024, 10_000_000)); + let ((left_bytes, left_rows), (right_bytes, right_rows)) = match build_side { + JoinSide::Left => (small, large), + JoinSide::Right => (large, small), + }; + let mut fields = vec![Field::new("left_key", DataType::Int32, false)]; + if let Some(mark_column) = join_type.mark_column() { + fields.push(Field::new(mark_column, DataType::Boolean, false)); + } + let mut builder = QueryPlanBuilder::new(); + let left = builder.add_source(estimated_local_source("left", "left_key", left_bytes, left_rows)); + let right = builder.add_source(estimated_local_source("right", "right_key", right_bytes, right_rows)); + let join = builder + .add_join( + left, + right, + JoinSpec::new( + join_type, + vec![JoinKey::new( + ColumnRef::input(0, "left_key", "left_key"), + ColumnRef::input(1, "right_key", "right_key"), + )], + ) + .with_output_projection(vec![JoinOutputColumn::left( + ColumnRef::input(0, "left_key", "left_key"), + "left_key", + )]) + .with_build_side(build_side), + Arc::new(Schema::new(fields)), + ) + .unwrap(); + (builder.finish_gpu_resident(join).unwrap(), left, right) +} + +fn discover_dynamic_filter_routes(plan: QueryPlan) -> NativeOptimizerResult { + let options = NativeOptimizerOptions { + dynamic_filter_route_discovery: DynamicFilterRouteDiscoveryOptions::default(), + ..NativeOptimizerOptions::disabled() + }; + optimize_query_plan( + plan, + &options, + rapids_analysis_context(PlanningPolicy::default()), + local_dimension_pruning_resolver(), + ) + .unwrap() +} + +fn assert_single_build_route_to_source(result: &NativeOptimizerResult, build_side: JoinSide, consumer_source: NodeId) { + let routes = result.plan.dynamic_filter_routes(); + assert_eq!(routes.len(), 1, "{routes:?}"); + assert!( + matches!(routes[0].producer, DynamicFilterProducer::HashJoinBuild { build_side: side, .. } if side == build_side) + ); + assert!(matches!( + routes[0].consumer, + DynamicFilterConsumer::LocalParquetSource { source_node_id } if source_node_id == consumer_source + )); + assert_eq!(routes[0].consumer_lowering, DynamicFilterConsumerLowering::ReaderPredicate); +} + +fn assert_join_shape_rejected(result: &NativeOptimizerResult) { + let discovery = &result.diagnostics.dynamic_filter_route_discovery; + assert_eq!(discovery.considered_count, 1); + assert_eq!(discovery.route_count, 0); + assert_eq!(discovery.candidates[0].rejected_reason, Some(DynamicFilterRouteRejectedReason::UnsupportedJoinShape)); + assert!(result.plan.dynamic_filter_routes().is_empty()); +} + +// A semi join drops the rows of either input that match nothing on the other, +// so the build keys may prune the probe source whichever side builds. A mark +// join keeps every left row and marks it, so its left input never consumes. +#[test] +fn native_optimizer_routes_semi_join_build_keys_to_probe_source_but_not_into_mark_join_left_input() { + let (plan, left, _) = build_small_left_only_output_join_plan(JoinType::LeftSemi, JoinSide::Right); + assert_single_build_route_to_source(&discover_dynamic_filter_routes(plan), JoinSide::Right, left); + + let (plan, _, right) = build_small_left_only_output_join_plan(JoinType::LeftSemi, JoinSide::Left); + assert_single_build_route_to_source(&discover_dynamic_filter_routes(plan), JoinSide::Left, right); + + let (plan, _, _) = + build_small_left_only_output_join_plan(JoinType::LeftMark { mark_column: "mark".to_string() }, JoinSide::Right); + assert_join_shape_rejected(&discover_dynamic_filter_routes(plan)); +} + +// An anti join keeps exactly the left rows that match nothing, so only its +// right input may lose unmatched rows. Under `NOT IN` semantics a right null or +// an empty right input changes the result, so neither input consumes. +#[test] +fn native_optimizer_routes_anti_join_build_keys_only_into_its_right_input() { + let (plan, _, right) = build_small_left_only_output_join_plan(JoinType::LEFT_ANTI, JoinSide::Left); + assert_single_build_route_to_source(&discover_dynamic_filter_routes(plan), JoinSide::Left, right); + + let (plan, _, _) = build_small_left_only_output_join_plan(JoinType::LEFT_ANTI, JoinSide::Right); + assert_join_shape_rejected(&discover_dynamic_filter_routes(plan)); + + let (plan, _, _) = build_small_left_only_output_join_plan(JoinType::LeftAnti { null_aware: true }, JoinSide::Left); + assert_join_shape_rejected(&discover_dynamic_filter_routes(plan)); +} diff --git a/crates/algeon-query-engine/tests/native_execution_tests_gpu/aggregate_join_pushdown.rs b/crates/algeon-query-engine/tests/native_execution_tests_gpu/aggregate_join_pushdown.rs index bd6f3439c..97e18447e 100644 --- a/crates/algeon-query-engine/tests/native_execution_tests_gpu/aggregate_join_pushdown.rs +++ b/crates/algeon-query-engine/tests/native_execution_tests_gpu/aggregate_join_pushdown.rs @@ -3,6 +3,7 @@ use std::{path::Path, sync::Arc}; use arrow::array::{Int32Array, Int64Array}; use query_engine::{dimension_pruning::local_dimension_pruning_resolver, rapids_backend::rapids_analysis_context}; use query_planning::{ + dynamic_filter::DynamicFilterConsumer, expr::ColumnRef, optimizer::{AggregateJoinPushdownRejectedReason, NativeOptimizerChange, optimize_query_plan}, plan::QueryPlan, @@ -413,7 +414,15 @@ fn bounded_native_query_filters_left_semi_groups_before_the_final_aggregate() { pushdown.candidates[0].rejected_reason, Some(AggregateJoinPushdownRejectedReason::LeftSemiAlreadyFiltersFinalGroups) ); - assert_eq!(optimizer_result.plan, plan, "the semi filter must stay directly below the grouped aggregate"); + assert_eq!( + (optimizer_result.plan.nodes(), optimizer_result.plan.sink()), + (plan.nodes(), plan.sink()), + "the semi filter must stay directly below the grouped aggregate" + ); + let routes = optimizer_result.plan.dynamic_filter_routes(); + assert_eq!(routes.len(), 1, "the semi join build keys should reduce the detail source: {routes:?}"); + let detail_source = plan.nodes()[0].id; + assert_eq!(routes[0].consumer, DynamicFilterConsumer::LocalParquetSource { source_node_id: detail_source }); let context = backend_context_with_grant(SEMI_GROUP_FILTER_GRANT_BYTES); let (result, _metrics, peak_bytes) = execute_admitted_with_peak(&context, &optimizer_result.plan); diff --git a/crates/algeon-query-engine/tests/native_execution_tests_gpu/bounded_joins/root_joins/hash/dynamic_filter.rs b/crates/algeon-query-engine/tests/native_execution_tests_gpu/bounded_joins/root_joins/hash/dynamic_filter.rs index 4b18c38cd..65a7a44be 100644 --- a/crates/algeon-query-engine/tests/native_execution_tests_gpu/bounded_joins/root_joins/hash/dynamic_filter.rs +++ b/crates/algeon-query-engine/tests/native_execution_tests_gpu/bounded_joins/root_joins/hash/dynamic_filter.rs @@ -60,6 +60,94 @@ fn bounded_native_query_applies_hash_join_build_dynamic_filter_to_probe_source() fs::remove_dir_all(dir).unwrap(); } +// The `x IN (SELECT ... WHERE ...)` shape: a left semi join that builds a +// filtered right input. The build keys prune the probe source's row groups, +// and the result matches the same plan executed without routes. +fn bounded_native_query_applies_semi_join_build_dynamic_filter_to_probe_source() { + let runtime = backend_runtime(); + let dir = unique_temp_dir("query-engine-bounded-root-semi-join-dynamic-filter"); + let left_path = dir.join("lineitem.parquet"); + let right_path = dir.join("orders.parquet"); + let left_schema = write_dynamic_filter_probe_parquet(&left_path); + write_join_right_keys_parquet(&runtime, &right_path, &[1, 2, 2, 3]); + let right_schema = Arc::new(Schema::new(vec![Field::new("order_id", DataType::Int32, false)])); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("order_id", DataType::Int32, false), + Field::new("qty", DataType::Int32, false), + ])); + + let mut builder = QueryPlanBuilder::new(); + let left = builder.add_source(SourceNode::LocalParquet(LocalParquetSource::from_files( + vec![ + LocalParquetFile::new(&left_path) + .with_file_size_bytes(fs::metadata(&left_path).unwrap().len()) + .with_row_group_rows(vec![2, 2, 2]), + ], + Arc::clone(&left_schema), + vec!["order_id".to_string(), "qty".to_string()], + ))); + let right = builder.add_source(SourceNode::LocalParquet(LocalParquetSource::from_files( + vec![LocalParquetFile::new(&right_path).with_file_size_bytes(fs::metadata(&right_path).unwrap().len())], + Arc::clone(&right_schema), + vec!["order_id".to_string()], + ))); + let right = builder + .add_filter(right, Expr::binary(Expr::column("order_id"), BinaryOp::Eq, Expr::literal(LiteralValue::Int32(2)))) + .unwrap(); + let join = builder + .add_join( + left, + right, + JoinSpec::new(JoinType::LeftSemi, vec![join_order_id_key()]) + .with_output_projection(vec![ + JoinOutputColumn::left(ColumnRef::input(0, "order_id", "order_id"), "order_id"), + JoinOutputColumn::left(ColumnRef::input(0, "qty", "qty"), "qty"), + ]) + .with_build_side(JoinSide::Right), + Arc::clone(&output_schema), + ) + .unwrap(); + let plan = builder.finish_with_host_sink(join).unwrap(); + let optimize = |route_discovery| { + let options = NativeOptimizerOptions { + dynamic_filter_route_discovery: route_discovery, + join_build_side_selection: JoinBuildSideSelectionOptions::disabled(), + ..NativeOptimizerOptions::default() + }; + optimize_query_plan( + plan.clone(), + &options, + rapids_analysis_context(PlanningPolicy::default()), + local_dimension_pruning_resolver(), + ) + .unwrap() + }; + let routed = optimize(DynamicFilterRouteDiscoveryOptions::default()); + let unrouted = optimize(DynamicFilterRouteDiscoveryOptions::disabled()); + assert_eq!(routed.diagnostics.dynamic_filter_route_discovery.route_count, 1); + assert!(unrouted.plan.dynamic_filter_routes().is_empty()); + let context = backend_context_with_grant(1024 * 1024); + + let mut outputs = Vec::new(); + for (optimized, row_groups_committed) in [(&routed, 2), (&unrouted, 4)] { + let result = execute_query_plan(&context, &optimized.plan).unwrap(); + let output = result.output().as_host_arrow().unwrap(); + let mut rows = int32_array(output, 0) + .values() + .iter() + .copied() + .zip(int32_array(output, 1).values().iter().copied()) + .collect::>(); + rows.sort_unstable(); + assert_eq!(result.metrics().row_groups_committed, row_groups_committed); + outputs.push(rows); + } + assert_eq!(outputs[0], vec![(2, 20), (2, 21)]); + assert_eq!(outputs[0], outputs[1]); + + fs::remove_dir_all(dir).unwrap(); +} + #[test] fn dynamic_filter_cases_policy_1_1() { run_named_cases(&[ @@ -67,6 +155,10 @@ fn dynamic_filter_cases_policy_1_1() { "bounded_joins::root_joins::hash::dynamic_filter::bounded_native_query_applies_hash_join_build_dynamic_filter_to_probe_source", bounded_native_query_applies_hash_join_build_dynamic_filter_to_probe_source, ), + ( + "bounded_joins::root_joins::hash::dynamic_filter::bounded_native_query_applies_semi_join_build_dynamic_filter_to_probe_source", + bounded_native_query_applies_semi_join_build_dynamic_filter_to_probe_source, + ), ( "bounded_joins::root_joins::hash::dynamic_filter::bounded_native_query_applies_dynamic_filter_to_source_under_full_sort_probe", bounded_native_query_applies_dynamic_filter_to_source_under_full_sort_probe, diff --git a/crates/algeon-query-planning/src/dynamic_filter.rs b/crates/algeon-query-planning/src/dynamic_filter.rs index bb87edc30..919b8978b 100644 --- a/crates/algeon-query-planning/src/dynamic_filter.rs +++ b/crates/algeon-query-planning/src/dynamic_filter.rs @@ -13,7 +13,8 @@ //! is complete before the consumer starts, and a consumer whose route produced //! no payload runs unfiltered. Correctness never depends on a payload: each one //! covers every non-null producer key, so applying it removes only rows that -//! could not have matched. +//! could not have matched, and discovery places a consumer only on a join input +//! whose unmatched rows the join discards. use std::{collections::BTreeMap, fmt}; diff --git a/crates/algeon-query-planning/src/optimizer/dynamic_filter.rs b/crates/algeon-query-planning/src/optimizer/dynamic_filter.rs index d6608657e..a044eaafa 100644 --- a/crates/algeon-query-planning/src/optimizer/dynamic_filter.rs +++ b/crates/algeon-query-planning/src/optimizer/dynamic_filter.rs @@ -1,24 +1,27 @@ //! Dynamic-filter route discovery. //! -//! The rule matches an inner single-key hash join under `NullsDoNotMatch` whose -//! build key traces to one source column, or to a grouping-only aggregate's key -//! set, and whose probe key traces to a local Parquet source column, to a filter -//! above it, or to the join's own probe input. In the rewrite, no node changes. -//! The pass replaces the plan's [`DynamicFilterRoute`] annotations with routes -//! that name a producer (the hash build or the aggregate key set) and a consumer; -//! execution applies the produced key payload at the consumer before rows reach -//! the join. A route changes no result because such a join emits a probe row only -//! when its key appears in the build, so removing probe rows whose key cannot -//! appear removes nothing the join would keep. The consumer must not sit upstream -//! of the producer, a source read by other consumers must be proven constrained -//! by the same key set (`keyset_constraints`), and the route must not close a -//! dependency cycle with earlier routes (`cycle`). +//! The rule matches a single-key hash join under `NullsDoNotMatch` that discards +//! probe rows matching no build row (an inner or semi join, or an anti join that +//! builds its left input), whose build key traces to one source column, or to a +//! grouping-only aggregate's key set, and whose probe key traces to a local +//! Parquet source column, to a filter above it, or to the join's own probe +//! input. In the rewrite, no node changes. The pass replaces the plan's +//! [`DynamicFilterRoute`] annotations with routes that name a producer (the hash +//! build or the aggregate key set) and a consumer; execution applies the +//! produced key payload at the consumer before rows reach the join. A route +//! changes no result because such a join keeps nothing of a probe row whose key +//! does not appear in the build. An outer join's preserved input, an anti join's +//! left input and a mark join's left input keep unmatched rows and never +//! consume. The consumer must not sit upstream of the producer, a source read by +//! other consumers must be proven constrained by the same key set +//! (`keyset_constraints`), and the route must not close a dependency cycle with +//! earlier routes (`cycle`). //! //! Discovery is the last step of the pass table because `QueryPlan::rewrite` //! drops a route whose endpoint node changed; a rule running afterwards would -//! lose routes. Equalities in an inner join's residual predicate add further -//! candidates (`residual_equality`). Rejection reasons and the candidate -//! diagnostics fields are report contracts. +//! lose routes. Equalities in a join's residual predicate add further candidates +//! for the inputs whose unmatched rows the join discards (`residual_equality`). +//! Rejection reasons and the candidate diagnostics fields are report contracts. use std::collections::BTreeSet; @@ -153,7 +156,10 @@ pub struct DynamicFilterRouteCandidateDiagnostics { pub enum DynamicFilterRouteRejectedReason { /// Pass was disabled. PassDisabled, - /// Join shape cannot produce a route yet. + /// The join keeps probe rows that match no build row (an outer join's + /// preserved input, an anti or mark join's left input), matches null keys, + /// is null-aware, has other than one key, has no output columns, or reuses + /// a prebuilt build artifact. UnsupportedJoinShape, /// Build-side key lineage could not be resolved to one source column. ProducerKeyLineageUnavailable, @@ -202,17 +208,40 @@ fn producer_input_is_literal_batch_exact_keyset(plan: &QueryPlan, node_id: NodeI } } -/// The join shape whose probe rows a build key payload may remove. Under -/// null-equal matching a null probe key matches null build keys, and no payload -/// kind represents null, so only `NullsDoNotMatch` joins qualify. +/// The join shape whose probe rows a build key payload may remove: the join +/// discards probe rows that match no build row, and its single key never +/// matches null. Under null-equal matching a null probe key matches null build +/// keys, and no payload kind represents null, so only `NullsDoNotMatch` joins +/// qualify. fn join_shape_supported(join: &JoinNode) -> bool { - join.spec.join_type == JoinType::Inner + let probe_side = match join.spec.build_side { + JoinSide::Left => JoinSide::Right, + JoinSide::Right => JoinSide::Left, + }; + join_type_discards_unmatched_rows(&join.spec.join_type, probe_side) && join.spec.keys.len() == 1 && join.spec.null_equality == JoinNullEquality::NullsDoNotMatch && !join.spec.output_projection.is_empty() && join.spec.build_artifact_id.is_none() } +/// Whether a join of `join_type` drops every row of its `side` input that +/// matches no row of the other input, so that removing such rows before the +/// join leaves its output unchanged. Inner and semi joins drop unmatched rows +/// of both inputs. Left outer, anti and mark joins keep every left row, as +/// output or as its mark, and a right row that matches nothing contributes to +/// none of them. A null-aware anti join reads the right input's nulls and its +/// emptiness, so neither of its inputs qualifies. +fn join_type_discards_unmatched_rows(join_type: &JoinType, side: JoinSide) -> bool { + match join_type { + JoinType::Inner | JoinType::LeftSemi => true, + JoinType::Left | JoinType::LeftAnti { null_aware: false } | JoinType::LeftMark { .. } => { + side == JoinSide::Right + }, + JoinType::LeftAnti { null_aware: true } | JoinType::Full | JoinType::Cross => false, + } +} + fn subtree_contains_node(plan: &QueryPlan, root: NodeId, target: NodeId) -> bool { let mut stack = vec![root]; let mut visited = BTreeSet::new(); diff --git a/crates/algeon-query-planning/src/optimizer/dynamic_filter/keyset_constraints.rs b/crates/algeon-query-planning/src/optimizer/dynamic_filter/keyset_constraints.rs index cb83eea13..869cdc1fc 100644 --- a/crates/algeon-query-planning/src/optimizer/dynamic_filter/keyset_constraints.rs +++ b/crates/algeon-query-planning/src/optimizer/dynamic_filter/keyset_constraints.rs @@ -4,17 +4,19 @@ //! A local Parquet source with several parents may take a source-level dynamic //! filter only when every path upward from it, following the key column through //! filters, renaming projections, sorts, grouping-only aggregates and join output -//! projections, reaches an inner `NullsDoNotMatch` join that equates the column -//! with the same producer key set. Rows the filter removes then match nothing on -//! any path and contribute to no output. A limit or window ends a path +//! projections, reaches a `NullsDoNotMatch` join that equates the column with the +//! same producer key set and discards the column's input rows that match +//! nothing. Rows the filter removes then match nothing on any path and +//! contribute to no output. A limit or window ends a path //! negatively because it selects rows by position, so removing rows below it //! changes its output even when a join above would constrain the key. use std::collections::BTreeSet; use super::{ - AggregateGroupKeyLineage, AggregateNode, JoinNode, JoinNullEquality, JoinSide, JoinType, NodeId, PlanNode, - QueryPlan, aggregate_group_key_output_name, direct_projection_column_name, subtree_contains_node, unique_lineage, + AggregateGroupKeyLineage, AggregateNode, JoinNode, JoinNullEquality, JoinSide, NodeId, PlanNode, QueryPlan, + aggregate_group_key_output_name, direct_projection_column_name, join_type_discards_unmatched_rows, + subtree_contains_node, unique_lineage, }; /// Whether every consumer of `source_column` at `source_node_id` is constrained @@ -218,10 +220,10 @@ fn join_constrains_or_carries_column( ) } -/// Whether `join` equates `column` on `child_side` with the producer key: the -/// other input must contain the producer key-set node and the other key's -/// unique lineage must be the producer's source column, so the two sides are -/// filtered by the same values. +/// Whether `join` equates `column` on `child_side` with the producer key and +/// drops `child_side` rows that match nothing: the other input must contain the +/// producer key-set node and the other key's unique lineage must be the +/// producer's source column, so the two sides are filtered by the same values. #[expect(clippy::too_many_arguments, reason = "planning inputs are distinct proof and constraint values")] fn join_directly_constrains_column( plan: &QueryPlan, @@ -233,7 +235,9 @@ fn join_directly_constrains_column( producer_source_node_id: NodeId, producer_source_column: &str, ) -> bool { - if join.spec.join_type != JoinType::Inner || join.spec.null_equality != JoinNullEquality::NullsDoNotMatch { + if !join_type_discards_unmatched_rows(&join.spec.join_type, child_side) + || join.spec.null_equality != JoinNullEquality::NullsDoNotMatch + { return false; } join.spec.keys.iter().any(|key| { diff --git a/crates/algeon-query-planning/src/optimizer/dynamic_filter/residual_equality.rs b/crates/algeon-query-planning/src/optimizer/dynamic_filter/residual_equality.rs index 062bf88b0..a75146e8a 100644 --- a/crates/algeon-query-planning/src/optimizer/dynamic_filter/residual_equality.rs +++ b/crates/algeon-query-planning/src/optimizer/dynamic_filter/residual_equality.rs @@ -1,10 +1,12 @@ -//! Route candidates from column equalities in an inner join's residual. +//! Route candidates from column equalities in a join's residual. //! //! A mixed join evaluates `left.a = right.b` inside its residual. It does not use //! that equality as a hash key. The equality constrains the consumer side as a //! key would, so a supported hash build below the producer side that carries the //! same source column can supply a payload to the consumer's local Parquet -//! source. The residual join is never its own producer. Candidates are reported +//! source. Only an input whose unmatched rows the join discards consumes, as for +//! key routes; SQL equality never holds for null, whatever the key null +//! equality. The residual join is never its own producer. Candidates are reported //! only when the producer side is already filtered or is smaller than the //! consumer by rows or bytes; otherwise the producer-before-consumer dependency //! would cost more than the rows it removes. Only equalities that are conjuncts @@ -13,9 +15,10 @@ use super::{ AggregateGroupKeyLineage, BinaryOp, ColumnQualifier, DynamicFilterExecutionOptions, DynamicFilterProducer, Expr, - JoinNode, JoinSide, JoinType, LocalParquetSource, NodeId, PlanNode, QueryPlan, SourceColumnLineage, SourceNode, + JoinNode, JoinSide, LocalParquetSource, NodeId, PlanNode, QueryPlan, SourceColumnLineage, SourceNode, dedup_lineages, estimated_output_bytes, estimated_output_rows, join_shape_supported, join_side_index, - lineages_for_column, local_parquet_source, subtree_contains_node, unique_lineage, + join_type_discards_unmatched_rows, lineages_for_column, local_parquet_source, subtree_contains_node, + unique_lineage, }; /// One producer and consumer pairing derived from a residual equality, before @@ -37,9 +40,6 @@ pub(super) fn residual_equality_route_candidates( join_node_id: NodeId, join: &JoinNode, ) -> Vec { - if join.spec.join_type != JoinType::Inner { - return Vec::new(); - } let Some(residual) = &join.spec.residual_predicate else { return Vec::new(); }; @@ -47,26 +47,30 @@ pub(super) fn residual_equality_route_candidates( collect_residual_input_equalities(residual, &mut equalities); let mut candidates = Vec::new(); for equality in equalities { - candidates.extend(residual_equality_route_candidates_for_orientation( - plan, - parent_counts, - join_node_id, - join, - equality.left_column.as_str(), - equality.right_column.as_str(), - JoinSide::Left, - JoinSide::Right, - )); - candidates.extend(residual_equality_route_candidates_for_orientation( - plan, - parent_counts, - join_node_id, - join, - equality.right_column.as_str(), - equality.left_column.as_str(), - JoinSide::Right, - JoinSide::Left, - )); + if join_type_discards_unmatched_rows(&join.spec.join_type, JoinSide::Right) { + candidates.extend(residual_equality_route_candidates_for_orientation( + plan, + parent_counts, + join_node_id, + join, + equality.left_column.as_str(), + equality.right_column.as_str(), + JoinSide::Left, + JoinSide::Right, + )); + } + if join_type_discards_unmatched_rows(&join.spec.join_type, JoinSide::Left) { + candidates.extend(residual_equality_route_candidates_for_orientation( + plan, + parent_counts, + join_node_id, + join, + equality.right_column.as_str(), + equality.left_column.as_str(), + JoinSide::Right, + JoinSide::Left, + )); + } } candidates } diff --git a/website/docs/design/gpu-execution.mdx b/website/docs/design/gpu-execution.mdx index 9a4a528df..11853593e 100644 --- a/website/docs/design/gpu-execution.mdx +++ b/website/docs/design/gpu-execution.mdx @@ -56,6 +56,19 @@ chunks. Empty and dense range payloads retain no key column. cuDF orders membership reads against the retained producer stream and retains completion dependencies until consumption ends. +Route discovery places a consumer only on a join input whose unmatched rows +the join discards. Inner and semi joins discard them on both inputs. A left +outer, anti or mark join discards them on its right input alone; each left row +reaches the output, or the mark column, whether it matches or not. A +null-aware anti join reads the nulls and the emptiness of its right input and +takes no route on either input. Only single-key joins under +`NullsDoNotMatch` produce a key route. Native execution builds the right input +of every keyed left outer and mark join. Build-side selection may move a keyed +semi or anti join's build to its left input, and the route then reduces the +right input. A key route at a left outer or mark join, or at an anti join that +builds its right input, would consume the left input and is rejected as +`UnsupportedJoinShape`. + Integer bounds and distinct counts classify empty, dense range, small-IN-list and larger exact keysets without copying the key column to the host. Resident exact payload metadata carries the existing kind, count and range diagnostics; From 8ec4056d7b7a0730760894f24bb3966905d01708 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Thu, 22 Oct 2026 00:27:13 +0000 Subject: [PATCH 6/8] fix(lowering): lower expression-keyed merges and InterleaveExec natively A merge of sorted partitions was rejected unless every sort key was a plain column, although a native sort lowers expression keys, so `ORDER BY value / global_pr DESC LIMIT 20` at the default partition count ran mixed with 23 host boundaries. Orderings are now compared by full PhysicalSortExpr equality and the merge accepts any key the native sort lowers. InterleaveExec, DataFusion's UNION ALL over identically hash-partitioned inputs, had no native form; a ChEMBL activity-cliff query fell back with 25 host boundaries and took 39.9 s. It now lowers through the native union path and runs fully native in 1.93 s with the same counts. Co-Authored-By: Claude Opus 5.5 --- .../src/gpu_coverage/tests.rs | 23 --- .../src/native/lowering/union.rs | 1 - .../lowering_tests/plan_shape_remedies.rs | 153 ++++++++++++------ .../src/native/lowering_tests/union_inputs.rs | 35 ++++ .../algeon-datafusion/src/native/normalize.rs | 42 ++--- .../src/native/normalize/types.rs | 43 ++++- .../native_report/diagnostics/shape_family.rs | 12 +- .../tests/native_exec_tests_gpu/main.rs | 86 ++++++++++ website/docs/design/native-lowering.mdx | 5 +- .../docs/reference/gpu-coverage-report.mdx | 2 +- website/docs/reference/native-sql-support.mdx | 2 +- 11 files changed, 297 insertions(+), 107 deletions(-) diff --git a/crates/algeon-datafusion/src/gpu_coverage/tests.rs b/crates/algeon-datafusion/src/gpu_coverage/tests.rs index b77533c46..25a02e4fa 100644 --- a/crates/algeon-datafusion/src/gpu_coverage/tests.rs +++ b/crates/algeon-datafusion/src/gpu_coverage/tests.rs @@ -338,29 +338,6 @@ fn candidate_outcomes_name_the_engine_gap_when_one_partition_keeps_the_rejection ); } -/// The branches of a `UNION ALL` stay separate partitions at one partition, so -/// DataFusion still merges them to order the rows, and the validation report -/// carries the engine gap instead of the setting that cannot help. -#[tokio::test] -async fn validate_query_names_the_engine_gap_for_a_merge_of_union_branches_at_one_partition() { - let ctx = validation_context(bounded_optimizer_config()); - ctx.sql("SET datafusion.execution.target_partitions = 1").await.unwrap().collect().await.unwrap(); - - let coverage = - validate_query(&ctx, "SELECT a, b FROM (SELECT 1 AS a, 2 AS b UNION ALL SELECT 3, 4) ORDER BY b + 1") - .await - .unwrap(); - - let rejected = coverage - .candidate_outcomes - .iter() - .find(|outcome| outcome.selection == CandidateOutcomeKind::NotSupported) - .expect("the merge of union branches is rejected"); - assert_eq!(rejected.reason_code.as_deref(), Some("native_sort_preserving_merge_unsupported")); - let remedy = rejected.remedy.as_ref().expect("a rejected candidate carries a remedy"); - assert_eq!((remedy.code, remedy.kind), ("unsupported_operator_gap", QueryGpuRemedyKind::EngineGap)); -} - #[tokio::test] async fn validate_query_preserves_the_installed_final_plan_policy() { let ctx = diff --git a/crates/algeon-datafusion/src/native/lowering/union.rs b/crates/algeon-datafusion/src/native/lowering/union.rs index f64fae3df..de1e1b9f4 100644 --- a/crates/algeon-datafusion/src/native/lowering/union.rs +++ b/crates/algeon-datafusion/src/native/lowering/union.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; -use datafusion::physical_plan::ExecutionPlan; use query_planning::plan::QueryPlanBuilder as NativeQueryPlanBuilder; use super::{ diff --git a/crates/algeon-datafusion/src/native/lowering_tests/plan_shape_remedies.rs b/crates/algeon-datafusion/src/native/lowering_tests/plan_shape_remedies.rs index b64ef8754..e761ad85c 100644 --- a/crates/algeon-datafusion/src/native/lowering_tests/plan_shape_remedies.rs +++ b/crates/algeon-datafusion/src/native/lowering_tests/plan_shape_remedies.rs @@ -33,7 +33,6 @@ fn rejection_reason(result: &NativeLoweringResult) -> Option { async fn target_partitions_one_avoids_the_plan_shapes_whose_remedy_names_it() { let dir = three_file_table("algeon-datafusion-plan-shape-remedies"); let cases = [ - ("SELECT key, payload FROM bench ORDER BY payload + 1", CapabilityReason::NativeSortPreservingMergeUnsupported), ( "SELECT key, MEDIAN(m00) AS median_m00 FROM bench GROUP BY key", CapabilityReason::NativeDistributedAggregateUnsupported, @@ -118,72 +117,128 @@ async fn distributed_aggregate_remedy_is_an_engine_gap_when_one_partition_cannot fs::remove_dir_all(dir).unwrap(); } -/// A merge of `UNION ALL` branches stays at one partition, and a sort key -/// without a native form is rejected by the single sort DataFusion plans there. -/// Both carry an engine gap at either partition count, and the single sort -/// reports the expression it cannot lower. The last query unions on the -/// collected side of a join, which DataFusion merges into one partition, so the -/// setting still removes the merge. +fn plan_has_expression_keyed_merge(plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::() + .is_some_and(|merge| merge.expr().iter().any(|sort_expr| sort_expr.expr.downcast_ref::().is_none())) + || plan.children().into_iter().any(|child| plan_has_expression_keyed_merge(child.as_ref())) +} + +/// One native sort by the merge's keys reproduces a merge of sorted +/// partitions, so a merge whose key is an expression the native sort accepts +/// lowers at any partition count. That holds when the merge reads `UNION ALL` +/// branches directly or through either side of a join. A key without a native +/// form is rejected as the expression it is at either partition count. #[tokio::test] -async fn sort_preserving_merge_remedy_is_an_engine_gap_when_one_partition_cannot_lower_it() { - let dir = three_file_table("algeon-datafusion-merge-gap-remedy"); +async fn expression_keyed_merges_lower_at_any_partition_count() { + let dir = three_file_table("algeon-datafusion-expression-keyed-merge"); let union_branches = "SELECT key, payload FROM bench UNION ALL SELECT key, payload FROM bench"; - let gap = ("unsupported_operator_gap", QueryGpuRemedyKind::EngineGap); - let setting = ("set_target_partitions_one", QueryGpuRemedyKind::ConfigAction); let cases = [ - ( - format!("SELECT key, payload FROM ({union_branches}) ORDER BY payload + 1"), - gap, - Some((CapabilityReason::NativeSortPreservingMergeUnsupported, gap)), - ), - ( - format!("SELECT key, payload FROM ({union_branches}) ORDER BY payload + 1 LIMIT 5"), - gap, - Some((CapabilityReason::NativeSortPreservingMergeUnsupported, gap)), - ), + ("SELECT key, payload FROM bench ORDER BY payload + 1".to_string(), None), + (format!("SELECT key, payload FROM ({union_branches}) ORDER BY payload + 1"), None), + (format!("SELECT key, payload FROM ({union_branches}) ORDER BY payload + 1 LIMIT 5"), None), ( "SELECT a.key, a.payload FROM bench a JOIN (SELECT key FROM bench UNION ALL SELECT key FROM bench) b \ ON a.key = b.key ORDER BY a.payload + 1" .to_string(), - gap, - Some((CapabilityReason::NativeSortPreservingMergeUnsupported, gap)), - ), - ( - "SELECT key, payload FROM bench ORDER BY random()".to_string(), - gap, - Some((CapabilityReason::UnsupportedExpression, ("rewrite_expression", QueryGpuRemedyKind::UserAction))), + None, ), ( "SELECT b.key, a.payload FROM (SELECT key FROM bench UNION ALL SELECT key FROM bench) b \ JOIN bench a ON a.key = b.key ORDER BY a.payload + 1" .to_string(), - setting, None, ), + ( + "SELECT key, payload FROM bench ORDER BY random()".to_string(), + Some(("rewrite_expression", QueryGpuRemedyKind::UserAction)), + ), ]; - for (sql, remedy_at_four, at_one) in cases { - let plan = sql_plan(&dir, SessionConfig::new().with_target_partitions(4), &sql).await; + for (sql, rejection) in cases { + for target_partitions in [4, 1] { + let plan = sql_plan(&dir, SessionConfig::new().with_target_partitions(target_partitions), &sql).await; + let context = format!("target_partitions = {target_partitions}: {sql}"); + if target_partitions > 1 { + assert!(plan_has_expression_keyed_merge(plan.as_ref()), "{context}"); + } + let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap(); + match rejection { + None => assert_eq!(result.selected_count(), 1, "{context}: {:?}", result.report.outcomes()), + Some(expected) => { + assert_eq!(rejection_reason(&result), Some(CapabilityReason::UnsupportedExpression), "{context}"); + let remedy = rejected_candidate_remedy(&result); + assert_eq!((remedy.code, remedy.kind), expected, "{context}"); + }, + } + } + } + fs::remove_dir_all(dir).unwrap(); +} + +/// A merge's fetch and an outer offset apply to the one native sort by the +/// merge's expression key, so the native plan keeps the rows DataFusion's +/// merge returns. +#[tokio::test] +async fn expression_keyed_merge_fetch_limits_the_native_sort() { + let dir = three_file_table("algeon-datafusion-expression-keyed-merge-fetch"); + for (sql, expected) in [ + ("SELECT key, payload FROM bench ORDER BY payload * 2 DESC LIMIT 2", LimitSpec::new(0, Some(2))), + ("SELECT key, payload FROM bench ORDER BY payload * 2 DESC LIMIT 2 OFFSET 1", LimitSpec::new(1, Some(2))), + ] { + let plan = sql_plan(&dir, SessionConfig::new().with_target_partitions(4), sql).await; + assert!(plan_has_expression_keyed_merge(plan.as_ref()), "{sql}"); + let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap(); - let context = format!("target_partitions = 4: {sql}"); - assert_eq!( - rejection_reason(&result), - Some(CapabilityReason::NativeSortPreservingMergeUnsupported), - "{context}" - ); - let remedy = rejected_candidate_remedy(&result); - assert_eq!((remedy.code, remedy.kind), remedy_at_four, "{context}"); - let plan = sql_plan(&dir, SessionConfig::new().with_target_partitions(1), &sql).await; + assert_eq!(result.selected_count(), 1, "{sql}: {:?}", result.report.outcomes()); + let native_plan = result.selected_plan(); + let limits = native_plan + .nodes() + .iter() + .filter_map(|node| match &node.kind { + PlanNode::Limit(limit) => Some(limit), + _ => None, + }) + .collect::>(); + let [limit] = limits.as_slice() else { + panic!("{sql}: expected one native limit:\n{native_plan}"); + }; + assert_eq!(limit.spec, expected, "{sql}"); + let PlanNode::Sort(sort) = &native_plan.node(limit.input).expect("limit input exists").kind else { + panic!("{sql}: the limit reads the sort:\n{native_plan}"); + }; + assert!(matches!(sort.spec.expressions.as_slice(), [key] if key.descending), "{sql}: {:?}", sort.spec); + } + fs::remove_dir_all(dir).unwrap(); +} + +/// A per-partition sort that keeps fewer rows than the merge above it returns +/// has no single native sort, so the merge is rejected. Over one scan, one +/// partition leaves a single sort, and the remedy names the setting. Over +/// `UNION ALL` branches, which stay partitions at one partition, the merge and +/// its rejection remain, so the remedy is the engine gap. +#[tokio::test] +async fn merge_over_a_smaller_per_partition_fetch_names_one_partition_unless_it_reads_union_branches() { + let dir = three_file_table("algeon-datafusion-merge-partition-fetch-remedy"); + let scan = sql_plan(&dir, SessionConfig::new().with_target_partitions(4), "SELECT key, payload FROM bench").await; + assert!(scan.properties().output_partitioning().partition_count() > 1); + let union: Arc = UnionExec::try_new(vec![Arc::clone(&scan), Arc::clone(&scan)]).unwrap(); + let setting = ("set_target_partitions_one", QueryGpuRemedyKind::ConfigAction); + let gap = ("unsupported_operator_gap", QueryGpuRemedyKind::EngineGap); + for (input, expected) in [(scan, setting), (union, gap)] { + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("payload", &input.schema()).unwrap(), + ArrowSortOptions::new(false, false), + )]) + .unwrap(); + let sort: Arc = + Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true).with_fetch(Some(1))); + let plan: Arc = Arc::new(SortPreservingMergeExec::new(ordering, sort)); + let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap(); - let context = format!("target_partitions = 1: {sql}"); - match at_one { - Some((reason, expected)) => { - assert_eq!(rejection_reason(&result), Some(reason), "{context}"); - let remedy = rejected_candidate_remedy(&result); - assert_eq!((remedy.code, remedy.kind), expected, "{context}"); - }, - None => assert_eq!(result.selected_count(), 1, "{context}: {:?}", result.report.outcomes()), - } + + assert_eq!(rejection_reason(&result), Some(CapabilityReason::NativeSortPreservingMergeUnsupported)); + let remedy = rejected_candidate_remedy(&result); + assert_eq!((remedy.code, remedy.kind), expected); } fs::remove_dir_all(dir).unwrap(); } diff --git a/crates/algeon-datafusion/src/native/lowering_tests/union_inputs.rs b/crates/algeon-datafusion/src/native/lowering_tests/union_inputs.rs index cbe5adfa0..01696cb07 100644 --- a/crates/algeon-datafusion/src/native/lowering_tests/union_inputs.rs +++ b/crates/algeon-datafusion/src/native/lowering_tests/union_inputs.rs @@ -1,4 +1,7 @@ +use datafusion::physical_plan::displayable; + use super::*; +use crate::native::ParquetReadStrategyConfig; /// The plan builder interns identical nodes, so UNION ALL branches that read /// the same table, or the same aggregate CTE, become one native node that the @@ -37,3 +40,35 @@ async fn unions_reading_one_native_node_twice_plan_natively() { fs::remove_dir_all(dir).unwrap(); } + +/// DataFusion plans a `UNION ALL` whose branches share one hash partitioning +/// as an `InterleaveExec`, which emits every branch row once like a union. +/// The candidate keeps the operator's name and lowers to one native union. +#[tokio::test] +async fn interleaved_union_of_partitioned_joins_lowers_to_a_native_union() { + let dir = unique_temp_dir("algeon-datafusion-native-lowering-interleave"); + for index in 0..3 { + write_numeric_parquet_file(&dir.join(format!("part{index}.parquet"))); + } + let mut config = SessionConfig::new().with_target_partitions(4); + config.options_mut().optimizer.hash_join_single_partition_threshold = 0; + config.options_mut().optimizer.hash_join_single_partition_threshold_rows = 0; + let sql = "SELECT key, COUNT(*) AS n FROM (\ + SELECT a.key FROM bench a JOIN bench b ON a.key = b.key WHERE a.payload >= 10 \ + UNION ALL SELECT a.key FROM bench a JOIN bench b ON a.key = b.key WHERE a.payload < 10) \ + GROUP BY key"; + let plan = sql_plan(&dir, config, sql).await; + let tree = displayable(plan.as_ref()).indent(false).to_string(); + assert!(tree.contains("InterleaveExec"), "{tree}"); + + let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap(); + + let [CandidateOutcome::Selected { candidate }] = result.report.outcomes() else { + panic!("{tree}\noutcomes={:?}", result.report.outcomes()); + }; + assert!(candidate.shape.contains_operator("InterleaveExec"), "{}", candidate.shape); + let unions = result.selected_plan().nodes().iter().filter(|node| matches!(node.kind, PlanNode::Union(_))).count(); + assert_eq!(unions, 1, "{}", result.selected_plan()); + + fs::remove_dir_all(dir).unwrap(); +} diff --git a/crates/algeon-datafusion/src/native/normalize.rs b/crates/algeon-datafusion/src/native/normalize.rs index 67421fffd..42b13f73b 100644 --- a/crates/algeon-datafusion/src/native/normalize.rs +++ b/crates/algeon-datafusion/src/native/normalize.rs @@ -13,7 +13,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use datafusion::{ datasource::{memory::MemorySourceConfig, source::DataSourceExec}, - physical_expr::{LexOrdering, PhysicalSortExpr, expressions::Column}, + physical_expr::LexOrdering, physical_plan::{ ExecutionPlan, ExecutionPlanProperties, Partitioning, aggregates::AggregateExec, @@ -27,7 +27,7 @@ use datafusion::{ projection::ProjectionExec, repartition::RepartitionExec, sorts::{sort::SortExec, sort_preserving_merge::SortPreservingMergeExec}, - union::UnionExec, + union::{InterleaveExec, UnionExec}, windows::{BoundedWindowAggExec, WindowAggExec}, }, }; @@ -52,7 +52,7 @@ pub(in crate::native) use types::{ NormalizedNativeGpuResidentSource, NormalizedNativeHashJoin, NormalizedNativeInput, NormalizedNativeLimit, NormalizedNativeLiteralSource, NormalizedNativeNestedLoopJoin, NormalizedNativePlan, NormalizedNativeProjection, NormalizedNativeSort, NormalizedNativeSortExprs, NormalizedNativeSortOrigin, NormalizedNativeSource, - NormalizedNativeUnion, NormalizedNativeWindow, NormalizedNativeWindowExec, + NormalizedNativeUnion, NormalizedNativeUnionExec, NormalizedNativeWindow, NormalizedNativeWindowExec, }; /// Normalizes the physical subtree rooted at `plan` into one native input @@ -86,6 +86,7 @@ fn root_operator_is_normalized(plan: &dyn ExecutionPlan) -> bool { || plan.downcast_ref::().is_some() || plan.downcast_ref::().is_some() || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() || gpu_resident_producer_input_name(plan).is_some() } @@ -117,7 +118,8 @@ pub(super) fn make_declared_ordered_hash_join_fetch_explicit<'a>( } /// Normalizes a merge of sorted partitions. One native sort over the -/// concatenated partitions reproduces the merge. When the partitions come from +/// concatenated partitions by the merge's keys reproduces the merge, so the +/// merge accepts any key the native sort lowers. When the partitions come from /// a per-partition sort, the merge only adds its fetch, which must not ask for /// more rows than each partition kept. fn normalize_sort_preserving_merge<'a>( @@ -130,9 +132,7 @@ fn normalize_sort_preserving_merge<'a>( blocker: Box::new(PlanBlocker::Operator(merge.name().to_string())), single_partition: merge_single_partition_outcome(merge), }; - if !sort_ordering_is_column_keyed(merge.expr()) - || !plan_partitions_match_ordering(merge.input().as_ref(), merge.expr()) - { + if !plan_partitions_match_ordering(merge.input().as_ref(), merge.expr()) { return Err(unsupported()); } let merge_fetch = merge.fetch().map(NativeLimitSpec::fetch); @@ -143,7 +143,7 @@ fn normalize_sort_preserving_merge<'a>( (Some(merge_fetch), Some(sort_fetch)) if sort_fetch < merge_fetch => return Err(unsupported()), _ => {}, } - if projection.is_some() || sort_orderings_match(merge.expr(), sort.expr()) { + if projection.is_some() || merge.expr() == sort.expr() { let sorted = normalize_native_input(merge.input().as_ref(), candidate, true)?; return Ok(with_merge_fetch(sorted, merge_fetch)); } @@ -456,7 +456,12 @@ fn normalize_native_input( right: Box::new(right), })); } - if let Some(union) = plan.downcast_ref::() { + let union = if let Some(union) = plan.downcast_ref::() { + Some(NormalizedNativeUnionExec::Union(union)) + } else { + plan.downcast_ref::().map(NormalizedNativeUnionExec::Interleave) + }; + if let Some(union) = union { let inputs = union .inputs() .iter() @@ -498,7 +503,8 @@ fn peel_repartition(plan: &dyn ExecutionPlan) -> &dyn ExecutionPlan { /// partitions or a hash repartition on the way ends that path: the partitions /// below no longer reach the operator above. DataFusion also merges the /// collected side of a join, so a union there ends at the merge and only a -/// union on the probe side counts. +/// union on the probe side counts. An `InterleaveExec` is not such a union: +/// each of its partitions combines the same hash partition of every input. fn reads_union_branches(stage: &dyn ExecutionPlan) -> bool { stage.is::() || stage.children().into_iter().any(|input| { @@ -521,19 +527,3 @@ fn ordered_multi_partition_memory_source_is_unsupported(source_exec: &DataSource .downcast_ref::() .is_some_and(|memory| source_exec.properties().output_ordering().is_some() && memory.partitions().len() > 1) } - -fn sort_orderings_match(left: &LexOrdering, right: &LexOrdering) -> bool { - if left.len() != right.len() { - return false; - } - left.iter().zip(right.iter()).all(|(left, right)| sort_expr_key(left) == sort_expr_key(right)) -} - -fn sort_ordering_is_column_keyed(ordering: &LexOrdering) -> bool { - !ordering.is_empty() && ordering.iter().all(|sort_expr| sort_expr.expr.downcast_ref::().is_some()) -} - -fn sort_expr_key(sort_expr: &PhysicalSortExpr) -> Option<(usize, &str, bool, bool)> { - let column = sort_expr.expr.downcast_ref::()?; - Some((column.index(), column.name(), sort_expr.options.descending, sort_expr.options.nulls_first)) -} diff --git a/crates/algeon-datafusion/src/native/normalize/types.rs b/crates/algeon-datafusion/src/native/normalize/types.rs index abc244fa5..bc5d6d158 100644 --- a/crates/algeon-datafusion/src/native/normalize/types.rs +++ b/crates/algeon-datafusion/src/native/normalize/types.rs @@ -4,7 +4,7 @@ use arrow::datatypes::SchemaRef; use datafusion::{ datasource::source::DataSourceExec, physical_expr::window::WindowExpr, - physical_plan::{empty::EmptyExec, placeholder_row::PlaceholderRowExec}, + physical_plan::{empty::EmptyExec, placeholder_row::PlaceholderRowExec, union::InterleaveExec}, }; use query_planning::expr::LimitSpec; @@ -389,16 +389,51 @@ impl NormalizedNativeNestedLoopJoin<'_> { } } -/// DataFusion physical union with recursively normalized children. +/// DataFusion `UNION ALL` with recursively normalized children. #[derive(Debug, Clone)] pub(in crate::native) struct NormalizedNativeUnion<'a> { - pub(in crate::native) union: &'a UnionExec, + pub(in crate::native) union: NormalizedNativeUnionExec<'a>, pub(in crate::native) inputs: Vec>, } +/// The DataFusion operator of a `UNION ALL`. Both emit every input row once +/// and differ only in how they assign rows to partitions, which native +/// execution does not keep. +#[derive(Debug, Clone, Copy)] +pub(in crate::native) enum NormalizedNativeUnionExec<'a> { + /// Each input partition stays a separate output partition. + Union(&'a UnionExec), + /// Output partition `k` combines partition `k` of every input, which + /// share one hash or range partitioning. + Interleave(&'a InterleaveExec), +} + +impl<'a> NormalizedNativeUnionExec<'a> { + pub(in crate::native) fn inputs(self) -> &'a [Arc] { + match self { + Self::Union(union) => union.inputs(), + Self::Interleave(interleave) => interleave.inputs(), + } + } + + pub(in crate::native) fn schema(&self) -> SchemaRef { + match self { + Self::Union(union) => union.schema(), + Self::Interleave(interleave) => interleave.schema(), + } + } + + fn name(&self) -> &'static str { + match self { + Self::Union(_) => "UnionExec", + Self::Interleave(_) => "InterleaveExec", + } + } +} + impl NormalizedNativeUnion<'_> { fn shape(&self) -> CandidateShape { - CandidateShape::new("UnionExec", self.inputs.iter().map(NormalizedNativeInput::shape).collect()) + CandidateShape::new(self.union.name(), self.inputs.iter().map(NormalizedNativeInput::shape).collect()) } } diff --git a/crates/algeon-datafusion/src/native_report/diagnostics/shape_family.rs b/crates/algeon-datafusion/src/native_report/diagnostics/shape_family.rs index 8fc334fa9..276275cb4 100644 --- a/crates/algeon-datafusion/src/native_report/diagnostics/shape_family.rs +++ b/crates/algeon-datafusion/src/native_report/diagnostics/shape_family.rs @@ -69,7 +69,7 @@ fn input_family(shape: &CandidateShape) -> &'static str { (false, false, false) => {}, _ => return "mixed_join", } - if shape.contains_operator("UnionExec") { + if shape.contains_operator("UnionExec") || shape.contains_operator("InterleaveExec") { return "union"; } let mut sources = Vec::new(); @@ -137,6 +137,16 @@ mod tests { )), "sort_union_input" ); + assert_eq!( + family(unary( + "AggregateExec", + CandidateShape::new( + "InterleaveExec", + vec![CandidateShape::leaf("LocalParquetScan"), CandidateShape::leaf("LocalParquetScan")] + ) + )), + "global_union_aggregate" + ); assert_eq!( family(unary( "AggregateExec", diff --git a/crates/algeon-datafusion/tests/native_exec_tests_gpu/main.rs b/crates/algeon-datafusion/tests/native_exec_tests_gpu/main.rs index 93bcadfc3..f63fc750a 100644 --- a/crates/algeon-datafusion/tests/native_exec_tests_gpu/main.rs +++ b/crates/algeon-datafusion/tests/native_exec_tests_gpu/main.rs @@ -1421,6 +1421,92 @@ fn whole_input_aggregate_windows_match_datafusion() { fs::remove_dir_all(dir).unwrap(); } +/// At four partitions DataFusion merges sorted partitions by an expression +/// key, and plans a `UNION ALL` of joins that share the grouping key's hash +/// partitioning as an `InterleaveExec`. Each query lowers whole under +/// NativeRequired and returns DataFusion's rows. +#[test] +fn multi_partition_expression_merge_and_interleaved_union_match_datafusion() { + const FILES: i32 = 3; + const ROWS_PER_FILE: i32 = 2_000; + let dir = unique_temp_dir("algeon-datafusion-native-exec-merge-interleave"); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Float64, false), + Field::new("weight", DataType::Float64, false), + ])); + for file in 0..FILES { + let ids = (file * ROWS_PER_FILE)..((file + 1) * ROWS_PER_FILE); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(ids.clone())), + Arc::new(Int32Array::from_iter_values(ids.clone().map(|id| id % 997))), + Arc::new(Float64Array::from_iter_values(ids.clone().map(|id| f64::from((id * 7919) % 1009)))), + Arc::new(Float64Array::from_iter_values(ids.map(|id| f64::from(id % 13 + 1)))), + ], + ) + .unwrap(); + let file = File::create(dir.join(format!("part{file}.parquet"))).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + let config = || { + let mut config = SessionConfig::new().with_target_partitions(4); + config.options_mut().optimizer.hash_join_single_partition_threshold = 0; + config.options_mut().optimizer.hash_join_single_partition_threshold_rows = 0; + config + }; + let required = SessionContext::new_with_state( + SessionStateBuilder::new() + .with_default_features() + .with_config(config()) + .with_algeon_native(native_optimizer_config().with_execution_mode(AlgeonGpuExecutionMode::NativeRequired)) + .build(), + ); + let cpu = SessionContext::new_with_config(config()); + let cases = [ + ("SortPreservingMergeExec", "SELECT id, value, weight FROM input ORDER BY value / weight DESC, id LIMIT 20"), + ( + "InterleaveExec", + "SELECT key, COUNT(*) AS n, SUM(value) AS s FROM (\ + SELECT a.key, a.value FROM input a JOIN input b ON a.key = b.key WHERE a.weight >= 7 \ + UNION ALL SELECT a.key, a.value FROM input a JOIN input b ON a.key = b.key WHERE a.weight < 7) u \ + GROUP BY key ORDER BY key", + ), + ]; + let runtime = Builder::new_multi_thread().worker_threads(4).enable_all().build().unwrap(); + runtime.block_on(async { + for ctx in [&required, &cpu] { + ctx.register_parquet("input", dir.to_str().unwrap(), ParquetReadOptions::default()).await.unwrap(); + } + for (operator, query) in cases { + let cpu_plan = cpu.sql(query).await.unwrap().create_physical_plan().await.unwrap(); + let cpu_tree = displayable(cpu_plan.as_ref()).indent(false).to_string(); + assert!(cpu_tree.contains(operator), "{query}\n{cpu_tree}"); + let native_plan = required + .sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap_or_else(|error| panic!("NativeRequired plans {query}: {error}")); + assert!( + below_native_admission(&native_plan).downcast_ref::().is_some(), + "{query}\n{}", + displayable(native_plan.as_ref()).indent(false) + ); + let rows = |batches: Vec| pretty_format_batches(&batches).unwrap().to_string(); + let actual = rows(collect(native_plan, Arc::new(TaskContext::default())).await.unwrap()); + let expected = rows(collect(cpu_plan, Arc::new(TaskContext::default())).await.unwrap()); + assert_eq!(actual, expected, "{query}"); + } + }); + fs::remove_dir_all(dir).unwrap(); +} + /// DataFusion inserts a multi-partition partition-preserving SortExec under /// window aggregates when target_partitions > 1. The native planner must /// normalize this input into a native sort and lower the entire window plan diff --git a/website/docs/design/native-lowering.mdx b/website/docs/design/native-lowering.mdx index 3bc3cdda5..415b51be3 100644 --- a/website/docs/design/native-lowering.mdx +++ b/website/docs/design/native-lowering.mdx @@ -132,7 +132,10 @@ CUDA streams. Lowering reconstructs operations whose semantics span chunks, such as multi-pass aggregates (`AVG` retaining sums and counts for final division) and distributed -sort and limit merges. +sort and limit merges. A merge of sorted partitions lowers as one native sort by +the merge's keys, expression keys included. An `InterleaveExec`, which +DataFusion plans for a `UNION ALL` whose branches share a hash partitioning, +lowers as a union. A window with `ORDER BY` keys, or one that numbers rows by position inside partitions, reads its input ordered by its partition keys, then its `ORDER BY` diff --git a/website/docs/reference/gpu-coverage-report.mdx b/website/docs/reference/gpu-coverage-report.mdx index f07939823..144538c2a 100644 --- a/website/docs/reference/gpu-coverage-report.mdx +++ b/website/docs/reference/gpu-coverage-report.mdx @@ -175,7 +175,7 @@ inside an executable plan, and the statement is not rejected. | `reason_code` | `remedy_code` (`remedy_kind`) | When it applies | | --- | --- | --- | -| `native_sort_preserving_merge_unsupported` | `set_target_partitions_one` (`config_action`) | DataFusion merged sorted partitions because it planned several partitions. | +| `native_sort_preserving_merge_unsupported` | `set_target_partitions_one` (`config_action`) | DataFusion merged sorted partitions whose per-partition sort keeps fewer rows than the merge returns, or whose input does not report the merge's ordering. | | `native_distributed_aggregate_unsupported` | `set_target_partitions_one` (`config_action`) | DataFusion planned partial and final aggregate stages over several partitions. | | `native_local_limit_partition_unsupported` | `set_target_partitions_one` (`config_action`) | DataFusion planned a limit that works on each of several partitions. | | `native_distributed_aggregate_unsupported` | `unsupported_operator_gap` (`engine_gap`) | The stages remain at one partition, or the single-stage aggregate fails native validation. | diff --git a/website/docs/reference/native-sql-support.mdx b/website/docs/reference/native-sql-support.mdx index 43b6e65bf..865ea190f 100644 --- a/website/docs/reference/native-sql-support.mdx +++ b/website/docs/reference/native-sql-support.mdx @@ -19,7 +19,7 @@ reports the exact coverage; this page summarizes the rules. | Area | Native forms | | --- | --- | | Sources | Local and S3 object-store Parquet. Projected columns, row-group selection, partition columns, reader predicates and scan limits when their contracts are supported. | -| Row operations | Filters, projections, stable or unstable sorts, `LIMIT`/`OFFSET`, and schema-compatible `UNION ALL`. A distinct `UNION` and `SELECT DISTINCT` are grouped aggregates. | +| Row operations | Filters, projections, stable or unstable sorts, `LIMIT`/`OFFSET`, and schema-compatible `UNION ALL`. A merge of sorted partitions lowers as one sort by its keys, and a `UNION ALL` that DataFusion plans as `InterleaveExec` lowers as a union. A distinct `UNION` and `SELECT DISTINCT` are grouped aggregates. | | Aggregation | Global and grouped aggregation, admitted grouping sets, aggregate filters, and partial/final state when the state is explicitly modeled. | | Windows | `rank`, `row_number`, `lag`, `lead`, whole-partition aggregate windows, and cumulative `sum`/`min`/`max` over `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`. `rank`, `lag` and `lead` require `ORDER BY`; `row_number` does not. An empty `OVER ()` aggregates the whole input as one partition. | | Joins | Equi joins and admitted residual or keyless forms for inner, outer, semi, anti, mark and cross joins. Build side, output projection, key types and residual expressions can narrow an otherwise supported join kind. | From 21339a27181ca5a06c16f8fc0650fa8ca3bb877d Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Thu, 22 Oct 2026 00:27:13 +0000 Subject: [PATCH 7/8] fix(capability): stop reading left anti joins as EXCEPT DISTINCT An edge guard rejected a left anti join whose build is a group-only aggregate unless its probe input was proven distinct, and rejected any grouped build with measures. It began as a recognizer rule for one set-operation envelope and was promoted to generic edge capability, but every bounded anti path keeps each unmatched probe row and DataFusion plans EXCEPT DISTINCT with an explicit dedup. `LEFT JOIN (SELECT DISTINCT ...) d ... WHERE d.k IS NULL` over a join output fell back to the CPU for no reason. Both arms and their reasons (left_anti_build_requires_dedup_aggregate, left_anti_probe_requires_dedup_proven_input) are removed. A GPU case keeps duplicate probe rows equal to DataFusion. Co-Authored-By: Claude Opus 5.5 --- .../public-api/algeon-query-planning.txt | 2 - .../native_exec_tests_gpu/bounded_join.rs | 52 +++++++ .../fixtures_plans/projection_scalar_setop.rs | 10 +- .../bounded_join_admission/mod.rs | 2 +- .../bounded_join_admission/support.rs | 146 ++++++++---------- .../src/capability/aggregate.rs | 13 +- .../src/capability/edge.rs | 117 +------------- .../src/capability/graph.rs | 7 - .../src/capability/model.rs | 6 +- .../src/capability/rejection.rs | 8 - 10 files changed, 136 insertions(+), 227 deletions(-) diff --git a/.config/policies/public-api/algeon-query-planning.txt b/.config/policies/public-api/algeon-query-planning.txt index 9b7ed9ca1..d5980debd 100644 --- a/.config/policies/public-api/algeon-query-planning.txt +++ b/.config/policies/public-api/algeon-query-planning.txt @@ -192,8 +192,6 @@ pub query_planning::capability::Reason::DatafusionExactStatisticalGroupKeyUnsupp pub query_planning::capability::Reason::DatafusionExactUngroupedFloatSumUnsupported pub query_planning::capability::Reason::ExpressionLineageRequiresResolvedInputColumn pub query_planning::capability::Reason::LargeKeyedHashJoinStrategyUnsupported -pub query_planning::capability::Reason::LeftAntiBuildRequiresDedupAggregate -pub query_planning::capability::Reason::LeftAntiProbeRequiresDedupProvenInput pub query_planning::capability::Reason::MissingEdgeProducer pub query_planning::capability::Reason::OutputLineageAggregateFieldIndexOutOfRange pub query_planning::capability::Reason::OutputLineageFieldIndexOutOfRange diff --git a/crates/algeon-datafusion/tests/native_exec_tests_gpu/bounded_join.rs b/crates/algeon-datafusion/tests/native_exec_tests_gpu/bounded_join.rs index 64bad73cd..f6662c0e9 100644 --- a/crates/algeon-datafusion/tests/native_exec_tests_gpu/bounded_join.rs +++ b/crates/algeon-datafusion/tests/native_exec_tests_gpu/bounded_join.rs @@ -80,6 +80,54 @@ fn native_optimizer_rule_rewrites_and_executes_bounded_left_anti_join_plan() { fs::remove_dir_all(dir).unwrap(); } +/// `LEFT JOIN (SELECT DISTINCT ...) d ... WHERE d.key IS NULL` over an +/// inner-join probe becomes a LeftAnti with a group-only dedup build. The +/// probe repeats key 1 with identical rows, and the anti join keeps both. +fn native_optimizer_rule_executes_left_anti_over_join_probe_keeping_duplicate_rows() { + let dir = unique_temp_dir("algeon-datafusion-native-exec-bounded-left-anti-join-probe-duplicates"); + let left = dir.join("left.parquet"); + let right = dir.join("right.parquet"); + write_join_left_parquet_file(&left); + write_join_right_parquet_file(&right); + let plan = single_partition_sql_plan_with_parquet_tables( + &[("left_input", &left), ("right_input", &right)], + "SELECT p.key, p.r_payload \ + FROM (SELECT l.key, r.r_payload FROM left_input l JOIN right_input r ON r.key = l.key) p \ + LEFT JOIN (SELECT DISTINCT r.key FROM right_input r WHERE r.r_m01 > 1) d ON d.key = p.key \ + WHERE d.key IS NULL \ + ORDER BY p.key", + ); + let rule = AlgeonNativePhysicalOptimizerRule::new(native_optimizer_config()); + + let (optimized, diagnostics) = + rule.optimize_with_diagnostics(Arc::clone(&plan), &ConfigOptions::default()).unwrap(); + + let diagnostics = diagnostics.expect("native optimizer should record a planning report"); + let report = &diagnostics.report; + assert!( + below_native_admission(&optimized).downcast_ref::().is_some(), + "outcomes={:?}; plan={}", + report.outcomes(), + displayable(optimized.as_ref()).indent(false) + ); + assert_eq!(report.selected_count(), 1, "outcomes={:?}", report.outcomes()); + assert_eq!(report.not_selected_count(), 0, "outcomes={:?}", report.outcomes()); + assert_eq!( + diagnostics.native_optimizer.outer_join_null_rejection.applied_count, 1, + "{:?}", + diagnostics.native_optimizer.outer_join_null_rejection + ); + + let native_rows = + null_aware_anti_join_rows(&block_on(collect(optimized, Arc::new(TaskContext::default()))).unwrap()); + let datafusion_rows = + null_aware_anti_join_rows(&block_on(collect(plan, Arc::new(TaskContext::default()))).unwrap()); + assert_eq!(datafusion_rows, vec![(Some(1), 10), (Some(1), 10)]); + assert_eq!(native_rows, datafusion_rows); + + fs::remove_dir_all(dir).unwrap(); +} + fn null_aware_anti_join_rows(batches: &[RecordBatch]) -> Vec<(Option, i32)> { let mut rows = batches .iter() @@ -149,6 +197,10 @@ fn bounded_join_cases_1() { "native_optimizer_rule_rewrites_and_executes_bounded_left_anti_join_plan", native_optimizer_rule_rewrites_and_executes_bounded_left_anti_join_plan, ), + ( + "native_optimizer_rule_executes_left_anti_over_join_probe_keeping_duplicate_rows", + native_optimizer_rule_executes_left_anti_over_join_probe_keeping_duplicate_rows, + ), ]); } diff --git a/crates/algeon-query-engine/src/exec/bounded/compiled/execution_tests/fixtures_plans/projection_scalar_setop.rs b/crates/algeon-query-engine/src/exec/bounded/compiled/execution_tests/fixtures_plans/projection_scalar_setop.rs index e84048489..83635902a 100644 --- a/crates/algeon-query-engine/src/exec/bounded/compiled/execution_tests/fixtures_plans/projection_scalar_setop.rs +++ b/crates/algeon-query-engine/src/exec/bounded/compiled/execution_tests/fixtures_plans/projection_scalar_setop.rs @@ -802,8 +802,7 @@ fn projection_root_setop_dedup_scalar_plan_with_variant(variant: ProjectionScala /// Three group-only dedup aggregates combined by two stacked keyed LeftAnti /// joins under an outermost groupless `count(*)`, i.e. `count(*) from ((a /// except b) except c)`. Models a CHAIN of EXCEPT DISTINCTs. The second anti -/// join's probe side is the first anti join itself. That probe side is not a -/// dedup aggregate, so its dedup is proven transitively. +/// join's probe side is the first anti join itself, not a dedup aggregate. fn projection_root_setop_chained_dedup_scalar_plan() -> QueryPlan { let key_schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); let one_schema = Arc::new(Schema::new(vec![Field::new("one", DataType::Int64, false)])); @@ -928,10 +927,9 @@ fn composed_pipeline_refuses_keyless_anti_join_fanout_as_malformed() { #[test] fn composed_pipeline_admits_chained_except_distinct_dedup_scalar() { - // The setop-dedup proof moved off the plan-family recognizer: aggregate - // nodes carry `aggregate_state`, and each LeftAnti input edge checks it. A - // CHAIN of EXCEPT DISTINCTs stays admitted because every build is a - // group-only dedup aggregate and every probe is dedup-proven. + // Each EXCEPT DISTINCT carries its dedup as an explicit group-only + // aggregate, and every LeftAnti keeps per-probe-row semantics, so a CHAIN + // of them is admitted with no set-operation reading of its edges. let plan = projection_root_setop_chained_dedup_scalar_plan(); assert!( bounded_pipeline_admission(&plan).unwrap().is_some(), diff --git a/crates/algeon-query-engine/tests/native_contract_tests_cpu/bounded_join_admission/mod.rs b/crates/algeon-query-engine/tests/native_contract_tests_cpu/bounded_join_admission/mod.rs index 31f89a388..94f6306fe 100644 --- a/crates/algeon-query-engine/tests/native_contract_tests_cpu/bounded_join_admission/mod.rs +++ b/crates/algeon-query-engine/tests/native_contract_tests_cpu/bounded_join_admission/mod.rs @@ -11,7 +11,7 @@ use query_planning::{ bounded_join_unsupported_shape_facts, compiled_bounded_pipeline_support, compiled_bounded_query_plan_support, }, backend::AnalysisContext, - capability::{CapabilityDecision, EdgeConsumer, JoinForm, JoinStrategy, Reason}, + capability::{CapabilityDecision, EdgeConsumer, JoinForm, JoinStrategy}, error::{QueryError, QueryErrorCode, Result}, expr::Expr, mode::StatisticalAggregateExecutionMode, diff --git a/crates/algeon-query-engine/tests/native_contract_tests_cpu/bounded_join_admission/support.rs b/crates/algeon-query-engine/tests/native_contract_tests_cpu/bounded_join_admission/support.rs index 7a08e6bd8..67eb4db37 100644 --- a/crates/algeon-query-engine/tests/native_contract_tests_cpu/bounded_join_admission/support.rs +++ b/crates/algeon-query-engine/tests/native_contract_tests_cpu/bounded_join_admission/support.rs @@ -267,124 +267,114 @@ fn bounded_aggregate_state_and_max_rows_are_node_facts() { } #[test] -fn bounded_left_anti_build_edge_requires_dedup_aggregate() { - // EXCEPT ALL-shaped: the build side is a grouped aggregate with measures, - // which the bounded engine cannot do with multiplicity. - let key_schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); +fn bounded_left_anti_admits_join_output_probe_against_dedup_build() { + // `p LEFT JOIN (SELECT DISTINCT k ...) d ... WHERE d.k IS NULL` under a + // trailing count. The probe is an inner-join output that may repeat keys; + // native LeftAnti keeps every unmatched probe row, so neither edge depends + // on the probe being distinct. let mut builder = QueryPlanBuilder::new(); - let left = builder.add_source(SourceNode::LocalParquet(LocalParquetSource::new( - "/tmp/s3-anti-build-left.parquet", - Arc::clone(&key_schema), - vec!["k".into()], - ))); - let right = builder.add_source(SourceNode::LocalParquet(LocalParquetSource::new( - "/tmp/s3-anti-build-right.parquet", + let probe = compositional_inner_join_base(&mut builder); + let key_schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); + let probe_output_schema = Arc::new(Schema::new(vec![Field::new("left_id", DataType::Int32, false)])); + let one_schema = Arc::new(Schema::new(vec![Field::new("one", DataType::Int64, false)])); + let count_schema = Arc::new(Schema::new(vec![Field::new("count", DataType::Int64, false)])); + let build_source = builder.add_source(SourceNode::LocalParquet(LocalParquetSource::new( + "/tmp/left-anti-dedup-build.parquet", Arc::clone(&key_schema), vec!["k".into()], ))); - let grouped = builder - .add_aggregate( - right, - AggregateSpec::new( - vec![Expr::column("k")], - vec![AggregateExpr::new(AggregateFunction::Count, Expr::column("k"), "count")], - ), - Arc::new(Schema::new(vec![ - Field::new("k", DataType::Int32, false), - Field::new("count", DataType::Int64, false), - ])), - ) + let dedup = builder + .add_aggregate(build_source, AggregateSpec::new(vec![Expr::column("k")], Vec::new()), key_schema) .unwrap(); let anti = builder .add_join( - left, - grouped, + probe, + dedup, JoinSpec::new( JoinType::LEFT_ANTI, - vec![JoinKey::new(ColumnRef::input(0, "k", "k"), ColumnRef::input(1, "k", "k"))], + vec![JoinKey::new(ColumnRef::input(0, "left_id", "left_id"), ColumnRef::input(1, "k", "k"))], ) - .with_output_projection(vec![JoinOutputColumn::left(ColumnRef::input(0, "k", "k"), "k")]) + .with_output_projection(vec![JoinOutputColumn::left(ColumnRef::input(0, "left_id", "left_id"), "left_id")]) .with_build_side(JoinSide::Right), - Arc::clone(&key_schema), + probe_output_schema, ) .unwrap(); - let plan = builder.finish_with_host_sink(anti).unwrap(); - let analysis = plan.capability_analysis(rapids_analysis_context(PlanningPolicy::default())); + let count_input = builder + .add_projection(anti, projection_outputs(vec![Expr::literal(LiteralValue::Int64(1)).alias("one")], one_schema)) + .unwrap(); + let count = builder + .add_aggregate( + count_input, + AggregateSpec::new( + Vec::new(), + vec![AggregateExpr::new(AggregateFunction::Count, Expr::column("one"), "count")], + ), + count_schema, + ) + .unwrap(); + let plan = builder.finish_with_host_sink(count).unwrap(); - let build_edge = analysis - .edges - .iter() - .find(|edge| edge.producer == grouped && edge.consumer == EdgeConsumer::Node(anti)) - .expect("aggregate build edge exists"); - let CapabilityDecision::Rejected { reason, facts } = &build_edge.decision else { - panic!("grouped-measure anti build should reject the build edge"); - }; - assert_eq!(*reason, Reason::LeftAntiBuildRequiresDedupAggregate); - assert_eq!(facts.get("producer_aggregate_state").map(String::as_str), Some("plain")); + assert_left_anti_edges_satisfied_and_admitted(&plan, anti); } #[test] -fn bounded_left_anti_probe_edge_requires_dedup_proven_input() { - // EXCEPT DISTINCT-shaped but the probe side is a raw source: a trailing - // scalar count would count duplicates. +fn bounded_left_anti_admits_grouped_measure_build() { + // `NOT EXISTS (SELECT k, count(*) ... GROUP BY k)`: the anti build reads + // only key existence, so a measure beside the group key changes nothing. let key_schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); - let one_schema = Arc::new(Schema::new(vec![Field::new("one", DataType::Int64, false)])); - let scalar_schema = Arc::new(Schema::new(vec![Field::new("count", DataType::Int64, false)])); let mut builder = QueryPlanBuilder::new(); let left = builder.add_source(SourceNode::LocalParquet(LocalParquetSource::new( - "/tmp/s3-anti-probe-left.parquet", + "/tmp/left-anti-grouped-build-left.parquet", Arc::clone(&key_schema), vec!["k".into()], ))); let right = builder.add_source(SourceNode::LocalParquet(LocalParquetSource::new( - "/tmp/s3-anti-probe-right.parquet", + "/tmp/left-anti-grouped-build-right.parquet", Arc::clone(&key_schema), vec!["k".into()], ))); - let dedup = builder - .add_aggregate(right, AggregateSpec::new(vec![Expr::column("k")], Vec::new()), Arc::clone(&key_schema)) + let grouped = builder + .add_aggregate( + right, + AggregateSpec::new( + vec![Expr::column("k")], + vec![AggregateExpr::new(AggregateFunction::Count, Expr::column("k"), "count")], + ), + Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("count", DataType::Int64, false), + ])), + ) .unwrap(); let anti = builder .add_join( left, - dedup, + grouped, JoinSpec::new( JoinType::LEFT_ANTI, vec![JoinKey::new(ColumnRef::input(0, "k", "k"), ColumnRef::input(1, "k", "k"))], ) .with_output_projection(vec![JoinOutputColumn::left(ColumnRef::input(0, "k", "k"), "k")]) .with_build_side(JoinSide::Right), - Arc::clone(&key_schema), - ) - .unwrap(); - let count_input = builder - .add_projection(anti, projection_outputs(vec![Expr::literal(LiteralValue::Int64(1)).alias("one")], one_schema)) - .unwrap(); - let count = builder - .add_aggregate( - count_input, - AggregateSpec::new( - Vec::new(), - vec![AggregateExpr::new(AggregateFunction::Count, Expr::column("one"), "count")], - ), - Arc::clone(&scalar_schema), + key_schema, ) .unwrap(); - let projection = - builder.add_projection(count, projection_outputs(vec![Expr::column("count")], scalar_schema)).unwrap(); - let plan = builder.finish_with_host_sink(projection).unwrap(); - let analysis = plan.capability_analysis(rapids_analysis_context(PlanningPolicy::default())); + let plan = builder.finish_with_host_sink(anti).unwrap(); - let probe_edge = analysis - .edges - .iter() - .find(|edge| edge.producer == left && edge.consumer == EdgeConsumer::Node(anti)) - .expect("probe edge exists"); - let CapabilityDecision::Rejected { reason, facts } = &probe_edge.decision else { - panic!("non-dedup anti probe should reject the probe edge"); - }; - assert_eq!(*reason, Reason::LeftAntiProbeRequiresDedupProvenInput); - assert_eq!(facts.get("producer_aggregate_state").map(String::as_str), Some("none")); + assert_left_anti_edges_satisfied_and_admitted(&plan, anti); +} + +fn assert_left_anti_edges_satisfied_and_admitted(plan: &QueryPlan, anti: NodeId) { + let analysis = plan.capability_analysis(rapids_analysis_context(PlanningPolicy::default())); + let anti_edges = analysis.edges.iter().filter(|edge| edge.consumer == EdgeConsumer::Node(anti)).collect::>(); + assert_eq!(anti_edges.len(), 2, "the anti join has a probe and a build edge"); + for edge in anti_edges { + assert_eq!(edge.decision, CapabilityDecision::Satisfied, "{edge:?}"); + } + assert!( + bounded_pipeline_admission(plan).expect("bounded admission should not fail").is_some(), + "the left anti plan should be admitted" + ); } #[test] diff --git a/crates/algeon-query-planning/src/capability/aggregate.rs b/crates/algeon-query-planning/src/capability/aggregate.rs index c18814192..7ddac69f3 100644 --- a/crates/algeon-query-planning/src/capability/aggregate.rs +++ b/crates/algeon-query-planning/src/capability/aggregate.rs @@ -40,24 +40,13 @@ fn aggregate_expr_is_exact_distinct_count(expr: &AggregateExpr) -> bool { pub enum AggregateStateContract { /// No distinct or group-only dedup state is required. Plain, - /// SQL setop-style dedup over grouped keys that have no aggregate expressions. + /// Grouped keys with no aggregate expressions, the shape `SELECT DISTINCT` lowers to. DedupGroupOnly, /// Exact count-distinct state represented as `COUNT(DISTINCT col)` or /// native `NUnique(col)`. ExactDistinctCount, } -impl AggregateStateContract { - /// Stable diagnostic label for this aggregate state contract. - pub(crate) const fn as_str(self) -> &'static str { - match self { - Self::Plain => "plain", - Self::DedupGroupOnly => "dedup_group_only", - Self::ExactDistinctCount => "exact_distinct_count", - } - } -} - /// The state shape bounded admission may rely on for `spec`, or `None` when the /// spec has no mergeable state contract (see [`AggregateStateContract`]). pub(crate) fn aggregate_state_contract(spec: &AggregateSpec) -> Option { diff --git a/crates/algeon-query-planning/src/capability/edge.rs b/crates/algeon-query-planning/src/capability/edge.rs index 671acf3ba..2d91231e0 100644 --- a/crates/algeon-query-planning/src/capability/edge.rs +++ b/crates/algeon-query-planning/src/capability/edge.rs @@ -5,13 +5,9 @@ //! input, union position, sink materialization) and one decision. Rejections //! here are local to the edge; whole-plan admission lives in [`crate::admission`]. -use std::collections::BTreeSet; - use super::{ - JoinForm, JoinNode, JoinSide, JoinType, NodeId, PlanNode, QueryPlan, ResultSink, - aggregate::{AggregateStateContract, aggregate_state_contract}, - classify_join, - graph::{consumer_nodes, join_side_index, opposite_join_side}, + JoinSide, NodeId, PlanNode, QueryPlan, ResultSink, + graph::consumer_nodes, model::{ CapabilityDecision, EdgeConsumer, EdgeContract, EdgeContractKind, InputRole, ResultSinkKind, VectorInputRole, }, @@ -19,9 +15,12 @@ use super::{ rejection::Reason, }; -/// One edge per input slot of `kind`, in input order. The left-anti dedup -/// check refines only an edge whose base decision is satisfied, so a missing -/// producer keeps its own, more local, reason. +/// One edge per input slot of `kind`, in input order. +/// +/// Edge analysis infers no set operation from the producers around a join. A +/// native left anti join keeps every probe row with no build match, duplicates +/// included, and reads only key existence from its build input, so its edges +/// carry the base decision like every other join edge. pub(super) fn analyze_input_edges(plan: &QueryPlan, consumer_id: NodeId, kind: &PlanNode) -> Vec { kind.input_ids() .iter() @@ -31,111 +30,11 @@ pub(super) fn analyze_input_edges(plan: &QueryPlan, consumer_id: NodeId, kind: & let input_role = input_role(kind, input_index); let contract_kinds = edge_contract_kinds(plan, *producer, consumer, kind, input_role); let decision = edge_decision(plan, *producer, consumer); - let decision = match (&decision, kind) { - (CapabilityDecision::Satisfied, PlanNode::Join(join)) - if matches!(join.spec.join_type, JoinType::LeftAnti { .. }) => - { - left_anti_dedup_edge_decision(plan, *producer, input_role, join) - }, - _ => decision, - }; EdgeContract { producer: *producer, consumer, contract_kinds, input_role, decision } }) .collect() } -/// Set-operation guard for left-anti joins. -/// -/// An anti build against a dedup aggregate claims EXCEPT DISTINCT, so an -/// aggregate build must be one: a grouped aggregate with measures or -/// distinct state is EXCEPT ALL, which the bounded engine cannot do with -/// multiplicity. Sources, filters, and joins keep ordinary anti semantics, -/// which the engine executes directly, so only aggregate builds are policed. -/// For an except-shaped join the probe side must be dedup-proven, or a -/// trailing scalar count would count duplicates. -fn left_anti_dedup_edge_decision( - plan: &QueryPlan, - producer: NodeId, - input_role: InputRole, - join: &JoinNode, -) -> CapabilityDecision { - let producer_state = plan.node(producer).and_then(|node| match &node.kind { - PlanNode::Aggregate(aggregate) => aggregate_state_contract(&aggregate.spec), - _ => None, - }); - match input_role { - InputRole::JoinBuildInput { .. } => { - if matches!(plan.node(producer).map(|node| &node.kind), Some(PlanNode::Aggregate(_))) - && producer_state != Some(AggregateStateContract::DedupGroupOnly) - { - return CapabilityDecision::rejected( - Reason::LeftAntiBuildRequiresDedupAggregate, - [ - ("producer", producer.to_string()), - ( - "producer_aggregate_state", - producer_state.map_or("none", AggregateStateContract::as_str).to_string(), - ), - ], - ); - } - CapabilityDecision::Satisfied - }, - InputRole::JoinProbeInput { .. } => { - let build_input = join.inputs[join_side_index(join.spec.build_side)]; - let build_dedup = plan.node(build_input).is_some_and(|node| match &node.kind { - PlanNode::Aggregate(aggregate) => { - aggregate_state_contract(&aggregate.spec) == Some(AggregateStateContract::DedupGroupOnly) - }, - _ => false, - }); - if build_dedup && !anti_probe_input_is_dedup_proven(plan, producer, &mut BTreeSet::new()) { - return CapabilityDecision::rejected( - Reason::LeftAntiProbeRequiresDedupProvenInput, - [ - ("producer", producer.to_string()), - ( - "producer_aggregate_state", - producer_state.map_or("none", AggregateStateContract::as_str).to_string(), - ), - ], - ); - } - CapabilityDecision::Satisfied - }, - _ => CapabilityDecision::Satisfied, - } -} - -/// Structural proof that an anti-join probe input is a distinct relation: -/// either a group-only dedup aggregate, or a keyed filtered left-anti join -/// whose own build input is a dedup aggregate and whose probe input is -/// dedup-proven. Consults no row estimates. -fn anti_probe_input_is_dedup_proven(plan: &QueryPlan, node_id: NodeId, visiting: &mut BTreeSet) -> bool { - if !visiting.insert(node_id) { - return false; - } - match plan.node(node_id).map(|node| &node.kind) { - Some(PlanNode::Aggregate(aggregate)) => { - aggregate_state_contract(&aggregate.spec) == Some(AggregateStateContract::DedupGroupOnly) - }, - Some(PlanNode::Join(join)) - if matches!(join.spec.join_type, JoinType::LeftAnti { .. }) - && matches!(classify_join(node_id, join), Ok(JoinForm::KeyedFilteredSemiAnti)) => - { - let build_input = join.inputs[join_side_index(join.spec.build_side)]; - let probe_input = join.inputs[join_side_index(opposite_join_side(join.spec.build_side))]; - plan.node(build_input).is_some_and(|node| match &node.kind { - PlanNode::Aggregate(aggregate) => { - aggregate_state_contract(&aggregate.spec) == Some(AggregateStateContract::DedupGroupOnly) - }, - _ => false, - }) && anti_probe_input_is_dedup_proven(plan, probe_input, visiting) - }, - _ => false, - } -} - /// Contract kinds one edge carries: producer facts first (source format, join /// materialization), then the consumer's role facts. The shared-producer kind /// is added later by [`mark_shared_producer_edge_contracts`], once every edge diff --git a/crates/algeon-query-planning/src/capability/graph.rs b/crates/algeon-query-planning/src/capability/graph.rs index a396bc089..693dd9736 100644 --- a/crates/algeon-query-planning/src/capability/graph.rs +++ b/crates/algeon-query-planning/src/capability/graph.rs @@ -84,10 +84,3 @@ pub(super) const fn join_side_index(side: JoinSide) -> usize { JoinSide::Right => 1, } } - -pub(super) const fn opposite_join_side(side: JoinSide) -> JoinSide { - match side { - JoinSide::Left => JoinSide::Right, - JoinSide::Right => JoinSide::Left, - } -} diff --git a/crates/algeon-query-planning/src/capability/model.rs b/crates/algeon-query-planning/src/capability/model.rs index 9fb4b11b3..3cd0e406d 100644 --- a/crates/algeon-query-planning/src/capability/model.rs +++ b/crates/algeon-query-planning/src/capability/model.rs @@ -210,10 +210,8 @@ pub enum ResidualRequirement { /// /// Every edge requires its producer to exist in the plan and to yield a /// GPU-resident cudf frame. `decision` carries `missing_edge_producer` when the -/// producer is absent, which plan construction already rules out, and the -/// left-anti dedup reasons when an anti join's aggregate build or probe is not -/// a proven distinct relation; every other edge is satisfied here and judged -/// by the compiled pipeline contracts. +/// producer is absent, which plan construction already rules out; every other +/// edge is satisfied here and judged by the compiled pipeline contracts. #[derive(Debug, Clone, PartialEq, Eq)] pub struct EdgeContract { /// Producer node for this dependency. diff --git a/crates/algeon-query-planning/src/capability/rejection.rs b/crates/algeon-query-planning/src/capability/rejection.rs index a99ca1f9d..e5aa244f8 100644 --- a/crates/algeon-query-planning/src/capability/rejection.rs +++ b/crates/algeon-query-planning/src/capability/rejection.rs @@ -81,10 +81,6 @@ pub enum Reason { UnsupportedJoinResidualExpressionCapability, /// A sort-key expression cannot be evaluated by the backend. UnsupportedSortExpressionCapability, - /// The anti-join build lacks the required deduplicating aggregate. - LeftAntiBuildRequiresDedupAggregate, - /// The anti-join probe lacks proof that its input is deduplicated. - LeftAntiProbeRequiresDedupProvenInput, /// The producer referenced by an execution edge is absent from the plan. MissingEdgeProducer, /// A shared aggregate producer has a consumer outside the aggregate contract. @@ -220,8 +216,6 @@ impl Reason { Self::UnsupportedAggregateGroupExpressionCapability => "unsupported_aggregate_group_expression_capability", Self::UnsupportedJoinResidualExpressionCapability => "unsupported_join_residual_expression_capability", Self::UnsupportedSortExpressionCapability => "unsupported_sort_expression_capability", - Self::LeftAntiBuildRequiresDedupAggregate => "left_anti_build_requires_dedup_aggregate", - Self::LeftAntiProbeRequiresDedupProvenInput => "left_anti_probe_requires_dedup_proven_input", Self::MissingEdgeProducer => "missing_edge_producer", Self::AggregateSharedProducerConsumerNotAggregate => "aggregate_shared_producer_consumer_not_aggregate", Self::AggregateSharedProducerInputLineageUnresolved => "aggregate_shared_producer_input_lineage_unresolved", @@ -317,8 +311,6 @@ impl Reason { CapabilityRejectionCategory::UnsupportedExpressionFamily }, Self::UnsupportedSortExpressionCapability => CapabilityRejectionCategory::UnsupportedExpressionFamily, - Self::LeftAntiBuildRequiresDedupAggregate => CapabilityRejectionCategory::MissingEdgeContract, - Self::LeftAntiProbeRequiresDedupProvenInput => CapabilityRejectionCategory::MissingEdgeContract, Self::MissingEdgeProducer => CapabilityRejectionCategory::MissingEdgeContract, Self::AggregateSharedProducerConsumerNotAggregate => CapabilityRejectionCategory::MissingSubgraphInvariant, Self::AggregateSharedProducerInputLineageUnresolved => { From 82c56d23e3d983888646e5f90147cef7791d22b4 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Thu, 22 Oct 2026 00:27:13 +0000 Subject: [PATCH 8/8] fix(ci): scan only website/docs in the default prose check The default scan read the gitignored docs/ and note/ trees, which git and CI never see, so local private files failed the full gate. Its self-test for a checkout without those trees ran the real checkout's script, which scans the checkout containing it, and so read the same private files. The default target is website/docs; explicitly named paths still work. The self-test runs a copy of the script inside a scratch checkout with violations under docs/ and note/ and one under website/docs/. Co-Authored-By: Claude Opus 5.5 --- scripts/ci/check_prose_slop.sh | 19 +++++-------------- scripts/ci/tests/check_prose_slop_tests.sh | 21 ++++++++++++++++----- 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/scripts/ci/check_prose_slop.sh b/scripts/ci/check_prose_slop.sh index d4151f503..00a512319 100755 --- a/scripts/ci/check_prose_slop.sh +++ b/scripts/ci/check_prose_slop.sh @@ -23,8 +23,9 @@ USAGE = """Usage: check_prose_slop.sh [OPTIONS] [PATHS...] Checks markdown documentation for AI writing slop (em-dashes and banned phrases) and validates YAML front matter syntax. -When no paths or git flags are given, defaults to scanning: - docs/ note/ website/docs/ +When no paths or git flags are given, defaults to scanning website/docs/. +The gitignored docs/ and note/ trees are outside git and CI and are never +scanned by default. Options: --staged Check files currently staged in git @@ -33,14 +34,9 @@ Options: -h, --help Show this help message """ -DEFAULT_TARGETS = ["docs", "note", "website/docs"] +DEFAULT_TARGETS = ["website/docs"] DOC_EXTENSIONS = {".md", ".mdx", ".markdown", ".txt"} -# Exclude self-referential policy / plan documents by default -DEFAULT_EXCLUDED_PATHS = { - "docs/plan/2026-09-20-de-ai-slop-elimination-and-prevention-plan.md", -} - BANNED_PATTERNS = [ ("em-dash ('—')", re.compile(r"\u2014")), ("banned phrase 'delve'", re.compile(r"\bdelve(?:s|d|ing)?\b", re.IGNORECASE)), @@ -106,10 +102,7 @@ def collect_files(mode, targets, excludes): if p.suffix.lower() in DOC_EXTENSIONS and p.is_file(): files.append(p) else: - # docs/ and note/ are gitignored private trees, so a default target may be - # absent from a checkout; explicitly named paths must exist. - search_roots = targets if targets else [target for target in DEFAULT_TARGETS if Path(target).exists()] - for target in search_roots: + for target in targets or DEFAULT_TARGETS: p = Path(target) if p.is_file(): files.append(p) @@ -125,8 +118,6 @@ def collect_files(mode, targets, excludes): filtered = [] for f in sorted(set(files)): rel = f.as_posix() - if rel in DEFAULT_EXCLUDED_PATHS: - continue if any(exc.search(rel) for exc in excludes): continue filtered.append(f) diff --git a/scripts/ci/tests/check_prose_slop_tests.sh b/scripts/ci/tests/check_prose_slop_tests.sh index 75efc03fa..a8b5bd07c 100755 --- a/scripts/ci/tests/check_prose_slop_tests.sh +++ b/scripts/ci/tests/check_prose_slop_tests.sh @@ -32,13 +32,24 @@ if ! "${CHECK_SCRIPT}" "${scratch}/clean.md" >/dev/null 2>&1; then fail "clean file unexpectedly failed check_prose_slop" fi -# 1b. Default targets that are absent (gitignored private trees) are skipped. -mkdir -p "${scratch}/checkout/website/docs" +# 1b. The default scan reads website/docs/ only; the gitignored docs/ and +# note/ trees are never scanned. The script scans the checkout that contains +# it, so the scratch checkout gets its own copy. +mkdir -p "${scratch}/checkout/website/docs" "${scratch}/checkout/scripts/ci" \ + "${scratch}/checkout/docs" "${scratch}/checkout/note" cp "${scratch}/clean.md" "${scratch}/checkout/website/docs/clean.md" -if ! (cd "${scratch}/checkout" && "${CHECK_SCRIPT}" >/dev/null 2>&1); then - fail "default scan failed in a checkout without docs/ or note/" +printf 'A private draft \u2014 never scanned.\n' >"${scratch}/checkout/docs/draft.md" +printf 'A private note \u2014 never scanned.\n' >"${scratch}/checkout/note/draft.md" +cp "${CHECK_SCRIPT}" "${scratch}/checkout/scripts/ci/check_prose_slop.sh" +checkout_script="${scratch}/checkout/scripts/ci/check_prose_slop.sh" +if ! (cd "${scratch}/checkout" && "${checkout_script}" >/dev/null 2>&1); then + fail "default scan read docs/ or note/" fi -if (cd "${scratch}/checkout" && "${CHECK_SCRIPT}" missing.md >/dev/null 2>&1); then +printf 'A published page \u2014 scanned.\n' >"${scratch}/checkout/website/docs/page.md" +if (cd "${scratch}/checkout" && "${checkout_script}" >/dev/null 2>&1); then + fail "default scan missed a violation under website/docs/" +fi +if (cd "${scratch}/checkout" && "${checkout_script}" missing.md >/dev/null 2>&1); then fail "an explicitly named missing path unexpectedly passed" fi