comphead commented on code in PR #27278:
URL: https://github.com/apache/datafusion/pull/27278#discussion_r4540672650
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -360,6 +482,24 @@ pub(super) struct MaterializingSortMergeJoinStream {
/// since NULL-padded rows reach filter evaluation. `None` unless
/// `deferred_filtering`.
pub filter_schema: Option<SchemaRef>,
+ /// Whether this is an Inner join with a filter. Such a join drops every
pair
+ /// the filter rejects, so the filter runs before any output column is
+ /// materialized: on a long run of buffered rows as it is paired with a
+ /// streamed row (see [`Self::pair_filtered_run`]), and at freeze for the
+ /// remaining pairs (see [`Self::freeze_inner_filtered`]). Outer joins keep
+ /// a rejected pair for the deferred pipeline, which may null-join its row.
+ /// Only an Inner join filters early, so its streamed side is the left
input.
+ pub early_filtering: bool,
+ /// The indices of every buffered column, for materializing all of them
+ pub buffered_columns: Vec<usize>,
+ /// The buffered columns that the join filter reads, in the filter's order.
+ /// Empty unless `early_filtering`.
+ pub buffered_filter_columns: Vec<usize>,
+ /// How many pairs the Inner join's filter has evaluated at freeze so far
+ /// (see [`Self::filter_first_pays`])
+ pub freeze_filter_evaluated: usize,
Review Comment:
Done in 8012102c4: both counters are `u64`, and the per-freeze increments
are widened.
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -1771,33 +2008,215 @@ impl MaterializingSortMergeJoinStream {
Ok(())
}
- /// Materializes right-side columns across all matched chunks.
+ /// Freezes the matched pairs of an Inner join with a filter. The filter
runs
+ /// on just the columns it reads, for the pairs of the chunks that
+ /// [`Self::pair_filtered_run`] did not filter already, and only the pairs
that
+ /// pass have their output columns materialized. A selective filter over a
+ /// many-to-many join rejects most pairs, so materializing every column of
+ /// every pair before filtering would spend most of the freeze on rows that
+ /// are then dropped. When no chunk was filtered already and earlier
freezes
+ /// kept enough of their pairs, [`Self::filter_first_pays`] says the
opposite,
+ /// and the batch is materialized whole and then filtered.
+ fn freeze_inner_filtered(
+ &mut self,
+ matched_chunks: &[MatchedChunk],
+ total_matched_rows: usize,
+ ) -> Result<()> {
+ let Some(filter) = &self.filter else {
+ return internal_err!("early filtering requires a join filter");
+ };
+ debug_assert_eq!(
+ self.join_type,
+ JoinType::Inner,
+ "only an Inner join filters early, so the streamed side is the
left one"
+ );
+
+ // Evaluate the filter over the pairs of the chunks not filtered yet.
+ let unfiltered: Vec<&MatchedChunk> =
+ matched_chunks.iter().filter(|c| !c.filtered).collect();
+ if unfiltered.len() == matched_chunks.len() &&
!self.filter_first_pays(filter) {
+ return self.freeze_inner_filtered_whole(matched_chunks,
total_matched_rows);
+ }
+ let passed = if unfiltered.is_empty() {
+ None
+ } else {
+ let num_rows = unfiltered.iter().map(|c| c.len()).sum();
+ let streamed_indices =
+ concat_indices(unfiltered.iter().map(|c|
&c.streamed_indices))?;
+ let mut buffered_arrays = self
+ .materialize_right_columns(
+ &unfiltered,
+ num_rows,
+ &self.buffered_filter_columns,
+ )?
+ .into_iter();
+ let columns = filter
+ .column_indices()
+ .iter()
+ .map(|column_index| match column_index.side {
+ JoinSide::Left => take_column(
+ self.streamed_batch.batch.column(column_index.index),
+ &streamed_indices,
+ ),
+ JoinSide::Right => buffered_arrays.next().ok_or_else(|| {
+ internal_datafusion_err!("missing buffered filter
column")
+ }),
+ JoinSide::None => {
+ internal_err!("Unexpected JoinSide::None in filter")
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+ self.freeze_filter_evaluated += num_rows;
+ Some(evaluate_join_filter(
+ filter,
+ filter.schema(),
+ columns,
+ num_rows,
+ )?)
+ };
+
+ // Keep the pairs that passed, in order.
+ let mut kept: Vec<Cow<MatchedChunk>> =
Vec::with_capacity(matched_chunks.len());
Review Comment:
Done in 8012102c4. `Cow`, `MatchedChunk::filter` and the `Borrow` bound are
gone, and `materialize_right_columns` takes `&[MatchedChunk]` again. With
andygrove's change, the kept chunks are slices of the filtered index arrays, so
no clone is needed.
##########
datafusion/physical-plan/src/joins/sort_merge_join/tests.rs:
##########
@@ -8395,3 +8399,322 @@ async fn
left_join_with_partial_filter_preserves_streamed_order() -> Result<()>
);
Ok(())
}
+
+// ============================================================================
+// Early filtering of Inner joins
+// ============================================================================
+
+/// Rows for the early-filtering tests, sorted by the key column `names[0]`:
key
+/// `k` of `0..keys` has `rows_per_key(k)` rows. `names[1]` is the row number
and
+/// `names[2]` a pseudo-random value in `0..1000` that is NULL in every 37th
row.
+/// The rows are split into batches of `rows_per_batch`.
+fn early_filtering_batches(
+ names: [&str; 3],
+ keys: i32,
+ rows_per_key: impl Fn(i32) -> usize,
+ seed: u64,
+ rows_per_batch: usize,
+) -> (Vec<RecordBatch>, SchemaRef) {
+ let mut key_values = vec![];
+ let mut row_numbers = vec![];
+ let mut values = vec![];
+ let mut state = seed;
+ for key in 0..keys {
+ for _ in 0..rows_per_key(key) {
+ state = state
+ .wrapping_mul(6364136223846793005)
+ .wrapping_add(1442695040888963407);
+ let row = key_values.len() as i32;
+ key_values.push(key);
+ row_numbers.push(row);
+ values.push((row % 37 != 0).then_some(((state >> 33) % 1000) as
i32));
+ }
+ }
+ let schema = Arc::new(Schema::new(vec![
+ Field::new(names[0], DataType::Int32, false),
+ Field::new(names[1], DataType::Int32, false),
+ Field::new(names[2], DataType::Int32, true),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(key_values)),
+ Arc::new(Int32Array::from(row_numbers)),
+ Arc::new(Int32Array::from(values)),
+ ],
+ )
+ .unwrap();
+ let batches = (0..batch.num_rows())
+ .step_by(rows_per_batch)
+ .map(|offset| batch.slice(offset, rows_per_batch.min(batch.num_rows()
- offset)))
+ .collect();
+ (batches, schema)
+}
+
+/// Inputs for the early-filtering tests. Every third key has a group of 600
+/// buffered rows, long enough to be filtered as it is paired, and the 250-row
+/// buffered batches split each such group. The other keys have groups of one
or
+/// two buffered rows, whose pairs are filtered at freeze. Keys repeat up to
four
+/// times on the streamed side.
+fn early_filtering_inputs()
+-> ((Vec<RecordBatch>, SchemaRef), (Vec<RecordBatch>, SchemaRef)) {
+ let left =
+ early_filtering_batches(["a1", "b1", "c1"], 40, |k| 1 + (k % 4) as
usize, 1, 7);
+ let right = early_filtering_batches(
+ ["a2", "b2", "c2"],
+ 40,
+ |k| if k % 3 == 0 { 600 } else { (k % 3) as usize },
+ 2,
+ 250,
+ );
+ (left, right)
+}
+
+/// The filter `c2 > c1 AND b1 % 2 = b2 % 2`, whose intermediate columns are in
+/// the order `c2, c1, b1, b2`, so that they alternate between the join sides.
+fn interleaved_early_filter(left_schema: &Schema, right_schema: &Schema) ->
JoinFilter {
+ let parity = |name: &str, index: usize| -> PhysicalExprRef {
+ Arc::new(BinaryExpr::new(
+ Arc::new(Column::new(name, index)),
+ Operator::Modulo,
+ Arc::new(Literal::new(ScalarValue::Int32(Some(2)))),
+ ))
+ };
+ let expression = Arc::new(BinaryExpr::new(
+ Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("c2", 0)),
+ Operator::Gt,
+ Arc::new(Column::new("c1", 1)),
+ )),
+ Operator::And,
+ Arc::new(BinaryExpr::new(
+ parity("b1", 2),
+ Operator::Eq,
+ parity("b2", 3),
+ )),
+ ));
+ let column_indices = vec![
+ ColumnIndex {
+ index: 2,
+ side: JoinSide::Right,
+ },
+ ColumnIndex {
+ index: 2,
+ side: JoinSide::Left,
+ },
+ ColumnIndex {
+ index: 1,
+ side: JoinSide::Left,
+ },
+ ColumnIndex {
+ index: 1,
+ side: JoinSide::Right,
+ },
+ ];
+ let schema = Schema::new(vec![
+ right_schema.field_with_name("c2").unwrap().clone(),
+ left_schema.field_with_name("c1").unwrap().clone(),
+ left_schema.field_with_name("b1").unwrap().clone(),
+ right_schema.field_with_name("b2").unwrap().clone(),
+ ]);
+ JoinFilter::new(expression, column_indices, Arc::new(schema))
+}
+
+/// The filter `b2 <> b1 OR c2 < c1`, which nearly every pair passes, with its
+/// intermediate columns in the order `b2, b1, c2, c1`. It reads four of the
six
+/// output columns, so after the first freeze the join materializes whole
batches
+/// and filters them (see `filter_first_pays`).
+fn high_pass_early_filter(left_schema: &Schema, right_schema: &Schema) ->
JoinFilter {
+ let expression = Arc::new(BinaryExpr::new(
+ Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("b2", 0)),
+ Operator::NotEq,
+ Arc::new(Column::new("b1", 1)),
+ )),
+ Operator::Or,
+ Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("c2", 2)),
+ Operator::Lt,
+ Arc::new(Column::new("c1", 3)),
+ )),
+ ));
+ let column_indices = vec![
+ ColumnIndex {
+ index: 1,
+ side: JoinSide::Right,
+ },
+ ColumnIndex {
+ index: 1,
+ side: JoinSide::Left,
+ },
+ ColumnIndex {
+ index: 2,
+ side: JoinSide::Right,
+ },
+ ColumnIndex {
+ index: 2,
+ side: JoinSide::Left,
+ },
+ ];
+ let schema = Schema::new(vec![
+ right_schema.field_with_name("b2").unwrap().clone(),
+ left_schema.field_with_name("b1").unwrap().clone(),
+ right_schema.field_with_name("c2").unwrap().clone(),
+ left_schema.field_with_name("c1").unwrap().clone(),
+ ]);
+ JoinFilter::new(expression, column_indices, Arc::new(schema))
+}
+
+/// A filter that reads no columns and always evaluates to `value`.
+fn column_free_early_filter(value: bool) -> JoinFilter {
+ JoinFilter::new(
+ Arc::new(Literal::new(ScalarValue::Boolean(Some(value)))),
+ vec![],
+ Arc::new(Schema::empty()),
+ )
+}
+
+fn early_filtering_on(left_schema: &Schema, right_schema: &Schema) ->
Result<JoinOn> {
+ Ok(vec![(
+ Arc::new(Column::new_with_schema("a1", left_schema)?) as _,
+ Arc::new(Column::new_with_schema("a2", right_schema)?) as _,
+ )])
+}
+
+/// Asserts that the key column of `batches` is ascending, the order of the
+/// streamed side that the join maintains.
+fn assert_streamed_order(batches: &[RecordBatch]) {
+ let keys: Vec<i32> = batches
+ .iter()
+ .flat_map(|b| {
+ b.column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect();
+ assert!(
+ keys.is_sorted(),
Review Comment:
Done in 8012102c4: it asserts that `(b1, b2)` ascends.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
