jayzhan211 commented on code in PR #26793:
URL: https://github.com/apache/datafusion/pull/26793#discussion_r4430277660
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -3234,18 +3232,19 @@ fn build_in_memory(
}
// The maximum is known: avoid geometric growth and its excess
capacity.
let mut hashes_buffer = Vec::with_capacity(max_batch_rows);
- let mut offset = num_rows;
-
- // Update the hashmap starting from the last batch, so rows with
smaller
- // indices end up on top of the chains. Each batch is indexed at its
- // position in the forward concatenation below.
- for batch in batches.iter().rev() {
- offset -= batch.num_rows();
+ // `max_batch_rows` is 0 when the build side is empty
+ let chunk_rows = max_batch_rows.max(1);
+
+ // Update the hashmap from slices of at most `max_batch_rows` rows,
+ // starting from the last one, so rows with smaller indices end up on
+ // top of the chains. Each slice is indexed at its offset in `batch`.
+ for offset in (0..num_rows).step_by(chunk_rows).rev() {
+ let chunk = batch.slice(offset, chunk_rows.min(num_rows - offset));
Review Comment:
Agree with @Phoenix500526. I measured it with CollectLeft, 8192-row build
batches, each batch with its own `Dictionary(Int32, Utf8)` dictionary (release
build). At 256 batches the build goes from 8.9 ms to 98.7 ms with 1000-value
dictionaries, and from 85 ms to 1015 ms with 8192-value dictionaries, because
every slice shares the merged dictionary and `hash_dictionary` hashes all of it
on each call.
@dwsmith1983 a small fix that brings both cases back to 11.6 ms and 85.2 ms
is to make each slice at least as long as the longest key dictionary, and use
that length for the prepared scratch reservation too (`hash_dictionary` already
allocates that many hashes per call):
```diff
-use arrow::array::{Array, ArrayRef, BooleanBufferBuilder};
+use arrow::array::{Array, ArrayRef, AsArray, BooleanBufferBuilder};
@@
+ // Every slice shares the whole dictionary of a dictionary key,
which
+ // `create_hashes` hashes on each call: a slice at least as long as
the
+ // dictionary keeps that cost linear in the build rows.
+ let chunk_rows = left_values
+ .iter()
+ .filter_map(|values| values.as_any_dictionary_opt())
+ .map(|dict| dict.values().len())
+ .fold(max_batch_rows, usize::max);
+
let scratch_reservation = reservation.new_empty();
if prepared {
scratch_reservation.try_grow(prepared::prepared_scratch_bytes(
- max_batch_rows,
+ chunk_rows,
on_left.len(),
null_equality,
)?)?;
}
// The maximum is known: avoid geometric growth and its excess
capacity.
- let mut hashes_buffer = Vec::with_capacity(max_batch_rows);
- // `max_batch_rows` is 0 when the build side is empty
- let chunk_rows = max_batch_rows.max(1);
+ let mut hashes_buffer = Vec::with_capacity(chunk_rows);
+ // `chunk_rows` is 0 when the build side is empty
+ let chunk_rows = chunk_rows.max(1);
```
Could you also add a test for the branch where the dictionary is longer than
every batch? This one passes with the fix (it can't catch the slowdown, but it
covers order across per-batch dictionaries):
```rs
#[tokio::test]
async fn join_keeps_build_order_with_dictionary_keys() -> Result<()> {
let dict_type =
DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8));
let schema = Arc::new(Schema::new(vec![
Field::new("a1", DataType::Int32, false),
Field::new("b1", dict_type.clone(), true),
]));
let batch = |a: Vec<i32>,
keys: Vec<Option<i32>>,
values: Vec<&str>|
-> Result<RecordBatch> {
let b1 = DictionaryArray::<Int32Type>::try_new(
Int32Array::from(keys),
Arc::new(StringArray::from(values)),
)?;
Ok(RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(Int32Array::from(a)), Arc::new(b1)],
)?)
};
let batches = vec![
batch(vec![1, 2], vec![Some(0), None], vec!["x", "y", "z"])?,
batch(
vec![3, 4, 5],
vec![Some(2), Some(0), None],
vec!["x", "w", "y", "v"],
)?,
batch(vec![6, 7], vec![Some(1), Some(3)], vec!["u", "x", "t", "z"])?,
];
let right_schema = Arc::new(Schema::new(vec![
Field::new("a2", DataType::Int32, false),
Field::new("b2", dict_type, true),
]));
let right = RecordBatch::try_new(
Arc::clone(&right_schema),
vec![
Arc::new(Int32Array::from(vec![10, 30])),
Arc::new(DictionaryArray::<Int32Type>::try_new(
Int32Array::from(vec![0, 1]),
Arc::new(StringArray::from(vec!["x", "y"])),
)?),
],
)?;
let on = vec![(
Arc::new(Column::new_with_schema("b1", &schema)?) as _,
Arc::new(Column::new_with_schema("b2", &right_schema)?) as _,
)];
let left = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema),
None)?;
let right = TestMemoryExec::try_new_exec(&[vec![right]], right_schema,
None)?;
let (_, output, _) = join_collect(
left,
right,
on,
&JoinType::Inner,
NullEquality::NullEqualsNothing,
prepare_task_ctx(8192, false),
)
.await?;
assert_snapshot!(batches_to_string(&output), @r"
+----+----+----+----+
| a1 | b1 | a2 | b2 |
+----+----+----+----+
| 1 | x | 10 | x |
| 4 | x | 10 | x |
| 6 | x | 10 | x |
| 3 | y | 30 | y |
+----+----+----+----+
");
Ok(())
}
```
With this change, the loop comment ("slices of at most `max_batch_rows`
rows") and the `HashJoinExec` doc ("at most as long as the largest build-side
batch") would need the same wording update.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -107,8 +107,7 @@ const ARRAY_MAP_CREATED_COUNT_METRIC_NAME: &str =
"array_map_created_count";
/// (perfect hash join), returning the `(min, max)` key range to build it over.
///
/// On `Some`, the memory of the [`ArrayMap`] has already been added to
-/// `reservation`; the caller concatenates the build batches and creates the
map
-/// with [`ArrayMap::try_new`].
+/// `reservation`; the caller creates the map with [`ArrayMap::try_new`].
Review Comment:
Optional cleanup: `array_map_key_range` now always gets a single batch, and
under `NullEqualsNull` it evaluates the key expressions on it again. It could
take the key array that `build_in_memory` already evaluated, plus the row count:
```diff
- batches: &[RecordBatch],
- on_left: &[PhysicalExprRef],
+ key: &ArrayRef,
+ num_rows: usize,
@@
- if null_equality == NullEquality::NullEqualsNull {
- for batch in batches.iter() {
- let arrays = evaluate_expressions_to_arrays(on_left, batch)?;
- if arrays[0].null_count() > 0 {
- return Ok(None);
- }
- }
- }
+ if null_equality == NullEquality::NullEqualsNull && key.null_count() >
0 {
+ return Ok(None);
+ }
```
Callers would pass `&left_values[0], num_rows`. The candidate check that
uses `on_left` and `schema` can move to the call site, where `is_phj_candidate`
is already computed.
--
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]
