From cf512516f06449bb01d5fd847e5f33d8ec4a8457 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 19 Oct 2026 14:55:13 +0100 Subject: [PATCH 1/3] perf: evaluate a hash join filter conjunct by conjunct, per input row where it can gcs_cost_insights_v2 joins object paths to table locations with a broadcast hash join on (bucket, parts[0]) and a filter of 20 conjuncts `isnull(loc.parts[i]) OR loc.parts[i] = obj.parts[i]`. One key holds 11.7k locations, so the join meets 25.7e9 candidate pairs a night. For every pair apply_join_filter_to_indices took both whole `parts` lists into an intermediate batch and evaluated all 40 GetArrayItem and the whole AND tree: 3.6 us per pair on Comet against <= 0.13 us in Spark, 26 task-h instead of 0.9 in one stage. HashJoinStream now evaluates the filter with LazyJoinFilter: - the top-level AND is split and each conjunct is evaluated only over the pairs the earlier ones kept, like Spark's short-circuiting AND; - the largest non-volatile subexpressions that read one input alone (`loc.parts[3]`, `isnull(loc.parts[3])`, `obj.parts[3]`) are lifted out and a pair only copies their results and the bare columns its conjunct reads; - lifted build-side expressions are evaluated per pair until the stream has evaluated as many pairs as there are build rows, then once over the build batch and cached for the rest of the stream; if that fails on some row, they stay per pair, so only rows a pair reaches can raise errors; - without the cache, lifted expressions are evaluated once per run of pairs sharing a row (a probe row with many matches), otherwise per pair over the columns they read, taken once per lookup and filtered with the survivors. Pairs keep their order, so outer/semi/anti/mark handling is unchanged. Filters reading no column, and the symmetric hash join, keep the old path. New bench spark-expr/benches/hash_join_array_filter.rs (Comet ListExtract, the 20-conjunct location filter, batch 8192), base -> this commit: skewed, 12k locations x 200 objects of one key (2.4M pairs): 131 -> 44 ms unique keys (200k pairs, all pass every conjunct): 19.2 -> 17.2 ms Co-Authored-By: Claude Opus 5.5 (1M context) --- native/spark-expr/Cargo.toml | 4 + .../benches/hash_join_array_filter.rs | 221 ++++ .../src/joins/hash_join/exec.rs | 2 +- .../src/joins/hash_join/lazy_filter.rs | 1072 +++++++++++++++++ .../src/joins/hash_join/mod.rs | 1 + .../src/joins/hash_join/stream.rs | 25 +- .../apache/comet/exec/CometJoinSuite.scala | 67 ++ 7 files changed, 1387 insertions(+), 5 deletions(-) create mode 100644 native/spark-expr/benches/hash_join_array_filter.rs create mode 100644 native/vendor/datafusion-physical-plan/src/joins/hash_join/lazy_filter.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index bb85467b9d6..421bf75a04d 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -447,6 +447,10 @@ harness = false name = "smj_interval_filter" harness = false +[[bench]] +name = "hash_join_array_filter" +harness = false + [[bench]] name = "smj_streamed_filter" harness = false diff --git a/native/spark-expr/benches/hash_join_array_filter.rs b/native/spark-expr/benches/hash_join_array_filter.rs new file mode 100644 index 00000000000..07dfe3f7f1b --- /dev/null +++ b/native/spark-expr/benches/hash_join_array_filter.rs @@ -0,0 +1,221 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Broadcast hash join with a filter over array elements, shaped like +//! gcs_cost_insights_v2's table location lookup: build rows are table +//! locations split into path parts, probe rows are object paths, the join +//! keys are the bucket and the first part, and the filter is +//! `AND` over `i in 1..21` of +//! `loc.parts[i] IS NULL OR loc.parts[i] = obj.parts[i]`. +//! Reports time per candidate pair in the `pairs` throughput. + +use arrow::array::{ArrayRef, ListBuilder, RecordBatch, StringArray, StringBuilder}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use datafusion::common::{JoinSide, JoinType, NullEquality}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{lit, BinaryExpr, Column, IsNullExpr}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; +use datafusion::physical_plan::{collect, ExecutionPlan}; +use datafusion::prelude::SessionConfig; +use datafusion_comet_spark_expr::{create_query_context_map, ListExtract}; +use std::sync::Arc; +use tokio::runtime::Runtime; + +const DEPTH: usize = 21; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("bucket", DataType::Utf8, false), + Field::new("first", DataType::Utf8, true), + Field::new( + "parts", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + Field::new("payload", DataType::Utf8, true), + ])) +} + +/// `rows` rows of one or `keys` keys, `parts` given per row +fn table(rows: usize, keys: usize, parts_of: impl Fn(usize) -> Vec) -> RecordBatch { + let mut parts = ListBuilder::new(StringBuilder::new()); + let mut first = vec![]; + for i in 0..rows { + let p = parts_of(i); + first.push(format!("k{}", i % keys)); + for s in &p { + parts.values().append_value(s); + } + parts.append(true); + } + RecordBatch::try_new( + schema(), + vec![ + Arc::new(StringArray::from(vec!["joom-analytics-ml-infra"; rows])) as ArrayRef, + Arc::new(StringArray::from(first)), + Arc::new(parts.finish()), + Arc::new(StringArray::from( + (0..rows) + .map(|i| format!("payload-{i}")) + .collect::>(), + )), + ], + ) + .unwrap() +} + +/// Locations `tables/db{d}/t{t}` (some one level deeper); objects +/// `tables/db{d}/t{t}/date=.../hour=.../part-*.parquet`, so most pairs fail on part 1 +/// or 2 and every object has one location. +fn inputs(locations: usize, objects: usize, keys: usize) -> (Vec, Vec) { + let location = |i: usize| { + let mut p = vec![ + "tables".to_string(), + format!("analytics_db_{}", i % 97), + format!("events_table_{}_daily", i / 97), + ]; + if i.is_multiple_of(5) { + p.push("v2".to_string()); + } + p + }; + let build = table(locations, keys, location); + let probe: Vec = (0..objects.div_ceil(8192)) + .map(|b| { + let n = 8192.min(objects - b * 8192); + table(n, keys, |j| { + let i = (b * 8192 + j) * 7919 % locations; + let mut p = location(i); + p.push(format!("date=2026-10-{:02}", j % 28 + 1)); + p.push(format!("hour={:02}", j % 24)); + p.push(format!( + "part-{j:05}-6d2f9a1c-0b7e-4c1f-9a39-{:012x}.c000.zstd.parquet", + j * 7919 + )); + p + }) + }) + .collect(); + (vec![build], probe) +} + +fn item(column: usize, i: usize) -> Arc { + Arc::new(ListExtract::new( + Arc::new(Column::new("parts", column)), + lit(i as i32), + None, + false, + false, + None, + create_query_context_map(), + )) +} + +/// Over `(loc.parts, obj.parts)` +fn location_filter() -> JoinFilter { + let mut expr: Option> = None; + for i in 1..DEPTH { + let term: Arc = Arc::new(BinaryExpr::new( + Arc::new(IsNullExpr::new(item(0, i))), + Operator::Or, + Arc::new(BinaryExpr::new(item(0, i), Operator::Eq, item(1, i))), + )); + expr = Some(match expr { + None => term, + Some(e) => Arc::new(BinaryExpr::new(e, Operator::And, term)), + }); + } + let parts = schema().field(2).clone(); + JoinFilter::new( + expr.unwrap(), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![parts.clone(), parts])), + ) +} + +fn make_exec(batches: &[RecordBatch]) -> Arc { + MemorySourceConfig::try_new_exec(&[batches.to_vec()], schema(), None).unwrap() +} + +fn bench_hash_join_array_filter(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let mut group = c.benchmark_group("hash_join_array_filter"); + group.sample_size(10); + for (shape, locations, objects, keys) in [ + ("skewed_12k_locations", 12_000, 200, 1), + ("unique_keys", 12_000, 200_000, 12_000), + ] { + let (build, probe) = inputs(locations, objects, keys); + let pairs = (locations / keys) * objects; + group.throughput(Throughput::Elements(pairs as u64)); + group.bench_function(BenchmarkId::new(shape, pairs), |b| { + b.iter(|| { + let s = schema(); + let on = vec![ + ( + Arc::new(Column::new_with_schema("bucket", &s).unwrap()) as _, + Arc::new(Column::new_with_schema("bucket", &s).unwrap()) as _, + ), + ( + Arc::new(Column::new_with_schema("first", &s).unwrap()) as _, + Arc::new(Column::new_with_schema("first", &s).unwrap()) as _, + ), + ]; + let join = HashJoinExec::try_new( + make_exec(&build), + make_exec(&probe), + on, + Some(location_filter()), + &JoinType::Right, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(8192)), + ); + let rows: usize = rt.block_on(async { + let batches = collect(Arc::new(join), task_ctx).await.unwrap(); + batches.iter().map(|b| b.num_rows()).sum() + }); + assert!(rows >= objects); + rows + }) + }); + } + group.finish(); +} + +criterion_group!(benches, bench_hash_join_array_filter); +criterion_main!(benches); diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/exec.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/exec.rs index 08d209003ad..22fb6fabe9f 100644 --- a/native/vendor/datafusion-physical-plan/src/joins/hash_join/exec.rs +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/exec.rs @@ -1576,7 +1576,7 @@ impl ExecutionPlan for HashJoinExec { self.mode, self.null_aware, self.fetch, - ))) + )?)) } fn metrics(&self) -> Option { diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/lazy_filter.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/lazy_filter.rs new file mode 100644 index 00000000000..9ef28c6763f --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/lazy_filter.rs @@ -0,0 +1,1072 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Hash join filter evaluated conjunct by conjunct over the surviving pairs, +//! with the subexpressions over one input alone evaluated per input row. + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, BooleanArray, PrimitiveArray, RecordBatch, + RecordBatchOptions, UInt32Array, UInt64Array, downcast_array, new_null_array, +}; +use arrow::compute::{self, filter_record_batch, take}; +use arrow::datatypes::{Field, Schema, SchemaRef}; +use datafusion_common::cast::as_boolean_array; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion_common::{JoinSide, Result, internal_err}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::utils::{collect_columns, split_conjunction}; +use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, is_volatile}; + +use crate::joins::utils::JoinFilter; + +/// A hash join filter `c1 AND c2 AND ...` evaluated one conjunct at a time, +/// each over the pairs every earlier conjunct kept, like Spark's +/// short-circuiting `AND`. +/// +/// The largest non-volatile subexpressions that read the columns of one input +/// alone (e.g. `parts[3]` or `parts[3] IS NULL` of one side) are lifted out +/// and evaluated per input row rather than per pair: +/// - once per run of pairs sharing a row (a probe row with many matches); +/// - over the build batch once the stream has evaluated as many build-side +/// pairs as it has rows, cached for the rest of the stream; +/// - otherwise per pair, over the columns they read taken once per lookup. +/// +/// A conjunct copies per pair only the lifted results and bare columns it +/// reads. +#[derive(Debug)] +pub(super) struct LazyJoinFilter { + /// The conjuncts in their original order + conjuncts: Vec, + build: Side, + probe: Side, + build_cache: BuildCache, +} + +/// The lifted subexpressions of one input +#[derive(Debug)] +struct Side { + inputs: Vec, + /// The input's columns the inputs that are not bare columns read + projection: Vec, + /// Schema of `projection` + schema: SchemaRef, +} + +#[derive(Debug)] +struct SideInput { + /// Over `Side::schema` + expr: PhysicalExprRef, + /// Set when the input is a bare column of the input, to its index + column: Option, +} + +#[derive(Debug)] +struct Conjunct { + /// Over `schema` + expr: PhysicalExprRef, + schema: SchemaRef, + /// Source of each column of `schema` + inputs: Vec, +} + +#[derive(Debug, Clone, Copy)] +enum InputRef { + Build(usize), + Probe(usize), +} + +#[derive(Debug, Default)] +struct BuildCache { + /// The arrays of the build batch the fields below are for, kept alive so + /// the batch is recognized by pointer + key: Vec, + /// Pairs whose build-side inputs were evaluated without the cache + pairs_evaluated: usize, + /// Results of the build-side inputs that are not bare columns over every + /// build row + values: Option>>, + /// Evaluating them over every build row failed on some row; they are + /// evaluated for pairs only from then on, so only pairs raise errors + failed: bool, +} + +/// One side of the current pairs of an `apply` +struct Pairs<'a, T: ArrowPrimitiveType> { + batch: &'a RecordBatch, + indices: PrimitiveArray, + /// `Side::projection` taken for `indices`, once needed + projection: Option, +} + +impl LazyJoinFilter { + /// Returns `None` when the filter reads no column, or reads a column of + /// neither input. + pub fn try_new(filter: &JoinFilter, build_side: JoinSide) -> Result> { + let column_indices = filter.column_indices(); + if column_indices.is_empty() + || column_indices + .iter() + .any(|c| !matches!(c.side, JoinSide::Left | JoinSide::Right)) + { + return Ok(None); + } + let filter_schema = filter.schema(); + + // Lift the largest single-input subexpressions, bare columns included, + // numbered over the whole filter. + let mut lifted: Vec<(bool, PhysicalExprRef)> = vec![]; + let expression = Arc::clone(filter.expression()) + .transform_down(|expr| { + let columns = collect_columns(&expr); + let Some(side) = columns + .iter() + .next() + .map(|c| column_indices[c.index()].side) + else { + return Ok(Transformed::no(expr)); + }; + if columns + .iter() + .any(|c| column_indices[c.index()].side != side) + || (expr.downcast_ref::().is_none() && is_volatile(&expr)) + { + return Ok(Transformed::no(expr)); + } + let is_build = side == build_side; + let position = match lifted + .iter() + .position(|(b, e)| *b == is_build && **e == *expr) + { + Some(position) => position, + None => { + lifted.push((is_build, Arc::clone(&expr))); + lifted.len() - 1 + } + }; + Ok(Transformed::new( + Arc::new(Column::new(&format!("__lifted_{position}"), position)), + true, + TreeNodeRecursion::Jump, + )) + })? + .data; + + let fields = lifted + .iter() + .enumerate() + .map(|(position, (_, expr))| { + Ok(Field::new( + format!("__lifted_{position}"), + expr.data_type(filter_schema)?, + expr.nullable(filter_schema)?, + )) + }) + .collect::>>()?; + let side = |is_build: bool| -> Result<(Side, Vec>)> { + // Filter columns the side's non-bare-column inputs read + let mut read: Vec = lifted + .iter() + .filter(|(b, e)| *b == is_build && e.downcast_ref::().is_none()) + .flat_map(|(_, e)| collect_columns(e)) + .map(|c| c.index()) + .collect(); + read.sort_unstable(); + read.dedup(); + let mut projection: Vec = + read.iter().map(|&i| column_indices[i].index).collect(); + projection.sort_unstable(); + projection.dedup(); + let mut inputs = vec![]; + let mut positions = vec![None; lifted.len()]; + for (position, (b, expr)) in lifted.iter().enumerate() { + if *b != is_build { + continue; + } + let column = expr + .downcast_ref::() + .map(|c| column_indices[c.index()].index); + let expr = Arc::clone(expr) + .transform_up(|e| { + let Some(c) = e.downcast_ref::() else { + return Ok(Transformed::no(e)); + }; + let index = column_indices[c.index()].index; + let index = projection.binary_search(&index).unwrap_or(0); + Ok(Transformed::yes( + Arc::new(Column::new(c.name(), index)) as PhysicalExprRef + )) + })? + .data; + positions[position] = Some(inputs.len()); + inputs.push(SideInput { expr, column }); + } + let schema = Arc::new(Schema::new( + projection + .iter() + .map(|&index| { + let i = read + .iter() + .copied() + .find(|&i| column_indices[i].index == index) + .unwrap(); + filter_schema.field(i).clone() + }) + .collect::>(), + )); + Ok(( + Side { + inputs, + projection, + schema, + }, + positions, + )) + }; + let (build, build_positions) = side(true)?; + let (probe, probe_positions) = side(false)?; + let lifted_inputs: Vec = (0..lifted.len()) + .map(|i| match build_positions[i] { + Some(j) => InputRef::Build(j), + None => InputRef::Probe(probe_positions[i].unwrap()), + }) + .collect(); + + let conjuncts = split_conjunction(&expression) + .into_iter() + .map(|conjunct| { + let mut used: Vec = collect_columns(conjunct) + .iter() + .map(|c| c.index()) + .collect(); + used.sort_unstable(); + used.dedup(); + let expr = Arc::clone(conjunct) + .transform_up(|e| { + let Some(column) = e.downcast_ref::() else { + return Ok(Transformed::no(e)); + }; + let index = used.binary_search(&column.index()).unwrap(); + Ok(Transformed::yes( + Arc::new(Column::new(column.name(), index)) as PhysicalExprRef, + )) + })? + .data; + Ok(Conjunct { + expr, + schema: Arc::new(Schema::new( + used.iter().map(|&i| fields[i].clone()).collect::>(), + )), + inputs: used.iter().map(|&i| lifted_inputs[i]).collect(), + }) + }) + .collect::>>()?; + + Ok(Some(Self { + conjuncts, + build, + probe, + build_cache: BuildCache::default(), + })) + } + + /// Keeps the pairs for which every conjunct is true, in their order. + pub fn apply( + &mut self, + build_batch: &RecordBatch, + probe_batch: &RecordBatch, + build_indices: UInt64Array, + probe_indices: UInt32Array, + ) -> Result<(UInt64Array, UInt32Array)> { + let mut build_pairs = Pairs { + batch: build_batch, + indices: build_indices, + projection: None, + }; + let mut probe_pairs = Pairs { + batch: probe_batch, + indices: probe_indices, + projection: None, + }; + for k in 0..self.conjuncts.len() { + let num_pairs = build_pairs.indices.len(); + if num_pairs == 0 { + break; + } + let mask = self.evaluate_conjunct(k, &mut build_pairs, &mut probe_pairs)?; + if mask.null_count() == 0 && mask.true_count() == num_pairs { + continue; + } + build_pairs.retain(&mask)?; + probe_pairs.retain(&mask)?; + } + Ok((build_pairs.indices, probe_pairs.indices)) + } + + fn evaluate_conjunct( + &mut self, + k: usize, + build_pairs: &mut Pairs<'_, arrow::datatypes::UInt64Type>, + probe_pairs: &mut Pairs<'_, arrow::datatypes::UInt32Type>, + ) -> Result { + let num_pairs = build_pairs.indices.len(); + let conjunct = &self.conjuncts[k]; + if conjunct + .inputs + .iter() + .any(|i| matches!(i, InputRef::Build(j) if self.build.inputs[*j].column.is_none())) + { + self.build_cache + .prepare(&self.build, build_pairs.batch, num_pairs); + } + + let mut columns = Vec::with_capacity(conjunct.inputs.len()); + for input in &conjunct.inputs { + let column = match *input { + InputRef::Build(j) => match self.build_cache.values(j, build_pairs.batch) { + Some(values) => take_or_nulls(values, &build_pairs.indices)?, + None => build_pairs.evaluate(&self.build, j)?, + }, + InputRef::Probe(j) => probe_pairs.evaluate(&self.probe, j)?, + }; + columns.push(column); + } + + let batch = projected_batch(&conjunct.schema, columns, num_pairs)?; + let result = conjunct.expr.evaluate(&batch)?.into_array(num_pairs)?; + let mask = as_boolean_array(&result)?; + if mask.len() != num_pairs { + return internal_err!( + "join filter conjunct returned {} rows for {num_pairs} pairs", + mask.len() + ); + } + Ok(mask.clone()) + } +} + +impl BuildCache { + /// Evaluates the build-side inputs over every row of `batch` once the + /// pairs evaluated without the cache reach its number of rows. + fn prepare(&mut self, side: &Side, batch: &RecordBatch, num_pairs: usize) { + if !same_batch(&self.key, batch) { + *self = BuildCache { + key: batch.columns().to_vec(), + ..Default::default() + }; + } + if self.values.is_some() || self.failed { + return; + } + if self.pairs_evaluated + num_pairs < batch.num_rows() { + self.pairs_evaluated += num_pairs; + return; + } + let num_rows = batch.num_rows(); + let columns = side + .projection + .iter() + .map(|&c| Arc::clone(batch.column(c))) + .collect(); + let values = projected_batch(&side.schema, columns, num_rows).and_then(|batch| { + side.inputs + .iter() + .map(|input| match input.column { + Some(_) => Ok(None), + None => Ok(Some(input.expr.evaluate(&batch)?.into_array(num_rows)?)), + }) + .collect::>>() + }); + match values { + Ok(values) => self.values = Some(values), + Err(_) => self.failed = true, + } + } + + /// The results of input `j` over every row of `batch`, when cached + fn values(&self, j: usize, batch: &RecordBatch) -> Option<&ArrayRef> { + self.values + .as_ref() + .filter(|_| same_batch(&self.key, batch)) + .and_then(|values| values[j].as_ref()) + } +} + +impl Pairs<'_, T> { + /// Input `j` of `side` for every pair: a bare column taken per pair, + /// otherwise evaluated once per run of pairs of one row, or per pair over + /// the side's columns taken once for all inputs. + fn evaluate(&mut self, side: &Side, j: usize) -> Result { + let input = &side.inputs[j]; + if let Some(column) = input.column { + return take_or_nulls(self.batch.column(column), &self.indices); + } + if let Some((rows, positions)) = runs_of(&self.indices) { + let num_rows = rows.len(); + let columns = side + .projection + .iter() + .map(|&c| take_or_nulls(self.batch.column(c), &rows)) + .collect::>>()?; + let values = input + .expr + .evaluate(&projected_batch(&side.schema, columns, num_rows)?)? + .into_array(num_rows)?; + return Ok(take(values.as_ref(), &positions, None)?); + } + if self.projection.is_none() { + let columns = side + .projection + .iter() + .map(|&c| take_or_nulls(self.batch.column(c), &self.indices)) + .collect::>>()?; + self.projection = Some(projected_batch(&side.schema, columns, self.indices.len())?); + } + let projection = self.projection.as_ref().unwrap(); + input + .expr + .evaluate(projection)? + .into_array(projection.num_rows()) + } + + /// Keeps the pairs of `mask` + fn retain(&mut self, mask: &BooleanArray) -> Result<()> { + self.indices = downcast_array(compute::filter(&self.indices, mask)?.as_ref()); + if let Some(projection) = &self.projection { + self.projection = Some(filter_record_batch(projection, mask)?); + } + Ok(()) + } +} + +/// Whether `batch` holds the very arrays of `key` +fn same_batch(key: &[ArrayRef], batch: &RecordBatch) -> bool { + key.len() == batch.num_columns() + && key + .iter() + .zip(batch.columns()) + .all(|(a, b)| Arc::ptr_eq(a, b)) +} + +fn take_or_nulls(array: &ArrayRef, indices: &I) -> Result { + if array.is_empty() || indices.null_count() == indices.len() { + return Ok(new_null_array(array.data_type(), indices.len())); + } + Ok(take(array.as_ref(), indices, None)?) +} + +/// For indices with runs of equal consecutive values averaging two pairs or +/// more: the value of each run and, per pair, the position of its run. +fn runs_of( + indices: &PrimitiveArray, +) -> Option<(PrimitiveArray, UInt32Array)> { + if indices.null_count() > 0 || indices.len() < 2 { + return None; + } + let values = indices.values(); + let num_runs = 1 + values.windows(2).filter(|w| w[0] != w[1]).count(); + if num_runs * 2 > values.len() { + return None; + } + let mut rows = Vec::with_capacity(num_runs); + let mut positions = Vec::with_capacity(values.len()); + for &value in values.iter() { + if rows.last() != Some(&value) { + rows.push(value); + } + positions.push((rows.len() - 1) as u32); + } + Some(( + PrimitiveArray::::from_iter_values(rows), + UInt32Array::from(positions), + )) +} + +fn projected_batch( + schema: &SchemaRef, + columns: Vec, + num_rows: usize, +) -> Result { + Ok(RecordBatch::try_new_with_options( + Arc::clone(schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(num_rows)), + )?) +} + +#[cfg(test)] +mod tests { + use std::fmt::{self, Display, Formatter}; + use std::hash::{Hash, Hasher}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use arrow::array::{AsArray, Int32Array, ListBuilder, StringArray, StringBuilder}; + use arrow::datatypes::DataType; + use datafusion_common::ScalarValue; + use datafusion_common::{JoinType, NullEquality}; + use datafusion_execution::TaskContext; + use datafusion_expr::ColumnarValue; + use datafusion_expr::Operator; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions::{ + BinaryExpr, CaseExpr, CastExpr, IsNullExpr, Literal, binary, lit, + }; + + use super::*; + use crate::common; + use crate::joins::utils::{ColumnIndex, apply_join_filter_to_indices}; + use crate::joins::{HashJoinExec, NestedLoopJoinExec, PartitionMode}; + use crate::test::TestMemoryExec; + use crate::{ExecutionPlan, execution_plan::collect}; + + /// `child[index]` over a list of strings, null when out of range; counts + /// the rows it is evaluated over + #[derive(Debug)] + struct Item { + child: PhysicalExprRef, + index: usize, + rows: Arc, + } + + impl Hash for Item { + fn hash(&self, state: &mut H) { + self.child.hash(state); + self.index.hash(state); + } + } + + impl PartialEq for Item { + fn eq(&self, other: &Self) -> bool { + self.child.eq(&other.child) && self.index == other.index + } + } + + impl Eq for Item {} + + impl Display for Item { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{}[{}]", self.child, self.index) + } + } + + impl PhysicalExpr for Item { + fn data_type(&self, _: &Schema) -> Result { + Ok(DataType::Utf8) + } + + fn nullable(&self, _: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let lists = self.child.evaluate(batch)?.into_array(batch.num_rows())?; + let lists = lists.as_list::(); + self.rows.fetch_add(lists.len(), Ordering::Relaxed); + let values: StringArray = (0..lists.len()) + .map(|i| { + if lists.is_null(i) { + return None; + } + let list = lists.value(i); + let list = list.as_string::(); + (self.index < list.len() && list.is_valid(self.index)) + .then(|| list.value(self.index).to_string()) + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(values))) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + Ok(Arc::new(Item { + child: Arc::clone(&children[0]), + index: self.index, + rows: Arc::clone(&self.rows), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } + } + + struct Rng(u64); + + impl Rng { + fn next(&mut self, n: u64) -> u64 { + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (self.0 >> 33) % n + } + } + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, true), + Field::new( + "parts", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + Field::new("x", DataType::Int32, true), + Field::new("s", DataType::Utf8, true), + ])) + } + + /// `k` skewed to 0, `parts` of 0..6 strings over a small alphabet with + /// nulls, nullable `x`, and `s` a number as text, `"bad"` on some rows + fn table(rng: &mut Rng, rows: usize, bad_s: bool) -> RecordBatch { + let mut parts = ListBuilder::new(StringBuilder::new()); + for _ in 0..rows { + if rng.next(20) == 0 { + parts.append_null(); + continue; + } + for _ in 0..rng.next(7) { + match rng.next(5) { + 0 => parts.values().append_null(), + v => parts + .values() + .append_value(["a", "b", "c", "d"][v as usize - 1]), + } + } + parts.append(true); + } + let k: Int32Array = (0..rows) + .map(|_| (rng.next(10) != 0).then(|| (rng.next(10) as i32).saturating_sub(6))) + .collect(); + let x: Int32Array = (0..rows) + .map(|_| (rng.next(8) != 0).then(|| rng.next(20) as i32)) + .collect(); + let s: StringArray = (0..rows) + .map(|i| match rng.next(10) { + 0 => None, + 1 if bad_s && i % 2 == 1 => Some("bad".to_string()), + v => Some(v.to_string()), + }) + .collect(); + RecordBatch::try_new( + schema(), + vec![ + Arc::new(k), + Arc::new(parts.finish()) as ArrayRef, + Arc::new(x), + Arc::new(s), + ], + ) + .unwrap() + } + + fn col(name: &str, index: usize) -> PhysicalExprRef { + Arc::new(Column::new(name, index)) + } + + fn item(child: PhysicalExprRef, index: usize, rows: &Arc) -> PhysicalExprRef { + Arc::new(Item { + child, + index, + rows: Arc::clone(rows), + }) + } + + fn and(l: PhysicalExprRef, r: PhysicalExprRef) -> PhysicalExprRef { + Arc::new(BinaryExpr::new(l, Operator::And, r)) + } + + fn or(l: PhysicalExprRef, r: PhysicalExprRef) -> PhysicalExprRef { + Arc::new(BinaryExpr::new(l, Operator::Or, r)) + } + + fn op(l: PhysicalExprRef, o: Operator, r: PhysicalExprRef) -> PhysicalExprRef { + Arc::new(BinaryExpr::new(l, o, r)) + } + + fn is_null(e: PhysicalExprRef) -> PhysicalExprRef { + Arc::new(IsNullExpr::new(e)) + } + + /// Strict (erroring) string to int cast + fn to_int(e: PhysicalExprRef) -> PhysicalExprRef { + Arc::new(CastExpr::new(e, DataType::Int32, None)) + } + + /// Filter columns, in this order, over left (build) and right (probe) + /// inputs of `schema()`: l.k, r.k, l.parts, r.parts, l.x, r.x, l.s, r.s + fn filter(expr: PhysicalExprRef) -> JoinFilter { + let s = schema(); + let mut column_indices = vec![]; + let mut fields = vec![]; + for index in 0..4 { + for side in [JoinSide::Left, JoinSide::Right] { + column_indices.push(ColumnIndex { index, side }); + fields.push(s.field(index).clone()); + } + } + JoinFilter::new(expr, column_indices, Arc::new(Schema::new(fields))) + } + + /// The table location filter of gcs_cost_insights_v2: `AND` over + /// `i in 1..=depth` of `l.parts[i] IS NULL OR l.parts[i] = r.parts[i]` + fn parts_filter(depth: usize, rows: &Arc) -> PhysicalExprRef { + let mut e: Option = None; + for i in 1..=depth { + let l = item(col("l.parts", 2), i, rows); + let r = item(col("r.parts", 3), i, rows); + let term = or(is_null(Arc::clone(&l)), op(l, Operator::Eq, r)); + e = Some(match e { + None => term, + Some(e) => and(e, term), + }); + } + e.unwrap() + } + + fn filters(rows: &Arc) -> Vec<(&'static str, PhysicalExprRef)> { + let x_l = || col("l.x", 4); + let x_r = || col("r.x", 5); + vec![ + ("parts", parts_filter(5, rows)), + ( + "parts_and_x", + and(parts_filter(3, rows), op(x_l(), Operator::Lt, x_r())), + ), + ( + "or_top", + or( + op(op(x_l(), Operator::Plus, x_r()), Operator::Gt, lit(15i32)), + is_null(x_l()), + ), + ), + ( + "case", + and( + Arc::new( + CaseExpr::try_new( + None, + vec![( + op(x_r(), Operator::Gt, lit(9i32)), + op(x_l(), Operator::Lt, x_r()), + )], + Some(is_null(item(col("l.parts", 2), 0, rows))), + ) + .unwrap(), + ), + op(x_r(), Operator::NotEq, lit(7i32)), + ), + ), + ( + "nullable_conjuncts", + and( + and( + op(x_l(), Operator::GtEq, lit(3i32)), + op(x_l(), Operator::LtEq, x_r()), + ), + op( + item(col("l.parts", 2), 0, rows), + Operator::Eq, + item(col("r.parts", 3), 0, rows), + ), + ), + ), + ( + "literal_conjunct", + and( + lit(true), + and( + Arc::new(Literal::new(ScalarValue::Boolean(None))), + op(x_l(), Operator::Eq, x_r()), + ), + ), + ), + ("probe_only", op(x_r(), Operator::Gt, lit(4i32))), + ( + "build_only_and_probe_only", + and( + op(x_l(), Operator::Lt, lit(12i32)), + is_null(item(col("r.parts", 3), 2, rows)), + ), + ), + ( + "strict_cast_of_unpaired_build_rows", + and( + op(x_l(), Operator::Gt, lit(-1i32)), + op(to_int(col("l.s", 6)), Operator::Lt, x_r()), + ), + ), + ] + } + + /// Candidate pairs grouped by probe row, as the hash join produces them, + /// with a long run of one probe row + fn pairs(rng: &mut Rng, build_rows: usize, probe_rows: usize) -> (UInt64Array, UInt32Array) { + let mut b = vec![]; + let mut p = vec![]; + for probe in 0..probe_rows as u32 { + let n = if probe == 3 { + 3 * build_rows + } else { + rng.next(6) as usize + }; + for _ in 0..n { + b.push(rng.next(build_rows as u64)); + p.push(probe); + } + } + (UInt64Array::from(b), UInt32Array::from(p)) + } + + fn old_path( + build: &RecordBatch, + probe: &RecordBatch, + b: &UInt64Array, + p: &UInt32Array, + filter: &JoinFilter, + ) -> Result<(UInt64Array, UInt32Array)> { + apply_join_filter_to_indices( + build, + probe, + b.clone(), + p.clone(), + filter, + JoinSide::Left, + None, + JoinType::Inner, + ) + } + + /// Pairs whose build row has a non-numeric `s` + fn without_bad_build_s( + build: &RecordBatch, + b: &UInt64Array, + p: &UInt32Array, + ) -> (UInt64Array, UInt32Array) { + let s = build.column(3).as_string::(); + let keep: BooleanArray = b + .values() + .iter() + .map(|&i| Some(!(s.is_valid(i as usize) && s.value(i as usize) == "bad"))) + .collect(); + ( + downcast_array(compute::filter(b, &keep).unwrap().as_ref()), + downcast_array(compute::filter(p, &keep).unwrap().as_ref()), + ) + } + + #[test] + fn keeps_the_pairs_of_the_unlifted_filter_in_order() -> Result<()> { + let rows = Arc::new(AtomicUsize::new(0)); + let mut rng = Rng(7); + let build = table(&mut rng, 300, true); + for (name, expr) in filters(&rows) { + let filter = filter(expr); + let mut lazy = LazyJoinFilter::try_new(&filter, JoinSide::Left)?.unwrap(); + // Several calls over the same build batch: the first ones before, + // the later ones after the build-side cache is filled. + for call in 0..6 { + let probe = table(&mut rng, 40 + call * 10, false); + let (b, p) = pairs(&mut rng, build.num_rows(), probe.num_rows()); + let (b, p) = without_bad_build_s(&build, &b, &p); + let expected = old_path(&build, &probe, &b, &p, &filter)?; + let actual = lazy.apply(&build, &probe, b.clone(), p.clone())?; + assert_eq!(actual, expected, "filter {name}, call {call}"); + // Unsorted pairs: no runs to share + let order: UInt32Array = (0..b.len() as u32) + .map(|i| (i * 7919) % b.len() as u32) + .collect(); + if b.len() > 1 && b.len() % 7919 != 0 { + let b = downcast_array(take(&b, &order, None)?.as_ref()); + let p = downcast_array(take(&p, &order, None)?.as_ref()); + let expected = old_path(&build, &probe, &b, &p, &filter)?; + let actual = lazy.apply(&build, &probe, b, p)?; + assert_eq!(actual, expected, "filter {name}, call {call}, shuffled"); + } + } + } + Ok(()) + } + + #[test] + fn errors_only_for_pairs_the_filter_reaches() -> Result<()> { + let rows = Arc::new(AtomicUsize::new(0)); + let mut rng = Rng(11); + let build = table(&mut rng, 200, true); + let probe = table(&mut rng, 50, false); + let filter = filter(and( + op(col("r.x", 5), Operator::Lt, lit(100i32)), + op(to_int(col("l.s", 6)), Operator::Lt, col("r.x", 5)), + )); + let mut lazy = LazyJoinFilter::try_new(&filter, JoinSide::Left)?.unwrap(); + // Enough pairs to fill the build-side cache, which fails on "bad" + // build rows no pair reaches: no error, results as before. + for _ in 0..4 { + let (b, p) = pairs(&mut rng, build.num_rows(), probe.num_rows()); + let (b, p) = without_bad_build_s(&build, &b, &p); + let expected = old_path(&build, &probe, &b, &p, &filter)?; + assert_eq!(lazy.apply(&build, &probe, b, p)?, expected); + } + assert!(lazy.build_cache.failed); + // A pair of a "bad" build row still raises the cast error. + let s = build.column(3).as_string::(); + let bad = (0..build.num_rows()) + .find(|&i| s.is_valid(i) && s.value(i) == "bad") + .unwrap() as u64; + let b = UInt64Array::from(vec![0, bad]); + let p = UInt32Array::from(vec![0, 0]); + assert!(old_path(&build, &probe, &b, &p, &filter).is_err()); + assert!(lazy.apply(&build, &probe, b, p).is_err()); + let _ = rows; + Ok(()) + } + + #[test] + fn evaluates_lifted_subexpressions_per_row_not_per_pair() -> Result<()> { + let rows = Arc::new(AtomicUsize::new(0)); + let mut rng = Rng(5); + let build = table(&mut rng, 1000, false); + let probe = table(&mut rng, 4, false); + let filter = filter(parts_filter(20, &rows)); + let mut lazy = LazyJoinFilter::try_new(&filter, JoinSide::Left)?.unwrap(); + // 25 calls of 8000 pairs, all of probe row 1 (a skewed key) + let calls = 25; + let mut pairs_total = 0; + for call in 0..calls { + let b: UInt64Array = (0..8000u64).map(|i| (i * 13 + call) % 1000).collect(); + let p = UInt32Array::from(vec![1u32; 8000]); + pairs_total += b.len(); + let expected = old_path(&build, &probe, &b, &p, &filter)?; + let before = rows.load(Ordering::Relaxed); + let actual = lazy.apply(&build, &probe, b, p)?; + assert_eq!(actual, expected); + if call >= 1 { + // Cache filled: only one probe row per lifted probe item + assert!(rows.load(Ordering::Relaxed) - before <= 20, "call {call}"); + } + } + rows.store(0, Ordering::Relaxed); + let _ = old_path( + &build, + &probe, + &(0..8000u64).map(|i| i % 1000).collect(), + &UInt32Array::from(vec![1u32; 8000]), + &filter, + )?; + // The unlifted filter evaluates the items per pair (pre-selection of + // DataFusion's AND skips some of them) + assert!(rows.load(Ordering::Relaxed) >= 2 * 8000); + assert!(pairs_total > 0); + Ok(()) + } + + #[test] + fn not_used_without_filter_columns() -> Result<()> { + let f = JoinFilter::new(lit(true), vec![], Arc::new(Schema::empty())); + assert!(LazyJoinFilter::try_new(&f, JoinSide::Left)?.is_none()); + Ok(()) + } + + fn memory_exec(batches: Vec) -> Arc { + TestMemoryExec::try_new_exec(&[batches], schema(), None).unwrap() + } + + fn sorted_rows(batches: &[RecordBatch]) -> Vec { + let mut rows: Vec = arrow::util::pretty::pretty_format_batches(batches) + .unwrap() + .to_string() + .lines() + .filter(|l| l.starts_with('|')) + .skip(1) + .map(|l| l.to_string()) + .collect(); + rows.sort(); + rows + } + + /// The hash join with each filter, for every join type and batch size, + /// returns the rows of a nested loop join over `l.k = r.k AND filter`. + #[tokio::test] + async fn hash_join_matches_nested_loop_join_for_every_join_type() -> Result<()> { + let rows = Arc::new(AtomicUsize::new(0)); + let mut rng = Rng(3); + let left: Vec = (0..3).map(|_| table(&mut rng, 120, false)).collect(); + let right: Vec = (0..4).map(|_| table(&mut rng, 90, false)).collect(); + let s = schema(); + for (name, expr) in filters(&rows) { + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, + ] { + for batch_size in [7, 64, 8192] { + let ctx = Arc::new( + TaskContext::default().with_session_config( + datafusion_execution::config::SessionConfig::new() + .with_batch_size(batch_size), + ), + ); + let hash_join = HashJoinExec::try_new( + memory_exec(left.clone()), + memory_exec(right.clone()), + vec![( + Arc::new(Column::new_with_schema("k", &s)?) as _, + Arc::new(Column::new_with_schema("k", &s)?) as _, + )], + Some(filter(Arc::clone(&expr))), + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?; + let nested_loop_join = NestedLoopJoinExec::try_new( + memory_exec(left.clone()), + memory_exec(right.clone()), + Some(filter(and( + binary(col("l.k", 0), Operator::Eq, col("r.k", 1), &s)?, + Arc::clone(&expr), + ))), + &join_type, + None, + )?; + let expected = + sorted_rows(&collect(Arc::new(nested_loop_join), Arc::clone(&ctx)).await?); + let actual = sorted_rows(&common::collect(hash_join.execute(0, ctx)?).await?); + assert_eq!( + actual, expected, + "filter {name}, {join_type}, batch size {batch_size}" + ); + } + } + } + Ok(()) + } +} diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/mod.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/mod.rs index b915802ea40..95e3ddc575b 100644 --- a/native/vendor/datafusion-physical-plan/src/joins/hash_join/mod.rs +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/mod.rs @@ -22,6 +22,7 @@ pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState mod exec; mod inlist_builder; +mod lazy_filter; mod partitioned_hash_eval; mod shared_bounds; mod stream; diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/stream.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/stream.rs index 686939537e7..98708e0fabe 100644 --- a/native/vendor/datafusion-physical-plan/src/joins/hash_join/stream.rs +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/stream.rs @@ -29,6 +29,7 @@ use crate::joins::Map; use crate::joins::MapOffset; use crate::joins::PartitionMode; use crate::joins::hash_join::exec::JoinLeftData; +use crate::joins::hash_join::lazy_filter::LazyJoinFilter; use crate::joins::hash_join::shared_bounds::{ PartitionBounds, PartitionBuildData, SharedBuildAccumulator, }; @@ -301,6 +302,8 @@ pub(super) struct HashJoinStream { on_right: Vec, /// optional join filter filter: Option, + /// `filter` evaluated conjunct by conjunct, when it reads columns + lazy_filter: Option, /// type of the join (left, right, semi, etc) join_type: JoinType, /// right (probe) input @@ -487,16 +490,21 @@ impl HashJoinStream { mode: PartitionMode, null_aware: bool, fetch: Option, - ) -> Self { + ) -> Result { // Create output buffer with coalescing and optional fetch limit. let output_buffer = LimitedBatchCoalescer::new(Arc::clone(&schema), batch_size, fetch); + let lazy_filter = filter + .as_ref() + .and_then(|f| LazyJoinFilter::try_new(f, JoinSide::Left).transpose()) + .transpose()?; - Self { + Ok(Self { partition, schema, on_right, filter, + lazy_filter, join_type, right, random_state, @@ -514,7 +522,7 @@ impl HashJoinStream { mode, output_buffer, null_aware, - } + }) } /// Returns the next state after the build side has been fully collected @@ -848,7 +856,16 @@ impl HashJoinStream { .add_total(distinct_right_indices_count); // apply join filter if exists - let (left_indices, right_indices) = if let Some(filter) = &self.filter { + let (left_indices, right_indices) = if let Some(lazy_filter) = + self.lazy_filter.as_mut() + { + lazy_filter.apply( + build_side.left_data.batch(), + &state.batch, + left_indices, + right_indices, + )? + } else if let Some(filter) = &self.filter { apply_join_filter_to_indices( build_side.left_data.batch(), &state.batch, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 392ad516a00..d8f6f379a3e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -824,6 +824,73 @@ class CometJoinSuite extends CometTestBase { } } + test("Broadcast HashJoin with a join filter over array elements of both sides") { + withTempPath { dir => + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "100", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + // Table locations split into path parts, joined to object paths on the bucket and + // the first part with a prefix condition over the remaining parts; most locations + // share the first part "tables", so an object meets hundreds of candidates. + val objectsPath = s"${dir.getCanonicalPath}/objects" + spark + .range(0, 3000, 1, 2) + .selectExpr( + "id", + "concat('b', id % 2) AS bucket", + "CASE WHEN id % 31 = 0 THEN NULL " + + "WHEN id % 7 = 0 THEN concat('other/x', id % 5, '/f', id) " + + "ELSE concat('tables/db', id % 9, '/t', id % 13, '/date=', id % 3, " + + "CASE WHEN id % 4 = 0 THEN '' ELSE concat('/f', id) END) END AS key") + .selectExpr("*", "split(key, '/') AS parts") + .write + .parquet(objectsPath) + val locationsPath = s"${dir.getCanonicalPath}/locations" + spark + .range(0, 400, 1, 1) + .selectExpr( + "concat('db', id % 9, '_', id) AS table_name", + "concat('b', id % 2) AS bucket", + "CASE WHEN id % 53 = 0 THEN NULL WHEN id % 19 = 0 THEN 'other' " + + "WHEN id % 5 = 0 THEN concat('tables/db', id % 9) " + + "WHEN id % 5 = 1 THEN concat('tables/db', id % 9, '/t', id % 13, '/date=', id % 3) " + + "ELSE concat('tables/db', id % 9, '/t', id % 13) END AS prefix") + .selectExpr("*", "split(prefix, '/') AS parts") + .write + .parquet(locationsPath) + + withParquetTable(objectsPath, "objects") { + withParquetTable(locationsPath, "locations") { + val condition = ("o.bucket = l.bucket AND l.parts[0] = o.parts[0]" +: + (1 until 8).map(i => s"(l.parts[$i] IS NULL OR l.parts[$i] = o.parts[$i])")) + .mkString(" AND ") + val ops = Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec]) + for (joinType <- Seq("INNER", "LEFT")) { + checkSparkAnswerAndOperator( + sql( + "SELECT /*+ BROADCAST(l) */ o.id, o.key, l.table_name, size(l.parts) " + + s"FROM objects o $joinType JOIN locations l ON $condition"), + ops) + } + for (joinType <- Seq("LEFT SEMI", "LEFT ANTI")) { + checkSparkAnswerAndOperator( + sql( + "SELECT /*+ BROADCAST(l) */ o.id, o.key " + + s"FROM objects o $joinType JOIN locations l ON $condition"), + ops) + } + checkSparkAnswerAndOperator( + sql( + "SELECT /*+ BROADCAST(o) */ o.id, l.table_name " + + s"FROM objects o RIGHT JOIN locations l ON $condition"), + ops) + } + } + } + } + } + test("HashJoin without join filter") { withSQLConf( "spark.sql.join.forceApplyShuffledHashJoin" -> "true", From 6b2a6fdb53ea2415ad50d215281e2c26dcd0f3c6 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 19 Oct 2026 15:59:31 +0100 Subject: [PATCH 2/3] test: format the array join filter test Co-Authored-By: Claude Opus 5.5 (1M context) --- .../org/apache/comet/exec/CometJoinSuite.scala | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index d8f6f379a3e..0d0b080b1e5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -865,25 +865,23 @@ class CometJoinSuite extends CometTestBase { val condition = ("o.bucket = l.bucket AND l.parts[0] = o.parts[0]" +: (1 until 8).map(i => s"(l.parts[$i] IS NULL OR l.parts[$i] = o.parts[$i])")) .mkString(" AND ") - val ops = Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec]) + val ops = + Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec]) for (joinType <- Seq("INNER", "LEFT")) { checkSparkAnswerAndOperator( - sql( - "SELECT /*+ BROADCAST(l) */ o.id, o.key, l.table_name, size(l.parts) " + - s"FROM objects o $joinType JOIN locations l ON $condition"), + sql("SELECT /*+ BROADCAST(l) */ o.id, o.key, l.table_name, size(l.parts) " + + s"FROM objects o $joinType JOIN locations l ON $condition"), ops) } for (joinType <- Seq("LEFT SEMI", "LEFT ANTI")) { checkSparkAnswerAndOperator( - sql( - "SELECT /*+ BROADCAST(l) */ o.id, o.key " + - s"FROM objects o $joinType JOIN locations l ON $condition"), + sql("SELECT /*+ BROADCAST(l) */ o.id, o.key " + + s"FROM objects o $joinType JOIN locations l ON $condition"), ops) } checkSparkAnswerAndOperator( - sql( - "SELECT /*+ BROADCAST(o) */ o.id, l.table_name " + - s"FROM objects o RIGHT JOIN locations l ON $condition"), + sql("SELECT /*+ BROADCAST(o) */ o.id, l.table_name " + + s"FROM objects o RIGHT JOIN locations l ON $condition"), ops) } } From a757b48766d7208f1e86766b72715f3b66953947 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 19 Oct 2026 17:31:03 +0100 Subject: [PATCH 3/3] test: run the array join filter test with ANSI off, as on Spark 3.x Co-Authored-By: Claude Opus 5.5 (1M context) --- spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala | 2 ++ 1 file changed, 2 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 0d0b080b1e5..9b530f71e8d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -827,6 +827,8 @@ class CometJoinSuite extends CometTestBase { test("Broadcast HashJoin with a join filter over array elements of both sides") { withTempPath { dir => withSQLConf( + // `parts[i]` past the end is NULL, as in Spark 3.x defaults + SQLConf.ANSI_ENABLED.key -> "false", CometConf.COMET_BATCH_SIZE.key -> "100", SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {