Skip to content

refactor(asof_join): enhance match value handling and comparator caching #24519

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jayzhan211 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor(asof_join): enhance match value handling and comparator caching #24519

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
170 changes: 137 additions & 33 deletions datafusion/physical-plan/src/joins/asof_join.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,16 +72,19 @@ use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, new_null_array};
use arrow::array::{
Array, ArrayRef, DynComparator, RecordBatch, RecordBatchOptions, make_comparator,
new_null_array,
};
use arrow::buffer::NullBuffer;
use arrow::compute::{SortOptions, interleave};
use arrow::datatypes::{Schema, SchemaRef};
use datafusion_common::stats::Precision;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::utils::memory::RecordBatchMemoryCounter;
use datafusion_common::utils::normalize_float_zero_scalar;
use datafusion_common::utils::normalize_float_zero;
use datafusion_common::{
ColumnStatistics, JoinSide, JoinType, NullEquality, Result, ScalarValue, Statistics,
ColumnStatistics, JoinSide, JoinType, NullEquality, Result, Statistics,
assert_eq_or_internal_err, internal_err, plan_err, project_schema,
};
use datafusion_execution::TaskContext;
Expand Down Expand Up @@ -626,9 +629,12 @@ struct InputCursor {
key_arrays: Arc<[ArrayRef]>,
/// Rows whose equality keys are all non-NULL.
key_validity: Option<NullBuffer>,
/// Evaluated match values for `batch`.
/// Evaluated and normalized match values for `batch`.
match_array: Option<ArrayRef>,
/// Monotonic identity of the current key arrays.
/// Logical NULLs cached separately so row checks avoid constructing
/// `ScalarValue`s and still handle nested representations such as dictionaries.
match_validity: Option<NullBuffer>,
/// Monotonic identity of the current arrays.
key_batch_id: usize,
/// Current row within `batch`.
row: usize,
Expand All @@ -650,6 +656,7 @@ impl InputCursor {
key_arrays: Arc::from([]),
key_validity: None,
match_array: None,
match_validity: None,
key_batch_id: 0,
row: 0,
eof: false,
Expand All @@ -671,6 +678,7 @@ impl InputCursor {
self.key_arrays = Arc::from([]);
self.key_validity = None;
self.match_array = None;
self.match_validity = None;
self.row = 0;
if self.eof {
return Poll::Ready(Ok(false));
Expand All @@ -692,11 +700,15 @@ impl InputCursor {
self.key_validity =
matchable_join_keys(&key_arrays, NullEquality::NullEqualsNothing);
self.key_arrays = key_arrays.into();
self.match_array = Some(
self.match_expr
.evaluate(&batch)?
.into_array(batch.num_rows())?,
);
let match_array = self
.match_expr
.evaluate(&batch)?
.into_array(batch.num_rows())?;
self.match_validity = match_array.logical_nulls();
// Arrow's comparator distinguishes -0.0 from +0.0, while SQL does
// not. Normalizing once here is cheaper than fixing every row-level
// comparison and lets the cached comparator use the arrays directly.
self.match_array = Some(normalize_float_zero(&match_array));
self.key_batch_id += 1;
self.batch = Some(batch);
}
Expand All @@ -708,11 +720,10 @@ impl InputCursor {
.is_some_and(|validity| validity.is_null(self.row))
}

fn match_value(&self) -> Result<ScalarValue> {
let array = self.match_array.as_ref().ok_or_else(|| {
datafusion_common::internal_datafusion_err!("ASOF match array is missing")
})?;
ScalarValue::try_from_array(array, self.row).map(normalize_float_zero_scalar)
fn match_is_null(&self) -> bool {
self.match_validity
.as_ref()
.is_some_and(|validity| validity.is_null(self.row))
}

fn batch_row(&self) -> Result<(Arc<RecordBatch>, usize)> {
Expand Down Expand Up @@ -847,6 +858,11 @@ impl PendingRows {
/// candidate advances from `(A, 2)` to `(A, 6)` without rewinding the right
/// cursor. Cursors and the candidate survive input batch changes and output
/// flushes; a change of equality group clears the candidate before reuse.
///
/// The hot path avoids materializing scalar values: expressions are evaluated
/// once per batch, comparators are cached by batch identity, and both cursors
/// only move forward. Input batches and evaluated arrays remain shared through
/// `Arc`; output can therefore use a zero-copy slice when its rows are contiguous.
struct AsOfJoinStream {
/// Output schema used when pending row references are materialized.
schema: SchemaRef,
Expand All @@ -870,6 +886,10 @@ struct AsOfJoinStream {
input_group_comparator: Option<(usize, usize, JoinKeyComparator)>,
/// Cached comparator for the candidate and current left batches.
candidate_group_comparator: Option<(usize, usize, JoinKeyComparator)>,
/// Cached match-key comparator for the current right and left batches.
/// Building it performs type dispatch, so doing so once per batch pair avoids
/// repeating that work for every candidate comparison.
input_match_comparator: Option<(usize, usize, DynComparator)>,
/// Left row references accumulated for the next output batch.
pending_left: PendingRows,
/// Matched right row references, aligned with `pending_left`.
Expand Down Expand Up @@ -914,6 +934,7 @@ impl AsOfJoinStream {
group_sort_options,
input_group_comparator: None,
candidate_group_comparator: None,
input_match_comparator: None,
batch_size: batch_size.max(1),
metrics,
}
Expand Down Expand Up @@ -982,6 +1003,50 @@ impl AsOfJoinStream {
Ok(comparator.compare(candidate.row, self.left.row) != Ordering::Equal)
}

/// Compares the current right match value with the current left match value.
///
/// This always returns natural ascending order (`right.cmp(left)`), regardless
/// of scan direction. [`is_eligible`] interprets that order for the four ASOF
/// operators, keeping direction-specific logic out of the comparator cache.
fn compare_input_matches(&mut self) -> Result<Ordering> {
let _timer = self.metrics.baseline.elapsed_compute().timer();
let right_batch_id = self.right.key_batch_id;
let left_batch_id = self.left.key_batch_id;
if self
.input_match_comparator
.as_ref()
.is_none_or(|(right, left, _)| {
*right != right_batch_id || *left != left_batch_id
})
{
let right = self.right.match_array.as_ref().ok_or_else(|| {
datafusion_common::internal_datafusion_err!(
"ASOF right match array is missing"
)
})?;
let left = self.left.match_array.as_ref().ok_or_else(|| {
datafusion_common::internal_datafusion_err!(
"ASOF left match array is missing"
)
})?;
let comparator = make_comparator(
right.as_ref(),
left.as_ref(),
SortOptions {
descending: false,
nulls_first: true,
},
)?;
self.input_match_comparator =
Some((right_batch_id, left_batch_id, comparator));
}
let (_, _, comparator) = self
.input_match_comparator
.as_ref()
.expect("ASOF input match comparator must be initialized");
Ok(comparator(self.right.row, self.left.row))
}

/// Produces the next output batch without resetting the merge state.
///
/// Each left row first validates its equality group, then advances the right
Expand Down Expand Up @@ -1021,11 +1086,7 @@ impl AsOfJoinStream {
return Poll::Ready(None);
}

let left_match = {
let _timer = self.metrics.baseline.elapsed_compute().timer();
self.left.match_value()?
};
if left_match.is_null() || self.left.group_has_null() {
if self.left.match_is_null() || self.left.group_has_null() {
self.candidate = None;
self.candidate_group_comparator = None;
self.push_current_left(None)?;
Expand Down Expand Up @@ -1056,13 +1117,11 @@ impl AsOfJoinStream {
Ordering::Greater => break,
Ordering::Equal => {}
}
let _timer = self.metrics.baseline.elapsed_compute().timer();
let right_match = self.right.match_value()?;
if right_match.is_null() {
if self.right.match_is_null() {
self.right.advance();
continue;
}
if !is_eligible(self.op, &left_match, &right_match)? {
if !is_eligible(self.op, self.compare_input_matches()?) {
break;
}
let (batch, row) = self.right.batch_row()?;
Expand Down Expand Up @@ -1230,15 +1289,14 @@ fn validate_expr_side(expr: &PhysicalExprRef, schema: &Schema, name: &str) -> Re
Ok(())
}

fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result<bool> {
let ordering = right.try_cmp(left)?;
Ok(match op {
Operator::Gt => ordering == Ordering::Less,
Operator::GtEq => ordering != Ordering::Greater,
Operator::Lt => ordering == Ordering::Greater,
Operator::LtEq => ordering != Ordering::Less,
fn is_eligible(op: Operator, right_vs_left: Ordering) -> bool {
match op {
Operator::Gt => right_vs_left == Ordering::Less,
Operator::GtEq => right_vs_left != Ordering::Greater,
Operator::Lt => right_vs_left == Ordering::Greater,
Operator::LtEq => right_vs_left != Ordering::Less,
_ => unreachable!("ASOF match operator is validated by try_new"),
})
}
}

#[cfg(test)]
Expand All @@ -1249,8 +1307,9 @@ mod tests {
use super::*;
use crate::collect;
use crate::test::TestMemoryExec;
use arrow::array::{Int32Array, Int64Array, StringArray};
use arrow::array::{Float64Array, Int32Array, Int64Array, StringArray};
use arrow::datatypes::{DataType, Field};
use datafusion_common::ScalarValue;
use datafusion_common::test_util::batches_to_sort_string;
use datafusion_execution::config::SessionConfig;
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
Expand Down Expand Up @@ -1525,6 +1584,51 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn floating_match_treats_signed_zero_as_equal() -> Result<()> {
let left_batch = RecordBatch::try_from_iter(vec![
("ts", Arc::new(Float64Array::from(vec![0.0])) as ArrayRef),
("id", Arc::new(Int32Array::from(vec![1])) as ArrayRef),
])?;
let left_schema = left_batch.schema();
let left = TestMemoryExec::try_new_exec(
&[vec![left_batch]],
Arc::clone(&left_schema),
None,
)?;

let right_batch = RecordBatch::try_from_iter(vec![
("ts", Arc::new(Float64Array::from(vec![-0.0])) as ArrayRef),
("price", Arc::new(Int32Array::from(vec![42])) as ArrayRef),
])?;
let right_schema = right_batch.schema();
let right = TestMemoryExec::try_new_exec(
&[vec![right_batch]],
Arc::clone(&right_schema),
None,
)?;

let exec = Arc::new(AsOfJoinExec::try_new(
left,
right,
vec![],
AsOfMatchExpr::new(
Arc::new(PhysicalColumn::new("ts", 0)),
Operator::Gt,
Arc::new(PhysicalColumn::new("ts", 0)),
),
Some(vec![1, 3]),
)?);
let batches = collect(exec, Arc::new(TaskContext::default())).await?;
let price = batches[0]
.column(1)
.as_any()
.downcast_ref::<Int32Array>()
.expect("price must be Int32");
assert!(price.is_null(0));
Ok(())
}

#[tokio::test]
async fn complex_equality_and_match_expressions() -> Result<()> {
let left_batch = RecordBatch::try_from_iter(vec![
Expand Down