diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 520f4e2ec..d0a6b5e0b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -43,7 +43,7 @@ jobs: | jq -r '[.packages[].rust_version | select(. != null)] | unique | .[0]')" --profile minimal rustup default stable - - uses: taiki-e/install-action@ab6895351a88b908e96a4697860d4ffea5e57dbe # v2.87.27 + - uses: taiki-e/install-action@8cca0918f8471c33515202b73f56ab27f2bed283 # v2.87.29 with: tool: cargo-nextest,cargo-hack@0.6.45 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/fixture/demo/decision_models/README.md b/fixture/demo/decision_models/README.md new file mode 100644 index 000000000..9935a0fff --- /dev/null +++ b/fixture/demo/decision_models/README.md @@ -0,0 +1,109 @@ +# 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) + scripts/export_website.py writes website/static/data/merrec-decision-{bridges,twins}.json from the A and B captures + registry/ *.parquet registries and *.engine.parquet decision tables (local only) + captures/ [_