From 2dd0437da0a1206e723ac7435a3b2161f3b2b458 Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Mon, 19 Oct 2026 19:11:30 +0000 Subject: [PATCH 1/5] feat: evaluate single-input join condition parts before the join A join condition that reads both inputs stays on the join. Each part that reads one input and that the cuDF join AST cannot evaluate (LIKE, string functions, IS DISTINCT FROM) becomes a computed column of that input before the join, as an expression join key does. Only parts that cannot fail move, because DataFusion evaluates the condition on key-matched pairs alone. Statements that write LOWER(name) LIKE in a filter above an inner join are native under native_required with this change; 0.1.2 rejected them. Co-Authored-By: Claude Fable 5.1 --- .../expression_contract/join_residual.rs | 436 ++++++++-------- .../src/native/expression_contract/mod.rs | 4 +- .../src/native/expression_contract/tests.rs | 231 ++++++++- .../src/native/lowering/expr.rs | 17 +- .../src/native/lowering/join.rs | 6 +- .../src/native/lowering/join_residual.rs | 472 ++++++------------ .../src/native/lowering_tests.rs | 10 +- .../lowering_tests/full_gpu_gap_analysis.rs | 69 +++ .../join_shapes/residual_variants.rs | 68 ++- .../scalar_expressions.rs | 100 +++- website/docs/design/native-lowering.mdx | 8 + .../reference/native-evaluation-semantics.mdx | 20 +- website/docs/reference/native-sql-support.mdx | 2 +- 13 files changed, 877 insertions(+), 566 deletions(-) diff --git a/crates/algeon-datafusion/src/native/expression_contract/join_residual.rs b/crates/algeon-datafusion/src/native/expression_contract/join_residual.rs index 909995c75..7c8e56d68 100644 --- a/crates/algeon-datafusion/src/native/expression_contract/join_residual.rs +++ b/crates/algeon-datafusion/src/native/expression_contract/join_residual.rs @@ -16,26 +16,34 @@ use query_engine::rapids_backend::rapids_backend; use query_planning::{datatypes::is_supported_ast_cast, validation::is_integral_type}; use super::{ - ExpressionCapability, ExpressionRejection, ExpressionRejectionSubject, ScalarFunctionKind, ValidatedPhysicalExpr, - binary_operator_subject, is_checked_temporal_arithmetic, is_decimal_division, physical_expr_data_type, - validate_physical_expr_contract, + ExpressionCapability, ExpressionRejection, ExpressionRejectionSubject, ValidatedPhysicalExpr, + binary_operator_subject, contains_fallible_operation, fallible_operand, is_checked_temporal_arithmetic, + is_decimal_division, physical_expr_data_type, validate_physical_expr_contract, }; -/// Representation selected by the join-residual-specific cast contract. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(in crate::native) enum JoinResidualCastMode { - Ast, - Materialize, -} - /// Proof that a join filter satisfies both the general physical-expression -/// contract and the stricter native join residual vocabulary. +/// contract and the native join residual contract. +/// +/// The residual keeps the operators the cuDF join AST evaluates. Each other +/// subexpression that reads one input and cannot fail is a +/// [`JoinResidualMaterializationSite`]: lowering evaluates it on every row of +/// that input before the join and the residual reads the result as a column. #[derive(Debug)] pub(in crate::native) struct ValidatedJoinResidual<'a> { filter: &'a JoinFilter, expression: ValidatedPhysicalExpr, left_schema: SchemaRef, right_schema: SchemaRef, + materializations: Vec, +} + +/// A residual subexpression that lowering evaluates on one join input before +/// the join. Sites are distinct `Arc`s in residual pre-order, and lowering +/// substitutes each by pointer identity. +#[derive(Debug)] +pub(in crate::native) struct JoinResidualMaterializationSite { + pub(in crate::native) expr: Arc, + pub(in crate::native) side: JoinSide, } impl<'a> ValidatedJoinResidual<'a> { @@ -56,15 +64,11 @@ impl<'a> ValidatedJoinResidual<'a> { } pub(in crate::native) fn contains_decimal_division(&self) -> bool { - fn contains(expr: &Arc, schema: &SchemaRef) -> bool { - expr.downcast_ref::().is_some_and(|binary| is_decimal_division(binary, schema)) - || expr.children().into_iter().any(|child| contains(child, schema)) - } - contains(self.filter.expression(), self.filter.schema()) + contains_decimal_division(self.filter.expression(), self.filter.schema()) } - pub(in crate::native) fn single_side(&self, expr: &Arc) -> Result { - join_filter_expr_single_side(expr, self.filter) + pub(in crate::native) fn materializations(&self) -> &[JoinResidualMaterializationSite] { + &self.materializations } } @@ -74,17 +78,19 @@ pub(in crate::native) fn validate_join_residual<'a>( right_schema: &SchemaRef, ) -> Result, ExpressionRejection> { let expression = validate_physical_expr_contract(filter.expression(), filter.schema())?; - validate_join_residual_expr(filter.expression(), filter)?; let residual = filter.expression(); + let mut shape = ResidualShape { + filter, + post_join: contains_decimal_division(residual, filter.schema()), + materializations: Vec::new(), + }; + shape.visit(residual)?; validate_join_residual_side_indices(residual, filter, left_schema, right_schema) .map_err(|rejection| rejection.at_expression(residual))?; - let mut side = None; - let mut saw_left = false; - let mut saw_right = false; - collect_join_filter_expr_sides(residual, filter, &mut side, &mut saw_left, &mut saw_right, false) - .map_err(|rejection| rejection.at_expression(residual))?; - if !saw_left || !saw_right { + let (reads_left, reads_right) = + join_filter_expr_sides(residual, filter).map_err(|rejection| rejection.at_expression(residual))?; + if !reads_left || !reads_right { return Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualSides) .at_expression(residual)); } @@ -94,9 +100,177 @@ pub(in crate::native) fn validate_join_residual<'a>( expression, left_schema: Arc::clone(left_schema), right_schema: Arc::clone(right_schema), + materializations: shape.materializations, }) } +fn contains_decimal_division(expr: &Arc, schema: &SchemaRef) -> bool { + expr.downcast_ref::().is_some_and(|binary| is_decimal_division(binary, schema)) + || expr.children().into_iter().any(|child| contains_decimal_division(child, schema)) +} + +/// Walk of a residual that admits its shape and records its materialization +/// sites. +struct ResidualShape<'f> { + filter: &'f JoinFilter, + /// Whether lowering runs the residual as the checked filter over + /// key-matched rows, which it does for a decimal `/`. That filter + /// evaluates a fallible operation only on matched pairs, as DataFusion + /// does, and the operation stays in the residual. + post_join: bool, + materializations: Vec, +} + +impl ResidualShape<'_> { + fn visit(&mut self, expr: &Arc) -> Result<(), ExpressionRejection> { + self.visit_shape(expr).map_err(|rejection| rejection.at_expression(expr)) + } + + fn visit_children(&mut self, expr: &Arc) -> Result<(), ExpressionRejection> { + expr.children().into_iter().try_for_each(|child| self.visit(child)) + } + + fn visit_shape(&mut self, expr: &Arc) -> Result<(), ExpressionRejection> { + if expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() { + return Ok(()); + } + let schema = self.filter.schema(); + if is_inline_operator(expr, schema) { + return self.visit_children(expr); + } + let Some(side) = join_filter_expr_single_side(expr, self.filter) else { + return self.visit_cross_input(expr); + }; + // A pre-join value is computed for every row of its input, including + // rows that match no key and that DataFusion never filters. Only an + // expression that cannot fail moves before the join. + if !contains_fallible_operation(expr, schema) { + if !self.materializations.iter().any(|site| Arc::ptr_eq(&site.expr, expr)) { + self.materializations.push(JoinResidualMaterializationSite { expr: Arc::clone(expr), side }); + } + return Ok(()); + } + if self.post_join && !is_unchecked_ast_arithmetic(expr, schema) { + return self.visit_children(expr); + } + Err(fallible_operation_rejection(expr, schema)) + } + + /// Admits an expression over both inputs that the AST does not evaluate. + /// The planner evaluates an inner join residual with `CASE` or null-aware + /// comparison as a filter over key-matched pairs. + fn visit_cross_input(&mut self, expr: &Arc) -> Result<(), ExpressionRejection> { + let null_aware = expr + .downcast_ref::() + .is_some_and(|binary| matches!(binary.op(), Operator::IsDistinctFrom | Operator::IsNotDistinctFrom)); + if null_aware || expr.downcast_ref::().is_some() { + return self.visit_children(expr); + } + Err(cross_input_rejection(expr, self.filter.schema())) + } +} + +/// Whether the residual keeps `expr` as an operator over its operands: the +/// cuDF join AST evaluates it, or, for a decimal `/`, the checked post-join +/// filter does. The AST has no concatenation, null-aware comparison or bit +/// shift. +fn is_inline_operator(expr: &Arc, schema: &SchemaRef) -> bool { + if let Some(binary) = expr.downcast_ref::() { + return !matches!( + binary.op(), + Operator::StringConcat + | Operator::IsDistinctFrom + | Operator::IsNotDistinctFrom + | Operator::BitwiseShiftLeft + | Operator::BitwiseShiftRight + ) && !is_unchecked_ast_arithmetic(expr, schema); + } + if let Some(cast) = expr.downcast_ref::() { + return physical_expr_data_type(cast.expr(), schema) + .is_ok_and(|source| is_supported_ast_cast(&source, cast.cast_type(), rapids_backend())); + } + expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() + || expr.downcast_ref::().is_some() +} + +/// Whether `expr` is arithmetic the cuDF AST would evaluate without the checks +/// DataFusion applies. The AST has no zero-divisor or overflow checks and no +/// verified floating-point remainder, and it wraps timestamp arithmetic that +/// Arrow checks. Decimal division uses the checked post-join filter. +fn is_unchecked_ast_arithmetic(expr: &Arc, schema: &SchemaRef) -> bool { + let Some(binary) = expr.downcast_ref::() else { + return false; + }; + match binary.op() { + Operator::Modulo => true, + Operator::Plus | Operator::Minus => is_checked_temporal_arithmetic(binary, schema), + Operator::Divide => { + physical_expr_data_type(binary.left(), schema).map_or(true, |data_type| is_integral_type(&data_type)) + }, + _ => false, + } +} + +/// The rejection for `expr`, an expression over both inputs that neither the +/// AST nor the post-join filter evaluates. +fn cross_input_rejection(expr: &Arc, schema: &SchemaRef) -> ExpressionRejection { + if let Some(binary) = expr.downcast_ref::() { + return ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape) + .named_by_operator(binary_operator_subject(binary, schema)); + } + if let Some(cast) = expr.downcast_ref::() { + return cast_rejection(cast, schema); + } + if let Some(function) = expr.downcast_ref::() { + return ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualScalarFunction) + .with_subject(ExpressionRejectionSubject::ScalarFunction(function.name().to_string())); + } + ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape) +} + +/// The rejection for `expr`, a single-input expression whose evaluation can +/// fail, naming its first fallible operation. +fn fallible_operation_rejection(expr: &Arc, schema: &SchemaRef) -> ExpressionRejection { + let operation = first_fallible_operation(expr, schema).unwrap_or(expr); + let rejection = if let Some(binary) = operation.downcast_ref::() { + ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape) + .named_by_operator(binary_operator_subject(binary, schema)) + } else if let Some(cast) = operation.downcast_ref::() { + cast_rejection(cast, schema) + } else { + ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape) + }; + rejection.at_expression(operation) +} + +fn first_fallible_operation<'e>( + expr: &'e Arc, + schema: &SchemaRef, +) -> Option<&'e Arc> { + if fallible_operand(expr, schema).is_some() { + return Some(expr); + } + expr.children().into_iter().find_map(|child| first_fallible_operation(child, schema)) +} + +/// A cast the residual cannot evaluate. A cast to an AST target type reports +/// its type pair; the AST has no other cast target. +fn cast_rejection(cast: &CastExpr, schema: &SchemaRef) -> ExpressionRejection { + let rejection = if matches!(cast.cast_type(), DataType::Int64 | DataType::UInt64 | DataType::Float64) { + ExpressionRejection::unsupported_type(ExpressionCapability::CastPair) + } else { + ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualCastTarget) + }; + match physical_expr_data_type(cast.expr(), schema) { + Ok(source) => { + rejection.with_subject(ExpressionRejectionSubject::Cast { source, target: cast.cast_type().clone() }) + }, + Err(type_rejection) => type_rejection, + } +} + fn validate_join_residual_side_indices( expr: &Arc, filter: &JoinFilter, @@ -133,201 +307,31 @@ fn validate_join_residual_side_indices( Ok(()) } -pub(in crate::native) fn join_residual_cast_mode( - target_type: &DataType, -) -> Result { - if matches!(target_type, DataType::Int64 | DataType::UInt64 | DataType::Float64) { - return Ok(JoinResidualCastMode::Ast); - } - if matches!(target_type, DataType::Decimal128(_, _) | DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View) { - return Ok(JoinResidualCastMode::Materialize); - } - Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualCastTarget)) -} - -pub(in crate::native) fn validate_join_residual_scalar_function( - kind: ScalarFunctionKind, -) -> Result<(), ExpressionRejection> { - match kind { - ScalarFunctionKind::Substring | ScalarFunctionKind::Lower | ScalarFunctionKind::Upper => Ok(()), - _ => Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualScalarFunction)), - } -} - -fn validate_join_residual_expr(expr: &Arc, filter: &JoinFilter) -> Result<(), ExpressionRejection> { - validate_join_residual_shape(expr, filter).map_err(|rejection| rejection.at_expression(expr)) -} - -fn validate_join_residual_shape(expr: &Arc, filter: &JoinFilter) -> Result<(), ExpressionRejection> { - if expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() { - return Ok(()); - } - if let Some(binary) = expr.downcast_ref::() { - // The cudf AST has no zero-divisor or overflow checks and no verified - // floating-point remainder, and it wraps timestamp arithmetic that - // Arrow checks. Decimal division uses the checked GPU post-join - // evaluator, including CASE branch activation. - let ast_arithmetic_is_unchecked = match binary.op() { - Operator::Modulo => true, - Operator::Plus | Operator::Minus => is_checked_temporal_arithmetic(binary, filter.schema()), - Operator::Divide => physical_expr_data_type(binary.left(), filter.schema()) - .map_or(true, |data_type| is_integral_type(&data_type)), - _ => false, - }; - if ast_arithmetic_is_unchecked { - return Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape) - .named_by_operator(binary_operator_subject(binary, filter.schema()))); - } - validate_join_residual_expr(binary.left(), filter)?; - return validate_join_residual_expr(binary.right(), filter); - } - if let Some(not) = expr.downcast_ref::() { - return validate_join_residual_expr(not.arg(), filter); - } - if let Some(is_null) = expr.downcast_ref::() { - return validate_join_residual_expr(is_null.arg(), filter); - } - if let Some(is_not_null) = expr.downcast_ref::() { - return validate_join_residual_expr(is_not_null.arg(), filter); - } - if let Some(cast) = expr.downcast_ref::() { - if matches!(join_residual_cast_mode(cast.cast_type())?, JoinResidualCastMode::Materialize) { - join_filter_expr_single_side(expr, filter)?; - } else { - // The AST cast wraps or saturates where DataFusion fails the query, - // and it has no cast over temporal values. - let source_type = physical_expr_data_type(cast.expr(), filter.schema())?; - if !is_supported_ast_cast(&source_type, cast.cast_type(), rapids_backend()) { - return Err(ExpressionRejection::unsupported_type(ExpressionCapability::CastPair).with_subject( - ExpressionRejectionSubject::Cast { source: source_type, target: cast.cast_type().clone() }, - )); - } - } - return validate_join_residual_expr(cast.expr(), filter); - } - if let Some(in_list) = expr.downcast_ref::() { - validate_join_residual_expr(in_list.expr(), filter)?; - for item in in_list.list() { - validate_join_residual_expr(item, filter)?; - } - return Ok(()); - } - if let Some(case) = expr.downcast_ref::() { - if let Some(base) = case.expr() { - validate_join_residual_expr(base, filter)?; - } - for (when, then) in case.when_then_expr() { - validate_join_residual_expr(when, filter)?; - validate_join_residual_expr(then, filter)?; - } - if let Some(else_expr) = case.else_expr() { - validate_join_residual_expr(else_expr, filter)?; - } - return Ok(()); - } - if let Some(function) = expr.downcast_ref::() { - validate_join_residual_scalar_function(ScalarFunctionKind::from_scalar_function(function)?).map_err( - |rejection| rejection.with_subject(ExpressionRejectionSubject::ScalarFunction(function.name().to_string())), - )?; - join_filter_expr_single_side(expr, filter)?; - for arg in function.args() { - validate_join_residual_expr(arg, filter)?; - } - return Ok(()); - } - Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape)) -} - -fn join_filter_expr_single_side( - expr: &Arc, - filter: &JoinFilter, -) -> Result { - let mut side = None; - let mut saw_left = false; - let mut saw_right = false; - collect_join_filter_expr_sides(expr, filter, &mut side, &mut saw_left, &mut saw_right, true)?; - side.ok_or_else(|| ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualSides)) -} - -fn collect_join_filter_expr_sides( +/// Which join inputs `expr` reads, as `(left, right)`. +fn join_filter_expr_sides( expr: &Arc, filter: &JoinFilter, - side: &mut Option, - saw_left: &mut bool, - saw_right: &mut bool, - enforce_single_side: bool, -) -> Result<(), ExpressionRejection> { +) -> Result<(bool, bool), ExpressionRejection> { if let Some(column) = expr.downcast_ref::() { - let column_index = filter - .column_indices() - .get(column.index()) - .ok_or_else(|| ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape))?; - match column_index.side { - JoinSide::Left | JoinSide::Right => { - *saw_left |= column_index.side == JoinSide::Left; - *saw_right |= column_index.side == JoinSide::Right; - if enforce_single_side && side.is_some_and(|existing| existing != column_index.side) { - return Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualSides)); - } - *side = Some(column_index.side); - return Ok(()); + return match filter.column_indices().get(column.index()).map(|column_index| column_index.side) { + Some(JoinSide::Left) => Ok((true, false)), + Some(JoinSide::Right) => Ok((false, true)), + Some(JoinSide::None) | None => { + Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape)) }, - JoinSide::None => { - return Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape)); - }, - } - } - if expr.downcast_ref::().is_some() { - return Ok(()); - } - if let Some(binary) = expr.downcast_ref::() { - collect_join_filter_expr_sides(binary.left(), filter, side, saw_left, saw_right, enforce_single_side)?; - return collect_join_filter_expr_sides(binary.right(), filter, side, saw_left, saw_right, enforce_single_side); - } - if let Some(not) = expr.downcast_ref::() { - return collect_join_filter_expr_sides(not.arg(), filter, side, saw_left, saw_right, enforce_single_side); - } - if let Some(is_null) = expr.downcast_ref::() { - return collect_join_filter_expr_sides(is_null.arg(), filter, side, saw_left, saw_right, enforce_single_side); - } - if let Some(is_not_null) = expr.downcast_ref::() { - return collect_join_filter_expr_sides( - is_not_null.arg(), - filter, - side, - saw_left, - saw_right, - enforce_single_side, - ); - } - if let Some(cast) = expr.downcast_ref::() { - return collect_join_filter_expr_sides(cast.expr(), filter, side, saw_left, saw_right, enforce_single_side); - } - if let Some(in_list) = expr.downcast_ref::() { - collect_join_filter_expr_sides(in_list.expr(), filter, side, saw_left, saw_right, enforce_single_side)?; - for item in in_list.list() { - collect_join_filter_expr_sides(item, filter, side, saw_left, saw_right, enforce_single_side)?; - } - return Ok(()); - } - if let Some(case) = expr.downcast_ref::() { - if let Some(base) = case.expr() { - collect_join_filter_expr_sides(base, filter, side, saw_left, saw_right, enforce_single_side)?; - } - for (when, then) in case.when_then_expr() { - collect_join_filter_expr_sides(when, filter, side, saw_left, saw_right, enforce_single_side)?; - collect_join_filter_expr_sides(then, filter, side, saw_left, saw_right, enforce_single_side)?; - } - if let Some(else_expr) = case.else_expr() { - collect_join_filter_expr_sides(else_expr, filter, side, saw_left, saw_right, enforce_single_side)?; - } - return Ok(()); + }; } - if let Some(function) = expr.downcast_ref::() { - for arg in function.args() { - collect_join_filter_expr_sides(arg, filter, side, saw_left, saw_right, enforce_single_side)?; - } - return Ok(()); + expr.children().into_iter().try_fold((false, false), |(left, right), child| { + let (child_left, child_right) = join_filter_expr_sides(child, filter)?; + Ok((left || child_left, right || child_right)) + }) +} + +/// The one input `expr` reads, or `None` when it reads both or neither. +fn join_filter_expr_single_side(expr: &Arc, filter: &JoinFilter) -> Option { + match join_filter_expr_sides(expr, filter) { + Ok((true, false)) => Some(JoinSide::Left), + Ok((false, true)) => Some(JoinSide::Right), + _ => None, } - Err(ExpressionRejection::unsupported_expression(ExpressionCapability::JoinResidualShape)) } diff --git a/crates/algeon-datafusion/src/native/expression_contract/mod.rs b/crates/algeon-datafusion/src/native/expression_contract/mod.rs index 63fa221f5..41a6b0b09 100644 --- a/crates/algeon-datafusion/src/native/expression_contract/mod.rs +++ b/crates/algeon-datafusion/src/native/expression_contract/mod.rs @@ -14,9 +14,7 @@ pub(in crate::native) use capability::{ is_decimal128_to_float32_cast_via_float64, is_list_carry_type, is_narrow_decimal_carry, is_numeric_or_decimal_type, is_supported_to_timestamp_seconds_cast, is_supported_try_cast_pair, native_aggregate_function_kind, }; -pub(in crate::native) use join_residual::{ - JoinResidualCastMode, ValidatedJoinResidual, join_residual_cast_mode, validate_join_residual, -}; +pub(in crate::native) use join_residual::{ValidatedJoinResidual, validate_join_residual}; pub(crate) use rejection::{ ExpressionCapability, ExpressionRejection, ExpressionRejectionKind, ExpressionRejectionSubject, }; diff --git a/crates/algeon-datafusion/src/native/expression_contract/tests.rs b/crates/algeon-datafusion/src/native/expression_contract/tests.rs index 736fa0132..09484193c 100644 --- a/crates/algeon-datafusion/src/native/expression_contract/tests.rs +++ b/crates/algeon-datafusion/src/native/expression_contract/tests.rs @@ -26,8 +26,7 @@ use query_planning::expr::{ use super::{ ExpressionCapability, ExpressionRejectionSubject, ScalarFunctionKind, SupportedBinaryOperator, fallible_operand, - join_residual::validate_join_residual_scalar_function, join_residual_cast_mode, native_aggregate_function_kind, - validate_join_residual, validate_physical_expr_contract, + native_aggregate_function_kind, validate_join_residual, validate_physical_expr_contract, }; use crate::{capability::CapabilityReason, native::lowering::lower_physical_expr_for_test}; @@ -59,22 +58,227 @@ fn binary_capability_rejection_is_structured_and_preserves_public_reason() { assert_eq!(rejection.capability_reason(), CapabilityReason::UnsupportedExpression); } +/// The left input reads `name`, `n` and `t`, the right input `r` and `label`; +/// the filter schema lists them in that order. +fn join_residual_schemas() -> (Arc, Arc, Arc) { + let timestamp = DataType::Timestamp(TimeUnit::Millisecond, None); + let left = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("n", DataType::Int32, true), + Field::new("t", timestamp.clone(), true), + ])); + let right = + Arc::new(Schema::new(vec![Field::new("r", DataType::Int32, true), Field::new("label", DataType::Utf8, true)])); + let filter = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("n", DataType::Int32, true), + Field::new("t", timestamp, true), + Field::new("r", DataType::Int32, true), + Field::new("label", DataType::Utf8, true), + ])); + (left, right, filter) +} + +fn join_residual_filter(expr: Arc, filter_schema: &Arc) -> JoinFilter { + JoinFilter::new( + expr, + vec![ + ColumnIndex { index: 0, side: JoinSide::Left }, + ColumnIndex { index: 1, side: JoinSide::Left }, + ColumnIndex { index: 2, side: JoinSide::Left }, + ColumnIndex { index: 0, side: JoinSide::Right }, + ColumnIndex { index: 1, side: JoinSide::Right }, + ], + Arc::clone(filter_schema), + ) +} + #[test] -fn join_residual_cast_targets_are_owned_by_the_contract() { - assert_eq!(join_residual_cast_mode(&DataType::Int64).unwrap(), super::JoinResidualCastMode::Ast); - assert_eq!( - join_residual_cast_mode(&DataType::Decimal128(12, 2)).unwrap(), - super::JoinResidualCastMode::Materialize +fn join_residual_moves_single_input_expressions_outside_the_ast_before_the_join() { + let (left_schema, right_schema, schema) = join_residual_schemas(); + let column = |name: &str| col(name, &schema).unwrap(); + let text = |value: &str| lit(ScalarValue::Utf8(Some(value.to_string()))); + let int = |value: i32| lit(ScalarValue::Int32(Some(value))); + let like = |negated: bool, expr: Arc| -> Arc { + Arc::new(LikeExpr::new(negated, false, expr, text("%bar%"))) + }; + let abs = scalar_function("abs", math::abs(), vec![column("n")], DataType::Int32); + let single_input_case: Arc = Arc::new( + CaseExpr::try_new( + None, + vec![(binary(column("n"), Operator::Gt, int(0), &schema).unwrap(), column("n"))], + Some(int(0)), + ) + .unwrap(), ); - let rejection = join_residual_cast_mode(&DataType::Date32).unwrap_err(); + let or_right_positive = |site: Arc| { + let right_positive = binary(column("r"), Operator::Gt, int(0), &schema).unwrap(); + (Arc::clone(&site), binary(site, Operator::Or, right_positive, &schema).unwrap()) + }; + let compared = |site: Arc, operator: Operator, other: Arc| { + (Arc::clone(&site), binary(site, operator, other, &schema).unwrap()) + }; + let cases = vec![ + ("like", or_right_positive(like(false, column("name")))), + ("not like", or_right_positive(like(true, column("name")))), + ( + "like over lower", + or_right_positive(like( + false, + scalar_function("lower", string::lower(), vec![column("name")], DataType::Utf8), + )), + ), + ( + "is distinct from", + or_right_positive(binary(column("n"), Operator::IsDistinctFrom, int(5), &schema).unwrap()), + ), + ("abs", compared(abs, Operator::Gt, column("r"))), + ( + "concatenation", + compared( + binary(column("name"), Operator::StringConcat, text("!"), &schema).unwrap(), + Operator::Eq, + column("label"), + ), + ), + ("try cast", compared(Arc::new(TryCastExpr::new(column("name"), DataType::Int32)), Operator::Eq, column("r"))), + ( + "timestamp to bigint", + compared( + Arc::new(CastExpr::new(column("t"), DataType::Int64, None)), + Operator::Gt, + Arc::new(CastExpr::new(column("r"), DataType::Int64, None)), + ), + ), + ( + "integer division by a literal", + compared(binary(column("n"), Operator::Divide, int(2), &schema).unwrap(), Operator::Eq, column("r")), + ), + ("single-input case", compared(single_input_case, Operator::Eq, column("r"))), + ]; + for (label, (site, residual)) in cases { + let filter = join_residual_filter(residual, &schema); + let validated = validate_join_residual(&filter, &left_schema, &right_schema) + .unwrap_or_else(|error| panic!("{label} must validate: {error:?}")); + let sites = validated.materializations(); + assert_eq!(sites.len(), 1, "{label}: {sites:?}"); + assert!(Arc::ptr_eq(&sites[0].expr, &site), "{label} materializes {:?}", sites[0].expr); + assert_eq!(sites[0].side, JoinSide::Left, "{label}"); + } + + let ast_only = binary( + binary(binary(column("n"), Operator::Plus, int(1), &schema).unwrap(), Operator::Gt, column("r"), &schema) + .unwrap(), + Operator::And, + Arc::new(NotExpr::new(Arc::new(IsNullExpr::new(column("n"))))), + &schema, + ) + .unwrap(); + let filter = join_residual_filter(ast_only, &schema); + let validated = validate_join_residual(&filter, &left_schema, &right_schema).unwrap(); + assert!(validated.materializations().is_empty(), "the AST evaluates arithmetic, comparison and NOT"); +} + +#[test] +fn join_residual_keeps_operations_that_can_fail_out_of_pre_join_columns() { + let (left_schema, right_schema, schema) = join_residual_schemas(); + let column = |name: &str| col(name, &schema).unwrap(); + let to_decimal = |name: &str, precision: u8| -> Arc { + Arc::new(CastExpr::new(column(name), DataType::Decimal128(precision, 0), None)) + }; + let narrowing_subject = + ExpressionRejectionSubject::Cast { source: DataType::Int32, target: DataType::Decimal128(4, 0) }; + + // An Int32 does not always fit four digits, and a pre-join column would + // also cast rows that match no key. + let narrowing = binary(to_decimal("n", 4), Operator::Eq, to_decimal("r", 10), &schema).unwrap(); + let rejection = + validate_join_residual(&join_residual_filter(narrowing, &schema), &left_schema, &right_schema).unwrap_err(); assert_eq!(rejection.capability(), ExpressionCapability::JoinResidualCastTarget); + assert_eq!(rejection.subject(), &narrowing_subject); + assert_eq!(rejection.sql(), Some("CAST(n AS Decimal128(4, 0))")); + + let guarded: Arc = Arc::new( + CaseExpr::try_new( + None, + vec![( + binary(column("n"), Operator::Gt, lit(ScalarValue::Int32(Some(0))), &schema).unwrap(), + to_decimal("n", 4), + )], + None, + ) + .unwrap(), + ); + let guarded = binary(guarded, Operator::Eq, to_decimal("r", 10), &schema).unwrap(); + let rejection = + validate_join_residual(&join_residual_filter(guarded, &schema), &left_schema, &right_schema).unwrap_err(); + assert_eq!(rejection.subject(), &narrowing_subject, "the rejection names the cast inside the CASE"); + + // The checked post-join filter of a decimal division evaluates matched + // pairs only, so the narrowing cast stays in the residual there. + let left_cast = to_decimal("n", 4); + let right_cast = to_decimal("r", 10); + let quotient = binary(Arc::clone(&left_cast), Operator::Divide, Arc::clone(&right_cast), &schema).unwrap(); + let post_join = binary(quotient, Operator::Gt, lit(ScalarValue::Decimal128(Some(1), 10, 0)), &schema).unwrap(); + let filter = join_residual_filter(post_join, &schema); + let validated = validate_join_residual(&filter, &left_schema, &right_schema).unwrap(); + let sites = validated.materializations(); + assert_eq!(sites.len(), 1, "{sites:?}"); + assert!(Arc::ptr_eq(&sites[0].expr, &right_cast), "only the cast that always fits moves before the join"); } #[test] -fn join_residual_scalar_functions_add_a_stricter_contract_layer() { - validate_join_residual_scalar_function(ScalarFunctionKind::Substring).unwrap(); - let rejection = validate_join_residual_scalar_function(ScalarFunctionKind::Abs).unwrap_err(); +fn join_residual_rejects_cross_input_expressions_neither_evaluator_runs() { + let (left_schema, right_schema, schema) = join_residual_schemas(); + let column = |name: &str| col(name, &schema).unwrap(); + let validate = |expr: Arc| { + validate_join_residual(&join_residual_filter(expr, &schema), &left_schema, &right_schema).map(|_| ()) + }; + + let concatenation = binary( + binary(column("name"), Operator::StringConcat, column("label"), &schema).unwrap(), + Operator::Eq, + lit(ScalarValue::Utf8(Some("x".to_string()))), + &schema, + ) + .unwrap(); + let rejection = validate(concatenation).unwrap_err(); + assert_eq!(rejection.capability(), ExpressionCapability::JoinResidualShape); + assert!(matches!( + rejection.subject(), + ExpressionRejectionSubject::BinaryOperator { operator: Operator::StringConcat, .. } + )); + + let sum = binary(column("n"), Operator::Plus, column("r"), &schema).unwrap(); + let widened = binary( + Arc::new(CastExpr::new(sum, DataType::Decimal128(12, 0), None)), + Operator::Gt, + lit(ScalarValue::Decimal128(Some(5), 12, 0)), + &schema, + ) + .unwrap(); + let rejection = validate(widened).unwrap_err(); + assert_eq!(rejection.capability(), ExpressionCapability::JoinResidualCastTarget); + assert_eq!( + rejection.subject(), + &ExpressionRejectionSubject::Cast { source: DataType::Int32, target: DataType::Decimal128(12, 0) } + ); + + let difference = binary(column("n"), Operator::Minus, column("r"), &schema).unwrap(); + let distance = binary( + scalar_function("abs", math::abs(), vec![difference], DataType::Int32), + Operator::Lt, + lit(ScalarValue::Int32(Some(5))), + &schema, + ) + .unwrap(); + let rejection = validate(distance).unwrap_err(); assert_eq!(rejection.capability(), ExpressionCapability::JoinResidualScalarFunction); + assert_eq!(rejection.subject(), &ExpressionRejectionSubject::ScalarFunction("abs".to_string())); + + // An inner join runs a null-aware comparison as a post-join filter. + validate(binary(column("n"), Operator::IsDistinctFrom, column("r"), &schema).unwrap()) + .expect("a cross-input IS DISTINCT FROM stays in the residual"); } #[test] @@ -665,11 +869,6 @@ fn join_residual_ast_rejects_integer_division_and_casts_it_cannot_evaluate_or_ch binary(elapsed, Operator::Gt, lit(ScalarValue::DurationMillisecond(Some(5))), &filter_schema).unwrap(); assert_eq!(validate(elapsed_bound).unwrap_err().capability(), ExpressionCapability::JoinResidualShape); - // The cast cannot fail, but the AST has no cast over temporal values. - let temporal_cast = - binary(cast("t", DataType::Int64), Operator::Eq, col("r", &filter_schema).unwrap(), &filter_schema).unwrap(); - assert_eq!(validate(temporal_cast).unwrap_err().capability(), ExpressionCapability::CastPair); - let failing_cast = binary(cast("l", DataType::UInt64), Operator::Eq, cast("r", DataType::UInt64), &filter_schema).unwrap(); assert_eq!(validate(failing_cast).unwrap_err().capability(), ExpressionCapability::CastPair); diff --git a/crates/algeon-datafusion/src/native/lowering/expr.rs b/crates/algeon-datafusion/src/native/lowering/expr.rs index 82ff3c646..1e622d816 100644 --- a/crates/algeon-datafusion/src/native/lowering/expr.rs +++ b/crates/algeon-datafusion/src/native/lowering/expr.rs @@ -81,13 +81,23 @@ pub(super) fn lower_validated_physical_expr_with_bindings( let activation_nodes = Cell::new(0); lower_validated_physical_expr( validated, - LoweringScope { bindings, join_columns: None, activation: None, activation_nodes: &activation_nodes }, + LoweringScope { + bindings, + join_columns: None, + join_substitutions: &[], + activation: None, + activation_nodes: &activation_nodes, + }, ) } +/// Lowers a join filter expression whose column `i` reads `columns[&i]`. +/// Each subexpression in `substitutions`, matched by pointer identity, lowers +/// to its paired expression instead. pub(super) fn lower_validated_join_residual_expr( validated: &ValidatedPhysicalExpr, columns: &BTreeMap, + substitutions: &[(Arc, NativeExpr)], ) -> Result { let activation_nodes = Cell::new(0); lower_validated_physical_expr( @@ -95,6 +105,7 @@ pub(super) fn lower_validated_join_residual_expr( LoweringScope { bindings: None, join_columns: Some(columns), + join_substitutions: substitutions, activation: None, activation_nodes: &activation_nodes, }, @@ -106,6 +117,7 @@ pub(super) fn lower_validated_join_residual_expr( struct LoweringScope<'a> { bindings: Option<&'a [FieldBinding]>, join_columns: Option<&'a BTreeMap>, + join_substitutions: &'a [(Arc, NativeExpr)], /// Set inside a `CASE` branch or the right operand of an `AND` or `OR`, /// which DataFusion may evaluate on a subset of rows: the rows where this /// boolean is TRUE. Native evaluation visits every row, so an operand that @@ -150,6 +162,9 @@ fn lower_validated_physical_expr( ) -> Result { let expr = validated.expr(); let input_schema = validated.input_schema(); + if let Some((_, substitute)) = scope.join_substitutions.iter().find(|(site, _)| Arc::ptr_eq(site, expr)) { + return Ok(substitute.clone()); + } if let Some(column) = expr.downcast_ref::() { if let Some(columns) = scope.join_columns { let resolved = resolve_column(column, input_schema)?; diff --git a/crates/algeon-datafusion/src/native/lowering/join.rs b/crates/algeon-datafusion/src/native/lowering/join.rs index 81e4e1b21..0b37eda98 100644 --- a/crates/algeon-datafusion/src/native/lowering/join.rs +++ b/crates/algeon-datafusion/src/native/lowering/join.rs @@ -157,9 +157,9 @@ fn add_join_residual_materializations( input: LoweredNativeInput, materializations: &[JoinResidualMaterialization], ) -> Result { - // Residual casts that cannot live in the predicate AST are appended as - // synthetic pre-join projection columns on the original DataFusion side. - // Canonical input ordering happens after this step. + // Residual subexpressions that the contract moves before the join are + // appended as synthetic projection columns on the original DataFusion + // side. Canonical input ordering happens after this step. if materializations.is_empty() { return Ok(input); } diff --git a/crates/algeon-datafusion/src/native/lowering/join_residual.rs b/crates/algeon-datafusion/src/native/lowering/join_residual.rs index b6d4843aa..60b2b6f1c 100644 --- a/crates/algeon-datafusion/src/native/lowering/join_residual.rs +++ b/crates/algeon-datafusion/src/native/lowering/join_residual.rs @@ -1,22 +1,23 @@ -use std::{collections::BTreeSet, sync::Arc}; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::Arc, +}; use arrow::datatypes::{Field, SchemaRef}; use datafusion::{ common::JoinSide as DataFusionJoinSide, logical_expr::Operator, physical_expr::{ - PhysicalExpr, ScalarFunctionExpr, + PhysicalExpr, expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, Literal, NotExpr, }, }, physical_plan::joins::utils::JoinFilter, }; -use query_planning::{ - expr::{ - BinaryOp as NativeBinaryOp, CaseWhen as NativeCaseWhen, Expr as NativeExpr, LiteralValue as NativeLiteralValue, - }, - validation::{is_string_type, substring_stop_position}, +use query_planning::expr::{ + BinaryOp as NativeBinaryOp, CaseWhen as NativeCaseWhen, ColumnRef as NativeColumnRef, Expr as NativeExpr, + LiteralValue as NativeLiteralValue, }; use super::{ @@ -32,11 +33,7 @@ use crate::{ capability::CapabilityReason, native::{ LoweringFailure, - expression_contract::{ - JoinResidualCastMode, ScalarFunctionKind, ValidatedJoinResidual, case_expr_data_type, - declared_decimal_divide_output_type, join_residual_cast_mode, physical_expr_data_type, - }, - literal::literal_i32, + expression_contract::{ValidatedJoinResidual, case_expr_data_type}, }, }; @@ -106,8 +103,6 @@ pub(super) fn lower_join_filter( canonical: CanonicalJoin, bindings: Option>, ) -> Result { - // Expressions the cuDF join AST cannot represent directly may be materialized - // on one input side. The residual predicate must still reference both sides. let filter = contract.filter(); debug_assert!(Arc::ptr_eq(contract.expression().expr(), filter.expression())); let mut left_required = BTreeSet::new(); @@ -121,6 +116,23 @@ pub(super) fn lower_join_filter( canonical, bindings, }; + // Each materialization site becomes a synthetic column of its input, and + // the residual reads that column in its place. + let substitutions = contract + .materializations() + .iter() + .map(|site| { + let column = materialize_join_residual_expr( + &site.expr, + site.side, + context, + &mut left_required, + &mut right_required, + &mut materializations, + )?; + Ok((Arc::clone(&site.expr), column)) + }) + .collect::, LoweringFailure>>()?; let predicate = if contract.contains_decimal_division() { // Evaluate checked arithmetic after matching keys. Pre-join // materialization would also evaluate rows with no matching key. @@ -140,15 +152,9 @@ pub(super) fn lower_join_filter( Ok((index, column)) }) .collect::>()?; - lower_validated_join_residual_expr(contract.expression(), &columns)? + lower_validated_join_residual_expr(contract.expression(), &columns, &substitutions)? } else { - lower_join_filter_expr( - filter.expression(), - context, - &mut left_required, - &mut right_required, - &mut materializations, - )? + lower_join_filter_expr(filter.expression(), context, &substitutions, &mut left_required, &mut right_required)? }; debug_assert!( !left_required.is_empty() && !right_required.is_empty(), @@ -157,66 +163,53 @@ pub(super) fn lower_join_filter( Ok(LoweredJoinResidual { predicate, left_required, right_required, materializations }) } +/// Lowers the residual operators the cuDF join AST evaluates, or that an +/// inner join evaluates as a post-join filter. Every other single-input +/// subexpression is a contract materialization site and arrives here as a +/// substitution. fn lower_join_filter_expr( expr: &Arc, context: JoinFilterLoweringContext<'_>, + substitutions: &[(Arc, NativeExpr)], left_required: &mut BTreeSet, right_required: &mut BTreeSet, - materializations: &mut JoinResidualMaterializations, ) -> Result { + if let Some((_, column)) = substitutions.iter().find(|(site, _)| Arc::ptr_eq(site, expr)) { + return Ok(column.clone()); + } if let Some(column) = expr.downcast_ref::() { return lower_join_filter_column(column, context, left_required, right_required); } if let Some(literal) = expr.downcast_ref::() { return Ok(lower_literal(literal.value()).expect("validated join-residual literal must lower")); } + let mut lower = |child: &Arc| { + lower_join_filter_expr(child, context, substitutions, left_required, right_required) + }; if let Some(binary) = expr.downcast_ref::() { - return lower_join_filter_binary_expr( - binary, - context, - |expr| lower_join_filter_expr(expr, context, left_required, right_required, materializations), - "validated join-residual decimal division must have a supported output type", - ); + let left = lower(binary.left())?; + let right = lower(binary.right())?; + return Ok(match binary.op() { + Operator::IsDistinctFrom => NativeExpr::is_distinct_from(left, right), + Operator::IsNotDistinctFrom => NativeExpr::is_not_distinct_from(left, right), + op => NativeExpr::binary(left, lower_binary_operator(op), right), + }); } if let Some(not) = expr.downcast_ref::() { - return Ok(!lower_join_filter_expr(not.arg(), context, left_required, right_required, materializations)?); + return Ok(!lower(not.arg())?); } if let Some(is_null) = expr.downcast_ref::() { - return Ok(NativeExpr::is_null(lower_join_filter_expr( - is_null.arg(), - context, - left_required, - right_required, - materializations, - )?)); + return Ok(NativeExpr::is_null(lower(is_null.arg())?)); } if let Some(is_not_null) = expr.downcast_ref::() { - return Ok(NativeExpr::is_not_null(lower_join_filter_expr( - is_not_null.arg(), - context, - left_required, - right_required, - materializations, - )?)); + return Ok(NativeExpr::is_not_null(lower(is_not_null.arg())?)); } if let Some(cast) = expr.downcast_ref::() { - if matches!( - join_residual_cast_mode(cast.cast_type()).expect("validated join-residual cast target"), - JoinResidualCastMode::Materialize - ) { - // Some casts are valid native expressions but not accepted in the - // join residual AST. Materialize those on one input side and let - // the residual predicate reference the synthetic column. - return materialize_join_residual_expr(expr, context, left_required, right_required, materializations); - } - return Ok(NativeExpr::cast( - lower_join_filter_expr(cast.expr(), context, left_required, right_required, materializations)?, - cast.cast_type().clone(), - )); + return Ok(NativeExpr::cast(lower(cast.expr())?, cast.cast_type().clone())); } if let Some(in_list) = expr.downcast_ref::() { debug_assert!(!in_list.is_empty(), "empty IN list cannot carry a join-residual validation proof"); - let expr = lower_join_filter_expr(in_list.expr(), context, left_required, right_required, materializations)?; + let value = lower(in_list.expr())?; let list = in_list .list() .iter() @@ -226,76 +219,40 @@ fn lower_join_filter_expr( !literal.value().is_null(), "null IN entry cannot carry a join-residual validation proof" ); - lower_join_filter_expr(item, context, left_required, right_required, materializations) + lower(item) }) .collect::, LoweringFailure>>()?; return Ok(if in_list.negated() { - NativeExpr::not_in_list(expr, list) + NativeExpr::not_in_list(value, list) } else { - NativeExpr::in_list(expr, list) + NativeExpr::in_list(value, list) }); } if let Some(case) = expr.downcast_ref::() { - if context.contract.single_side(expr).is_ok() { - return materialize_join_residual_expr(expr, context, left_required, right_required, materializations); - } - return lower_join_filter_case_expr(case, context, left_required, right_required, materializations); - } - if expr.downcast_ref::().is_some() { - return materialize_join_residual_expr(expr, context, left_required, right_required, materializations); + return lower_join_filter_case_expr(case, context, &mut lower); } unreachable!("validated join residual must have a lowering branch") } -fn lower_join_filter_binary_expr( - binary: &BinaryExpr, - context: JoinFilterLoweringContext<'_>, - mut lower_child: impl FnMut(&Arc) -> Result, - decimal_division_expectation: &'static str, -) -> Result { - let left = lower_child(binary.left())?; - let right = lower_child(binary.right())?; - let op = lower_binary_operator(binary.op()); - if matches!(binary.op(), Operator::Divide) - && let Some(output_type) = - declared_decimal_divide_output_type(binary, context.filter.schema()).expect(decimal_division_expectation) - { - return Ok(NativeExpr::binary_with_output_type(left, op, right, output_type)); - } - Ok(NativeExpr::binary(left, op, right)) -} - fn lower_join_filter_case_expr( case: &CaseExpr, context: JoinFilterLoweringContext<'_>, - left_required: &mut BTreeSet, - right_required: &mut BTreeSet, - materializations: &mut JoinResidualMaterializations, + lower: &mut impl FnMut(&Arc) -> Result, ) -> Result { - let base_expr = case - .expr() - .map(|expr| lower_join_filter_expr(expr, context, left_required, right_required, materializations)) - .transpose()?; + let base_expr = case.expr().map(&mut *lower).transpose()?; let branches = case .when_then_expr() .iter() .map(|(when, then)| { let when = match &base_expr { - Some(base_expr) => NativeExpr::binary( - base_expr.clone(), - NativeBinaryOp::Eq, - lower_join_filter_expr(when, context, left_required, right_required, materializations)?, - ), - None => lower_join_filter_expr(when, context, left_required, right_required, materializations)?, + Some(base_expr) => NativeExpr::binary(base_expr.clone(), NativeBinaryOp::Eq, lower(when)?), + None => lower(when)?, }; - Ok(NativeCaseWhen::new( - when, - lower_join_filter_expr(then, context, left_required, right_required, materializations)?, - )) + Ok(NativeCaseWhen::new(when, lower(then)?)) }) .collect::, LoweringFailure>>()?; let else_expr = match case.else_expr() { - Some(else_expr) => lower_join_filter_expr(else_expr, context, left_required, right_required, materializations)?, + Some(else_expr) => lower(else_expr)?, None => { let data_type = case_expr_data_type(case, context.filter.schema()) .expect("validated join-residual CASE must have an output type"); @@ -305,28 +262,45 @@ fn lower_join_filter_case_expr( Ok(NativeExpr::case(branches, else_expr)) } +/// Appends `expr`, which reads only `side`, to that input as a synthetic +/// column and returns the residual's reference to it. The general expression +/// lowering produces the column, as it does for a join key expression. fn materialize_join_residual_expr( expr: &Arc, + side: DataFusionJoinSide, context: JoinFilterLoweringContext<'_>, left_required: &mut BTreeSet, right_required: &mut BTreeSet, materializations: &mut JoinResidualMaterializations, ) -> Result { - // Materializations are valid only when the whole expression belongs to one - // DataFusion side. Cross-side expressions must remain in the residual AST. - let side = context.contract.single_side(expr).expect("validated materialized residual must belong to one side"); let (schema, required) = match side { DataFusionJoinSide::Left => (context.left_schema, left_required), DataFusionJoinSide::Right => (context.right_schema, right_required), DataFusionJoinSide::None => unreachable!("validated materialized residual must belong to left or right"), }; - let native_expr = lower_join_filter_side_expr(expr, context, schema, side, required)?; + let mut filter_indices = BTreeSet::new(); + collect_required_field_indices(expr, context.filter.schema(), &mut filter_indices)?; + let mut columns = BTreeMap::new(); + for filter_index in filter_indices { + let column_index = context + .filter + .column_indices() + .get(filter_index) + .expect("validated residual column must have a side mapping"); + debug_assert_eq!(column_index.side, side, "a materialization site reads one input"); + let field = + schema.fields().get(column_index.index).expect("validated residual column must exist in its side schema"); + required.insert(column_index.index); + let native_name = side_native_name(context, side, column_index.index, field.name())?; + columns.insert(filter_index, NativeColumnRef::unqualified(native_name)); + } + let validated = context.contract.expression().trusted_descendant(expr); + let native_expr = lower_validated_join_residual_expr(&validated, &columns, &[])?; let side_materializations = materializations.side_mut(side).expect("validated materialization side must be left or right"); let alias = unique_synthetic_join_residual_alias(side, side_materializations.len(), schema, side_materializations); - let field = - physical_expr_field_from_validated(&context.contract.expression().trusted_descendant(expr), alias.as_str()) - .expect("validated materialized residual must have a native field"); + let field = physical_expr_field_from_validated(&validated, alias.as_str()) + .expect("validated materialized residual must have a native field"); side_materializations.push(JoinResidualMaterialization { alias: alias.clone(), expr: alias_if_needed(native_expr, alias.as_str()), @@ -339,204 +313,6 @@ fn materialize_join_residual_expr( Ok(NativeExpr::input_column(native_input_index, alias.as_str(), alias.as_str())) } -fn lower_join_filter_side_expr( - expr: &Arc, - context: JoinFilterLoweringContext<'_>, - side_schema: &SchemaRef, - side: DataFusionJoinSide, - required: &mut BTreeSet, -) -> Result { - if let Some(column) = expr.downcast_ref::() { - let column_index = context - .filter - .column_indices() - .get(column.index()) - .expect("validated residual column must have a side mapping"); - if column_index.side != side { - unreachable!("validated single-side residual cannot cross join inputs"); - } - let field = side_schema - .fields() - .get(column_index.index) - .expect("validated residual column must exist in its side schema"); - required.insert(column_index.index); - return Ok(NativeExpr::column(side_native_name(context, side, column_index.index, field.name())?)); - } - if let Some(literal) = expr.downcast_ref::() { - return Ok(lower_literal(literal.value()).expect("validated materialized literal must lower")); - } - if let Some(binary) = expr.downcast_ref::() { - return lower_join_filter_binary_expr( - binary, - context, - |expr| lower_join_filter_side_expr(expr, context, side_schema, side, required), - "validated materialized decimal division must have a supported output type", - ); - } - if let Some(not) = expr.downcast_ref::() { - return Ok(!lower_join_filter_side_expr(not.arg(), context, side_schema, side, required)?); - } - if let Some(is_null) = expr.downcast_ref::() { - return Ok(NativeExpr::is_null(lower_join_filter_side_expr( - is_null.arg(), - context, - side_schema, - side, - required, - )?)); - } - if let Some(is_not_null) = expr.downcast_ref::() { - return Ok(NativeExpr::is_not_null(lower_join_filter_side_expr( - is_not_null.arg(), - context, - side_schema, - side, - required, - )?)); - } - if let Some(cast) = expr.downcast_ref::() { - let source_type = physical_expr_data_type(cast.expr(), context.filter.schema()) - .expect("validated materialized cast source must have a type"); - if is_string_type(&source_type) && is_string_type(cast.cast_type()) { - return lower_join_filter_side_expr(cast.expr(), context, side_schema, side, required); - } - return Ok(NativeExpr::cast( - lower_join_filter_side_expr(cast.expr(), context, side_schema, side, required)?, - cast.cast_type().clone(), - )); - } - if let Some(in_list) = expr.downcast_ref::() { - debug_assert!(!in_list.is_empty(), "empty materialized IN list cannot carry a validation proof"); - let expr = lower_join_filter_side_expr(in_list.expr(), context, side_schema, side, required)?; - let list = in_list - .list() - .iter() - .map(|item| { - let literal = item.downcast_ref::().expect("validated materialized IN entries are literals"); - debug_assert!(!literal.value().is_null(), "null materialized IN entry cannot carry a validation proof"); - lower_join_filter_side_expr(item, context, side_schema, side, required) - }) - .collect::, LoweringFailure>>()?; - return Ok(if in_list.negated() { - NativeExpr::not_in_list(expr, list) - } else { - NativeExpr::in_list(expr, list) - }); - } - if let Some(case) = expr.downcast_ref::() { - return lower_join_filter_side_case_expr(case, context, side_schema, side, required); - } - if let Some(function) = expr.downcast_ref::() { - return lower_join_filter_side_scalar_function_expr(function, context, side_schema, side, required); - } - unreachable!("validated materialized residual must have a lowering branch") -} - -fn lower_join_filter_side_case_expr( - case: &CaseExpr, - context: JoinFilterLoweringContext<'_>, - side_schema: &SchemaRef, - side: DataFusionJoinSide, - required: &mut BTreeSet, -) -> Result { - let base_expr = - case.expr().map(|expr| lower_join_filter_side_expr(expr, context, side_schema, side, required)).transpose()?; - let branches = case - .when_then_expr() - .iter() - .map(|(when, then)| { - let when = match &base_expr { - Some(base_expr) => NativeExpr::binary( - base_expr.clone(), - NativeBinaryOp::Eq, - lower_join_filter_side_expr(when, context, side_schema, side, required)?, - ), - None => lower_join_filter_side_expr(when, context, side_schema, side, required)?, - }; - Ok(NativeCaseWhen::new(when, lower_join_filter_side_expr(then, context, side_schema, side, required)?)) - }) - .collect::, LoweringFailure>>()?; - let else_expr = match case.else_expr() { - Some(else_expr) => lower_join_filter_side_expr(else_expr, context, side_schema, side, required)?, - None => { - let data_type = case_expr_data_type(case, context.filter.schema()) - .expect("validated materialized CASE must have an output type"); - NativeExpr::literal(NativeLiteralValue::Null(data_type)) - }, - }; - Ok(NativeExpr::case(branches, else_expr)) -} - -fn lower_join_filter_side_scalar_function_expr( - function: &ScalarFunctionExpr, - context: JoinFilterLoweringContext<'_>, - side_schema: &SchemaRef, - side: DataFusionJoinSide, - required: &mut BTreeSet, -) -> Result { - match ScalarFunctionKind::from_scalar_function(function).expect("validated join residual scalar function") { - ScalarFunctionKind::Substring => { - lower_join_filter_side_substring_function(function, context, side_schema, side, required) - }, - ScalarFunctionKind::Lower => lower_join_filter_side_string_case_function( - function, - context, - side_schema, - side, - required, - NativeExpr::lower, - ), - ScalarFunctionKind::Upper => lower_join_filter_side_string_case_function( - function, - context, - side_schema, - side, - required, - NativeExpr::upper, - ), - _ => unreachable!("join residual proof excludes other scalar functions"), - } -} - -fn lower_join_filter_side_substring_function( - function: &ScalarFunctionExpr, - context: JoinFilterLoweringContext<'_>, - side_schema: &SchemaRef, - side: DataFusionJoinSide, - required: &mut BTreeSet, -) -> Result { - let args = function.args(); - debug_assert!((2..=3).contains(&args.len()), "invalid SUBSTRING arity cannot carry a residual proof"); - let start = literal_i32(&args[1]).expect("validated residual SUBSTRING start must be an i32 literal"); - debug_assert!(start >= 1, "invalid residual SUBSTRING start cannot carry a validation proof"); - let length = if args.len() == 3 { - Some(literal_i32(&args[2]).expect("validated residual SUBSTRING length must be an i32 literal")) - } else { - None - }; - debug_assert!( - length.is_none_or(|length| length >= 0 && substring_stop_position(start, length).is_some()), - "invalid residual SUBSTRING length cannot carry a validation proof" - ); - Ok(NativeExpr::substring( - lower_join_filter_side_expr(&args[0], context, side_schema, side, required)?, - start, - length, - )) -} - -fn lower_join_filter_side_string_case_function( - function: &ScalarFunctionExpr, - context: JoinFilterLoweringContext<'_>, - side_schema: &SchemaRef, - side: DataFusionJoinSide, - required: &mut BTreeSet, - constructor: fn(NativeExpr) -> NativeExpr, -) -> Result { - let [expr] = function.args() else { unreachable!("validated residual string case must have one argument") }; - Ok(constructor(lower_join_filter_side_expr(expr, context, side_schema, side, required)?)) -} - fn unique_synthetic_join_residual_alias( side: DataFusionJoinSide, index: usize, @@ -619,7 +395,9 @@ mod tests { logical_expr::{Operator, ScalarUDF}, physical_expr::{ PhysicalExpr, ScalarFunctionExpr, - expressions::{CaseExpr, CastExpr, IsNotNullExpr, IsNullExpr, NotExpr, binary, col, in_list, lit}, + expressions::{ + CaseExpr, CastExpr, IsNotNullExpr, IsNullExpr, LikeExpr, NotExpr, binary, col, in_list, lit, + }, }, physical_plan::joins::utils::{ColumnIndex, JoinFilter}, }; @@ -659,7 +437,7 @@ mod tests { left_schema: &SchemaRef, right_schema: &SchemaRef, filter_schema: &SchemaRef, - ) { + ) -> LoweredJoinResidual { let filter = JoinFilter::new( expr, vec![ColumnIndex { index: 0, side: JoinSide::Left }, ColumnIndex { index: 0, side: JoinSide::Right }], @@ -671,6 +449,24 @@ mod tests { .unwrap_or_else(|error| panic!("{label} must lower: {error:?}")); assert!(!lowered.left_required.is_empty(), "{label} must retain a left input dependency"); assert!(!lowered.right_required.is_empty(), "{label} must retain a right input dependency"); + lowered + } + + /// The single left-input materialization of `lowered`, which the residual reads by its alias. + fn left_materialization<'l>(label: &str, lowered: &'l LoweredJoinResidual) -> &'l NativeExpr { + let [materialization] = lowered.materializations.side(DataFusionJoinSide::Left).unwrap() else { + panic!("{label} must materialize one left-input expression"); + }; + let alias = materialization.alias.as_str(); + assert!( + lowered.predicate.any(|expr| matches!(expr, NativeExpr::Column(column) if column.name() == alias)), + "{label} residual must read {alias}: {}", + lowered.predicate + ); + let NativeExpr::Alias { expr, .. } = &materialization.expr else { + panic!("{label} materialization must carry its alias: {:?}", materialization.expr); + }; + expr } #[test] @@ -770,5 +566,59 @@ mod tests { let predicate = binary(function, Operator::Eq, right(), &schema).unwrap(); assert_join_residual_parity(&name, predicate, &left_schema, &right_schema, &schema); } + + let text = |value: &str| lit(ScalarValue::Utf8(Some(value.to_string()))); + let or_right_is_x = |expr| { + binary(expr, Operator::Or, binary(right(), Operator::Eq, text("x"), &schema).unwrap(), &schema).unwrap() + }; + let like = Arc::new(LikeExpr::new(false, false, left(), text("a%"))); + let lowered = assert_join_residual_parity("LIKE", or_right_is_x(like), &left_schema, &right_schema, &schema); + assert!(matches!(left_materialization("LIKE", &lowered), NativeExpr::Like { .. })); + + let concatenation = binary(left(), Operator::StringConcat, text("!"), &schema).unwrap(); + let predicate = binary(concatenation, Operator::Eq, right(), &schema).unwrap(); + let lowered = assert_join_residual_parity("||", predicate, &left_schema, &right_schema, &schema); + assert!(left_materialization("||", &lowered).any(|expr| matches!(expr, NativeExpr::Concat(_)))); + + for operator in [Operator::IsDistinctFrom, Operator::IsNotDistinctFrom] { + let single_input = binary(left(), operator, text("x"), &schema).unwrap(); + let label = format!("single-input {operator}"); + let lowered = + assert_join_residual_parity(&label, or_right_is_x(single_input), &left_schema, &right_schema, &schema); + assert!(matches!( + left_materialization(&label, &lowered), + NativeExpr::IsDistinctFrom { .. } | NativeExpr::IsNotDistinctFrom { .. } + )); + + let cross_input = binary(left(), operator, right(), &schema).unwrap(); + let lowered = assert_join_residual_parity( + &format!("cross-input {operator}"), + cross_input, + &left_schema, + &right_schema, + &schema, + ); + assert!(matches!( + lowered.predicate, + NativeExpr::IsDistinctFrom { .. } | NativeExpr::IsNotDistinctFrom { .. } + )); + } + + let (left_schema, right_schema, schema) = residual_schemas(DataType::Int32); + let left = || col("left", &schema).unwrap(); + let right = || col("right", &schema).unwrap(); + let halved = binary(left(), Operator::Divide, lit(ScalarValue::Int32(Some(2))), &schema).unwrap(); + let predicate = binary(halved, Operator::Eq, right(), &schema).unwrap(); + let lowered = assert_join_residual_parity( + "integer division by a literal", + predicate, + &left_schema, + &right_schema, + &schema, + ); + assert!(matches!( + left_materialization("integer division by a literal", &lowered), + NativeExpr::Binary { op: NativeBinaryOp::Divide, .. } + )); } } diff --git a/crates/algeon-datafusion/src/native/lowering_tests.rs b/crates/algeon-datafusion/src/native/lowering_tests.rs index dbb892247..d4b5684dc 100644 --- a/crates/algeon-datafusion/src/native/lowering_tests.rs +++ b/crates/algeon-datafusion/src/native/lowering_tests.rs @@ -453,15 +453,19 @@ fn assert_join_input_materializes_upper_utf8view_cast( if actual_alias == alias && matches!( expr.as_ref(), - Expr::Upper(expr) + Expr::Cast { expr, target_type: DataType::Utf8View } if matches!( expr.as_ref(), - Expr::Column(column) if column.name() == source_column + Expr::Upper(expr) + if matches!( + expr.as_ref(), + Expr::Column(column) if column.name() == source_column + ) ) ) ) }), - "join input projection should materialize upper({source_column}) as {alias}: {:?}", + "join input projection should materialize CAST(upper({source_column}) AS Utf8View) as {alias}: {:?}", projection.expressions().collect::>() ); } diff --git a/crates/algeon-datafusion/src/native/lowering_tests/full_gpu_gap_analysis.rs b/crates/algeon-datafusion/src/native/lowering_tests/full_gpu_gap_analysis.rs index 3bbc872b2..37f0a4b74 100644 --- a/crates/algeon-datafusion/src/native/lowering_tests/full_gpu_gap_analysis.rs +++ b/crates/algeon-datafusion/src/native/lowering_tests/full_gpu_gap_analysis.rs @@ -426,6 +426,74 @@ const EXISTING_COVERAGE_CASES: &[GpuGapCase] = &[ }, ]; +/// Join filters DataFusion keeps on the join because they read both inputs. +/// Their single-input parts outside the cuDF join AST are computed on that +/// input before the join, unless they can fail. +const JOIN_RESIDUAL_CASES: &[GpuGapCase] = &[ + GpuGapCase { + id: "join_residual_like_inner", + + sql: "SELECT b.id, s.id AS sid FROM bench b JOIN bench s ON b.key = s.key \ + AND (s.name LIKE 'a%' OR b.payload > 10)", + expected: ExpectedGpuCoverageOutcome::Selected, + }, + GpuGapCase { + id: "join_residual_lower_like_left", + + sql: "SELECT b.id, s.id AS sid FROM bench b LEFT JOIN bench s ON b.key = s.key \ + AND (LOWER(s.name) LIKE '%a%' OR b.payload > 10)", + expected: ExpectedGpuCoverageOutcome::Selected, + }, + GpuGapCase { + id: "join_residual_not_like_full", + + sql: "SELECT b.id, s.id AS sid FROM bench b FULL JOIN bench s ON b.key = s.key \ + AND (s.name NOT LIKE '%a%' OR b.payload > 10)", + expected: ExpectedGpuCoverageOutcome::Selected, + }, + GpuGapCase { + id: "join_residual_like_semi", + + sql: "SELECT b.id FROM bench b WHERE EXISTS \ + (SELECT 1 FROM bench s WHERE s.key = b.key AND (s.name LIKE 'a%' OR b.payload > 10))", + expected: ExpectedGpuCoverageOutcome::Selected, + }, + GpuGapCase { + id: "join_residual_like_anti", + + sql: "SELECT b.id FROM bench b WHERE NOT EXISTS \ + (SELECT 1 FROM bench s WHERE s.key = b.key AND (s.name LIKE 'a%' OR b.payload > 10))", + expected: ExpectedGpuCoverageOutcome::Selected, + }, + GpuGapCase { + id: "join_residual_null_aware_comparisons", + + sql: "SELECT b.id, s.id AS sid FROM bench b JOIN bench s ON b.key = s.key \ + AND ((s.payload IS DISTINCT FROM 5) OR b.name IS NOT DISTINCT FROM s.name)", + expected: ExpectedGpuCoverageOutcome::Selected, + }, + GpuGapCase { + id: "join_residual_concatenation", + + sql: "SELECT b.id, s.id AS sid FROM bench b JOIN bench s ON b.key = s.key \ + AND ((s.name || '!') = b.name OR b.payload > 10)", + expected: ExpectedGpuCoverageOutcome::Selected, + }, + // The cast fails on a payload above 999. DataFusion casts only key-matched + // pairs, and a pre-join column would cast every row. + GpuGapCase { + id: "join_residual_cast_that_can_fail", + + sql: "SELECT b.id FROM bench b JOIN bench s ON b.key = s.key \ + AND CAST(s.payload AS DECIMAL(3, 0)) > CAST(b.id AS DECIMAL(10, 0))", + expected: ExpectedGpuCoverageOutcome::NotSupported { + reason: CapabilityReason::UnsupportedExpression, + unsupported_category: "expression", + blocking: ("cast", "cast from Int64 to Decimal128(3, 0) in `CAST(payload AS Decimal128(3, 0))`"), + }, + }, +]; + #[tokio::test] async fn full_gpu_gap_analysis_sql_planning_coverage_is_stable() { let dir = unique_temp_dir("algeon-datafusion-full-gpu-gap-analysis"); @@ -468,6 +536,7 @@ fn gpu_gap_cases() -> Vec { SCALAR_EXPRESSION_CASES, ROOT_SOURCE_CASES, EXISTING_COVERAGE_CASES, + JOIN_RESIDUAL_CASES, ] .into_iter() .flatten() diff --git a/crates/algeon-datafusion/src/native/lowering_tests/join_shapes/residual_variants.rs b/crates/algeon-datafusion/src/native/lowering_tests/join_shapes/residual_variants.rs index dd98df35a..b106183f0 100644 --- a/crates/algeon-datafusion/src/native/lowering_tests/join_shapes/residual_variants.rs +++ b/crates/algeon-datafusion/src/native/lowering_tests/join_shapes/residual_variants.rs @@ -490,6 +490,8 @@ fn hash_join_filter_substring_materializes_inputs_before_residual_predicate() { let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap(); + // The plan reads one left SUBSTRING expression in both disjuncts, which + // lowers to one pre-join column. let join = rewritten_native_join(&result); let residual = join.spec.residual_predicate.as_ref().expect("join filter should lower as residual predicate"); assert!(matches!( @@ -509,7 +511,7 @@ fn hash_join_filter_substring_materializes_inputs_before_residual_predicate() { Expr::Binary { left, right, op: BinaryOp::NotEq, .. } if matches!( left.as_ref(), - Expr::Column(column) if column.to_string() == "input0.__algeon_native_join_residual_left_1" + Expr::Column(column) if column.to_string() == "input0.__algeon_native_join_residual_left_0" ) && matches!( right.as_ref(), @@ -523,12 +525,6 @@ fn hash_join_filter_substring_materializes_inputs_before_residual_predicate() { "__algeon_native_join_residual_left_0", "l_zip", ); - assert_join_input_materializes_substring( - result.selected_plan(), - join.inputs[0], - "__algeon_native_join_residual_left_1", - "l_zip", - ); assert_join_input_materializes_substring( result.selected_plan(), join.inputs[1], @@ -539,6 +535,64 @@ fn hash_join_filter_substring_materializes_inputs_before_residual_predicate() { fs::remove_dir_all(dir).unwrap(); } +#[tokio::test] +async fn sql_join_filter_like_is_computed_on_the_input_it_reads() { + let dir = unique_temp_dir("algeon-datafusion-native-lowering-sql-join-filter-like"); + let path = dir.join("bench.parquet"); + write_full_gpu_gap_analysis_parquet_file(&path); + + for (label, sql, join_type) in [ + ( + "inner", + "SELECT b.id, s.id AS sid FROM bench b JOIN bench s ON b.key = s.key \ + AND (s.name LIKE '%a%' OR b.payload > 10)", + NativeJoinType::Inner, + ), + ( + "left", + "SELECT b.id, s.id AS sid FROM bench b LEFT JOIN bench s ON b.key = s.key \ + AND (s.name LIKE '%a%' OR b.payload > 10)", + NativeJoinType::Left, + ), + ( + "semi", + "SELECT b.id FROM bench b WHERE EXISTS \ + (SELECT 1 FROM bench s WHERE s.key = b.key AND (s.name LIKE '%a%' OR b.payload > 10))", + NativeJoinType::LeftSemi, + ), + ] { + let plan = sql_plan(&path, direct_friendly_sql_session_config(), sql).await; + let result = lower_physical_plan_to_native(plan, &ParquetReadStrategyConfig::default()).unwrap(); + let join = rewritten_native_join(&result); + assert_eq!(join.spec.join_type, join_type, "{label}"); + let residual = join.spec.residual_predicate.as_ref().expect("join filter should lower as residual predicate"); + let mut synthetic = Vec::new(); + residual.any(|expr| { + if let Expr::Column(column) = expr + && column.name().starts_with("__algeon_native_join_residual_") + { + synthetic.push(column.name().to_string()); + } + false + }); + let [alias] = synthetic.as_slice() else { + panic!("{label} residual must read one pre-join column: {residual}"); + }; + let materializes_like = join.inputs.iter().any(|input| { + matches!( + &result.selected_plan().node(*input).expect("join input should exist").kind, + PlanNode::Projection(projection) if projection.expressions().any(|expr| matches!( + expr, + Expr::Alias { expr, alias: actual } if actual == alias && matches!(expr.as_ref(), Expr::Like { .. }) + )) + ) + }); + assert!(materializes_like, "{label}: the input that reads s.name must compute the LIKE as {alias}"); + } + + fs::remove_dir_all(dir).unwrap(); +} + #[test] fn hash_join_filter_upper_string_cast_materializes_input_before_residual_predicate() { let dir = unique_temp_dir("algeon-datafusion-native-lowering-hash-join-filter-upper-cast-materialized"); diff --git a/crates/algeon-datafusion/tests/native_runtime_pruning_tests_gpu/scalar_expressions.rs b/crates/algeon-datafusion/tests/native_runtime_pruning_tests_gpu/scalar_expressions.rs index 935d99ca2..9035cc2bd 100644 --- a/crates/algeon-datafusion/tests/native_runtime_pruning_tests_gpu/scalar_expressions.rs +++ b/crates/algeon-datafusion/tests/native_runtime_pruning_tests_gpu/scalar_expressions.rs @@ -8,7 +8,7 @@ use algeon_datafusion::{ use arrow::{ array::{ Array, AsArray, BooleanArray, Date32Array, Decimal128Array, Float32Array, Float64Array, Int8Array, Int16Array, - Int32Array, Int64Array, + Int32Array, Int64Array, StringArray, }, compute::{cast, concat_batches}, datatypes::{Float32Type, Float64Type, TimeUnit}, @@ -1287,6 +1287,13 @@ fn scalar_expressions_cases_3() { .expect("case runtime") .block_on(native_runtime_checked_failures_fail_the_query_like_cpu()) }), + ("native_runtime_join_filters_with_single_input_expressions_match_cpu", || { + TestRuntimeBuilder::new_current_thread() + .enable_all() + .build() + .expect("case runtime") + .block_on(native_runtime_join_filters_with_single_input_expressions_match_cpu()) + }), ]); } @@ -1563,6 +1570,97 @@ async fn native_runtime_decimal_join_residuals_match_cpu_without_host_boundaries fs::remove_dir_all(dir).unwrap(); } +fn write_join_filter_string_input_parquet(path: &std::path::Path) { + let schema = Arc::new(Schema::new(vec![ + Field::new("row_id", DataType::Int32, false), + Field::new("key", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + Field::new("amount", DataType::Int64, true), + ])); + // Each key holds two rows, so every row has one partner besides itself. Row 4 has no name and no amount. + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8])), + Arc::new(Int32Array::from(vec![1, 1, 2, 2, 3, 3, 4, 4])), + Arc::new(StringArray::from(vec![ + Some("Barbie"), + Some("ken"), + Some("barbie doll"), + None, + Some("Bar"), + Some("alpha"), + Some("Barbie"), + Some("zed"), + ])), + Arc::new(Int64Array::from(vec![Some(5), Some(60), Some(7), None, Some(70), Some(-1), Some(7), Some(100)])), + ], + ) + .unwrap(); + let mut writer = ArrowWriter::try_new(File::create(path).unwrap(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +/// Join filters that read both inputs stay on the join. Their single-input +/// string predicates and null-aware comparisons are computed on that input +/// before the join, and a cross-input null-aware comparison filters the matched +/// pairs of an inner join. +async fn native_runtime_join_filters_with_single_input_expressions_match_cpu() { + let dir = unique_temp_dir("algeon-datafusion-native-runtime-join-filter-strings"); + let path = dir.join("input.parquet"); + write_join_filter_string_input_parquet(&path); + + for (name, sql) in [ + ( + "inner-lower-like", + "SELECT l.row_id, r.row_id AS r_id FROM bench l JOIN bench r ON l.key = r.key \ + AND (LOWER(r.name) LIKE '%bar%' OR l.amount > 50) ORDER BY l.row_id, r_id", + ), + ( + "left-like", + "SELECT l.row_id, r.row_id AS r_id FROM bench l LEFT JOIN bench r ON l.key = r.key \ + AND r.row_id <> l.row_id AND (r.name LIKE 'B%' OR l.amount > 50) ORDER BY l.row_id, r_id", + ), + ( + "full-like", + "SELECT l.row_id, r.row_id AS r_id FROM bench l FULL JOIN bench r ON l.key = r.key \ + AND (r.name LIKE '%ie%' OR l.amount > 60) ORDER BY l.row_id, r_id", + ), + ( + "semi-not-like", + "SELECT row_id FROM bench l WHERE EXISTS (SELECT 1 FROM bench r WHERE r.key = l.key \ + AND r.row_id <> l.row_id AND (r.name NOT LIKE '%a%' OR l.amount < 0)) ORDER BY row_id", + ), + ( + "anti-not-like", + "SELECT row_id FROM bench l WHERE NOT EXISTS (SELECT 1 FROM bench r WHERE r.key = l.key \ + AND r.row_id <> l.row_id AND (r.name NOT LIKE '%a%' OR l.amount < 0)) ORDER BY row_id", + ), + ( + "single-input-concatenation-and-is-distinct-from", + "SELECT l.row_id, r.row_id AS r_id FROM bench l JOIN bench r ON l.key = r.key \ + AND ((r.name || '!') = 'Barbie!' OR (l.amount IS DISTINCT FROM 7)) ORDER BY l.row_id, r_id", + ), + ( + "cross-input-is-not-distinct-from", + "SELECT l.row_id, r.row_id AS r_id FROM bench l JOIN bench r ON l.key = r.key \ + AND l.row_id <> r.row_id AND ((l.name IS NOT DISTINCT FROM r.name) OR l.amount IS NULL) \ + ORDER BY l.row_id, r_id", + ), + ] { + let cpu = cpu_sql_output(&path, sql).await; + assert!(cpu.iter().map(RecordBatch::num_rows).sum::() > 0, "{name} must exercise output rows"); + let plan = sql_plan(&path, SessionConfig::new(), sql).await; + let optimization = optimize_with_pruning_template_requiring_native(plan); + assert_algeon_native_exec(&optimization.plan, name); + let native = collect(optimization.plan, Arc::new(TaskContext::default())).await.unwrap(); + assert_batches_equal(&cpu, &native); + } + + fs::remove_dir_all(dir).unwrap(); +} + async fn native_runtime_decimal_division_failures_fail_the_query_like_cpu() { let dir = unique_temp_dir("algeon-datafusion-native-runtime-decimal-division-failures"); let path = dir.join("input.parquet"); diff --git a/website/docs/design/native-lowering.mdx b/website/docs/design/native-lowering.mdx index f914b8915..3bc3cdda5 100644 --- a/website/docs/design/native-lowering.mdx +++ b/website/docs/design/native-lowering.mdx @@ -241,6 +241,14 @@ may still become GPU fragments. Explicit graph, vector and ML functions also form GPU fragments. Transfers at CPU/GPU boundaries are part of the cost of such a mixed plan. +A join condition that reads both inputs stays on the join. Each part of it +that reads one input and that the cuDF join AST cannot evaluate, such as +`LIKE`, a string function or `IS DISTINCT FROM`, becomes a column of that +input computed before the join, as an expression join key does. Only a part +that cannot fail moves, because DataFusion evaluates the condition on +key-matched pairs alone. [Native Evaluation Semantics](../reference/native-evaluation-semantics.mdx#where-checked-operations-are-not-evaluated) +states where a part with a checked operation runs. + Decimal division in a keyed inner join predicate uses the native GPU post-join filter. Equality keys first select candidate row pairs. The GPU column evaluator then checks decimal division and preserves `CASE` branch diff --git a/website/docs/reference/native-evaluation-semantics.mdx b/website/docs/reference/native-evaluation-semantics.mdx index 1926853ba..acabd175a 100644 --- a/website/docs/reference/native-evaluation-semantics.mdx +++ b/website/docs/reference/native-evaluation-semantics.mdx @@ -65,10 +65,22 @@ WHERE n >= 0 AND total / n > 1 -- fails on a row with n = 0, as on Data ## Where checked operations are not evaluated Integer `/`, any `%`, and timestamp `+` or `-` are not evaluated inside reader -predicates or join residual conditions. A reader predicate with one of them -stays a filter after the scan, and a join residual with one of them keeps the -join on DataFusion. A join residual with a decimal `/` keeps the join on -DataFusion as well. +predicates. A reader predicate with one of them stays a filter after the scan. + +DataFusion evaluates a join condition that reads both inputs only on the row +pairs whose equality keys match. Native lowering moves each part of the +condition that reads one input and contains no checked operation before the +join, where it is computed on every row of that input. Such a part cannot +fail, and the extra rows change no result. Integer `/` and `%` by a literal +other than `0` and `-1`, and floating-point `%`, can move this way. + +A part that reads one input and contains a checked operation stays in the +condition. When the condition belongs to an inner join with equality keys and +contains a decimal `/`, it is evaluated after matching, on the matched pairs +only, with the guards described above. Any other condition with a checked +operation keeps the join on DataFusion. Integer `/` or `%` with a column +divisor, `%` over both inputs, and timestamp `+` or `-` keep it there in every +join.
diff --git a/website/docs/reference/native-sql-support.mdx b/website/docs/reference/native-sql-support.mdx index 89e3b6889..43b6e65bf 100644 --- a/website/docs/reference/native-sql-support.mdx +++ b/website/docs/reference/native-sql-support.mdx @@ -33,7 +33,7 @@ DataFusion, the DataFusion baseline remains in place. | Expression | Native forms | | --- | --- | | Values | Columns and literals. | -| Arithmetic | `+`, `-`, `*`, `/`, `%` and unary minus. Decimal `%` is not native. Integer `/`, any `%` and timestamp `+`/`-` are not evaluated in reader predicates or join residuals. | +| Arithmetic | `+`, `-`, `*`, `/`, `%` and unary minus. Decimal `%` is not native. Integer `/`, any `%` and timestamp `+`/`-` are not evaluated in reader predicates; [join conditions](./native-evaluation-semantics.mdx#where-checked-operations-are-not-evaluated) evaluate them only in parts that cannot fail. | | Timestamp arithmetic | Two timezone-free timestamps of one unit subtract to a `Duration` of that unit. Such a timestamp plus or minus a duration gives a timestamp. Duration plus duration, date plus or minus duration, and date minus date are not native. | | Comparison and logic | Comparisons, `AND`, `OR`, `NOT`, null tests. | | Conditionals | `CASE`. | From b393e49d171f487c471c14a95940157f57c1ce6f Mon Sep 17 00:00:00 2001 From: Morris Tai Date: Mon, 19 Oct 2026 19:11:32 +0000 Subject: [PATCH 2/5] feat: add the decision-model demo over the Flight SQL server fixture/demo/decision_models runs the probe, decide, materialize, recompute loop against an unchanged server: a probe query bounds the candidate set, a decision backend answers a fixed question set per candidate (choice, noul, score), a policy bands the action probability, the typed answers return as a Parquet table, and the recompute runs natively over the full data. Six cases over MerRec, Amazon reviews and EMBER with a registry cache, replay from captures with zero API calls, non-model controls and comparators. The backend client reads only JEV_API_URL and JEV_API_KEY; captures record the backend name and model id only. Registries, captures and derived data stay local (MerRec CC BY-NC). Co-Authored-By: Claude Fable 5.1 --- fixture/demo/decision_models/README.md | 105 ++++ .../cases/a_brand_resolution/01_probe.sql | 28 + .../cases/a_brand_resolution/02_context.sql | 31 + .../cases/a_brand_resolution/03_recompute.sql | 29 + .../a_brand_resolution/03a_canonical.sql | 6 + .../a_brand_resolution/compare_captures.py | 100 ++++ .../a_brand_resolution/control_trigram.py | 133 +++++ .../cases/a_brand_resolution/policy.json | 6 + .../cases/a_brand_resolution/questions.json | 30 + .../scripts/cases/a_brand_resolution/run.py | 198 +++++++ .../cases/b_bridge_purification/01_probe.sql | 45 ++ .../b_bridge_purification/02_context.sql | 40 ++ .../b_bridge_purification/03_recompute.sql | 47 ++ .../04_recompute_canonical.sql | 48 ++ .../b_bridge_purification/compare_captures.py | 88 +++ .../control_denylist.json | 17 + .../cases/b_bridge_purification/policy.json | 1 + .../b_bridge_purification/questions.json | 30 + .../cases/b_bridge_purification/run.py | 307 ++++++++++ .../cases/c_complaint_causes/01_probe.sql | 47 ++ .../cases/c_complaint_causes/02_context.sql | 15 + .../cases/c_complaint_causes/03_recompute.sql | 25 + .../04_product_cause_shares.sql | 25 + .../c_complaint_causes/05_cause_shares.sql | 22 + .../06_anonymous_ranking.sql | 28 + .../07_low_rating_ranking.sql | 27 + .../08_helpful_weighted.sql | 31 + .../c_complaint_causes/09_safety_clusters.sql | 21 + .../c_complaint_causes/compare_captures.py | 115 ++++ .../c_complaint_causes/control_keywords.sql | 62 ++ .../control_keywords_products.sql | 90 +++ .../cases/c_complaint_causes/policy.json | 7 + .../cases/c_complaint_causes/questions.json | 40 ++ .../scripts/cases/c_complaint_causes/run.py | 428 ++++++++++++++ .../cases/d_barbie_mechanisms/01_probe.sql | 45 ++ .../d_barbie_mechanisms/02_bridge_items.sql | 39 ++ .../d_barbie_mechanisms/03_counterfactual.sql | 49 ++ .../d_barbie_mechanisms/04_decomposition.sql | 22 + .../d_barbie_mechanisms/compare_captures.py | 93 +++ .../cases/d_barbie_mechanisms/policy.json | 1 + .../cases/d_barbie_mechanisms/questions.json | 28 + .../scripts/cases/d_barbie_mechanisms/run.py | 345 +++++++++++ .../cases/e_pokemon_submarkets/01_probe.sql | 22 + .../02_forward_validation.sql | 44 ++ .../e_pokemon_submarkets/03_control_rules.sql | 24 + .../e_pokemon_submarkets/compare_captures.py | 93 +++ .../cases/e_pokemon_submarkets/policy.json | 1 + .../cases/e_pokemon_submarkets/questions.json | 38 ++ .../scripts/cases/e_pokemon_submarkets/run.py | 340 +++++++++++ .../scripts/cases/f_import_gate/01_probe.sql | 31 + .../cases/f_import_gate/02_coverage.sql | 21 + .../f_import_gate/03_precision_pairs.sql | 28 + .../cases/f_import_gate/04_seed_neighbors.sql | 24 + .../cases/f_import_gate/05_exact_pairs.sql | 23 + .../cases/f_import_gate/compare_captures.py | 93 +++ .../cases/f_import_gate/control_dll_list.json | 15 + .../scripts/cases/f_import_gate/policy.json | 1 + .../cases/f_import_gate/prepare_tables.py | 152 +++++ .../cases/f_import_gate/questions.json | 29 + .../scripts/cases/f_import_gate/run.py | 430 ++++++++++++++ .../cases/f_import_gate/variant_edges.sql | 14 + .../scripts/client/__init__.py | 8 + .../decision_models/scripts/client/algeon.py | 166 ++++++ .../decision_models/scripts/client/capture.py | 179 ++++++ .../decision_models/scripts/client/case.py | 124 ++++ .../scripts/client/decision_client.py | 553 ++++++++++++++++++ .../decision_models/scripts/client/policy.py | 48 ++ .../scripts/client/registry.py | 150 +++++ .../decision_models/scripts/tests/__init__.py | 0 .../scripts/tests/test_client.py | 240 ++++++++ 70 files changed, 5785 insertions(+) create mode 100644 fixture/demo/decision_models/README.md create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/01_probe.sql create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/02_context.sql create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/03_recompute.sql create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/03a_canonical.sql create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/compare_captures.py create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/control_trigram.py create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/policy.json create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/questions.json create mode 100644 fixture/demo/decision_models/scripts/cases/a_brand_resolution/run.py create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/01_probe.sql create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/02_context.sql create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/03_recompute.sql create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/04_recompute_canonical.sql create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/compare_captures.py create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/control_denylist.json create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/policy.json create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/questions.json create mode 100644 fixture/demo/decision_models/scripts/cases/b_bridge_purification/run.py create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/01_probe.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/02_context.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/03_recompute.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/04_product_cause_shares.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/05_cause_shares.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/06_anonymous_ranking.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/07_low_rating_ranking.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/08_helpful_weighted.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/09_safety_clusters.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/compare_captures.py create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/control_keywords.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/control_keywords_products.sql create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/policy.json create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/questions.json create mode 100644 fixture/demo/decision_models/scripts/cases/c_complaint_causes/run.py create mode 100644 fixture/demo/decision_models/scripts/cases/d_barbie_mechanisms/01_probe.sql create mode 100644 fixture/demo/decision_models/scripts/cases/d_barbie_mechanisms/02_bridge_items.sql create mode 100644 fixture/demo/decision_models/scripts/cases/d_barbie_mechanisms/03_counterfactual.sql create mode 100644 fixture/demo/decision_models/scripts/cases/d_barbie_mechanisms/04_decomposition.sql create mode 100644 fixture/demo/decision_models/scripts/cases/d_barbie_mechanisms/compare_captures.py create mode 100644 fixture/demo/decision_models/scripts/cases/d_barbie_mechanisms/policy.json create mode 100644 fixture/demo/decision_models/scripts/cases/d_barbie_mechanisms/questions.json create mode 100644 fixture/demo/decision_models/scripts/cases/d_barbie_mechanisms/run.py create mode 100644 fixture/demo/decision_models/scripts/cases/e_pokemon_submarkets/01_probe.sql create mode 100644 fixture/demo/decision_models/scripts/cases/e_pokemon_submarkets/02_forward_validation.sql create mode 100644 fixture/demo/decision_models/scripts/cases/e_pokemon_submarkets/03_control_rules.sql create mode 100644 fixture/demo/decision_models/scripts/cases/e_pokemon_submarkets/compare_captures.py create mode 100644 fixture/demo/decision_models/scripts/cases/e_pokemon_submarkets/policy.json create mode 100644 fixture/demo/decision_models/scripts/cases/e_pokemon_submarkets/questions.json create mode 100644 fixture/demo/decision_models/scripts/cases/e_pokemon_submarkets/run.py create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/01_probe.sql create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/02_coverage.sql create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/03_precision_pairs.sql create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/04_seed_neighbors.sql create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/05_exact_pairs.sql create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/compare_captures.py create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/control_dll_list.json create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/policy.json create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/prepare_tables.py create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/questions.json create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/run.py create mode 100644 fixture/demo/decision_models/scripts/cases/f_import_gate/variant_edges.sql create mode 100644 fixture/demo/decision_models/scripts/client/__init__.py create mode 100644 fixture/demo/decision_models/scripts/client/algeon.py create mode 100644 fixture/demo/decision_models/scripts/client/capture.py create mode 100644 fixture/demo/decision_models/scripts/client/case.py create mode 100644 fixture/demo/decision_models/scripts/client/decision_client.py create mode 100644 fixture/demo/decision_models/scripts/client/policy.py create mode 100644 fixture/demo/decision_models/scripts/client/registry.py create mode 100644 fixture/demo/decision_models/scripts/tests/__init__.py create mode 100644 fixture/demo/decision_models/scripts/tests/test_client.py diff --git a/fixture/demo/decision_models/README.md b/fixture/demo/decision_models/README.md new file mode 100644 index 000000000..b3c3f8d98 --- /dev/null +++ b/fixture/demo/decision_models/README.md @@ -0,0 +1,105 @@ +# Decision models in the loop + +Demo of one orchestration pattern over the existing Algeon server: a probe +query returns a bounded candidate set with its context, a decision model +answers a fixed question set for each candidate, the typed answers go back +into Algeon as a Parquet table, and a recompute query runs over the full data +with the corrected structure. Measured results are not part of the +repository: each case's `run.py --mode replay` reproduces them from a capture +file, and the capture files stay local (see "Layout"). + +## Loop contract + +- Probe and recompute are separate Flight SQL queries. No GPU grant is held + while the decision backend runs. +- The probe's row limit (`max_entities` in each case) is the only cost bound. +- Decisions are cached in a registry keyed by the SHA-256 of the entity state + text and of the canonical question JSON. A second run of the same case makes + zero backend calls; a new time window calls the backend only for entities + the registry has not seen. +- The recompute joins the typed decision table (`CREATE EXTERNAL TABLE ... + STORED AS PARQUET`) and records `host_datafusion_boundary_count` from the + query stats together with the `EXPLAIN GPU` summary. +- Each case has a control that uses no model, and a `policy.json` with the + accept, review and reject bands on the action `noul`. + +## Layout + +``` +fixture/demo/decision_models/ + README.md + scripts/client/ decision_client.py registry.py policy.py capture.py algeon.py case.py + scripts/cases/ a_brand_resolution/ b_bridge_purification/ c_complaint_causes/ + d_barbie_mechanisms/ e_pokemon_submarkets/ f_import_gate/ + scripts/tests/ unit tests (no server, no API) + registry/ *.parquet registries and *.engine.parquet decision tables (local only) + captures/ [_