kosiew commented on code in PR #25496:
URL: https://github.com/apache/datafusion/pull/25496#discussion_r4061262731
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2932,48 +3012,75 @@ async fn collect_left_input(
config.execution.perfect_hash_join_min_key_density,
null_equality,
)? {
- array_map_created_count.add(1);
- metrics.build_mem_used.add(array_map.size());
+ let batch = concat_build_batches(
+ &schema,
+ batches,
+ false,
+ inputs_reserved,
+ &mut reservation,
+ &metrics,
+ )?;
+ let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
Review Comment:
One thing worth noting is that `evaluate_expressions_to_arrays` can
materialize computed build keys before their retained buffers are counted and
reserved below. That leaves a temporary gap between pool accounting and actual
RSS.
I don't think this needs to block this PR. This PR addresses the unaccounted
retained concat and computed-key buffers, and the `RecordBatchMemoryCounter`
accounting below handles the latter once evaluation completes. Arbitrary
physical expressions and UDFs also don't have a reliable generic allocation
upper bound, and similar pre-evaluation temporary allocations already exist in
`update_hash` and the ArrayMap null check.
I think reservation-backed expression evaluation would be better handled as
a broader follow-up rather than expanding the scope here.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -7139,6 +7249,328 @@ mod tests {
Ok(())
}
+ /// Build side batches of `num_batches` x `num_rows` rows with distinct
+ /// buffers: an Int32 key, a Utf8 and a Utf8View payload.
+ fn concat_test_batches(num_batches: usize, num_rows: usize) ->
Vec<RecordBatch> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int32, false),
+ Field::new("s", DataType::Utf8, true),
+ Field::new("v", DataType::Utf8View, true),
+ ]));
+ (0..num_batches)
+ .map(|b| {
+ let start = (b * num_rows) as i32;
+ let keys = Int32Array::from_iter_values(start..start +
num_rows as i32);
+ let strings = (0..num_rows)
+ .map(|i| (i % 7 != 0).then(|| format!("string-{b}-{i}")))
+ .collect::<StringArray>();
+ let views = (0..num_rows)
+ .map(|i| {
+ (i % 5 != 0).then(|| format!("a long string view value
{b}-{i}"))
+ })
+ .collect::<StringViewArray>();
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(keys), Arc::new(strings), Arc::new(views)],
+ )
+ .unwrap()
+ })
+ .collect()
+ }
+
+ /// Reserves `batches` the way `collect_left_input` does.
+ fn reserve_inputs(
+ batches: &[RecordBatch],
+ pool: &Arc<dyn MemoryPool>,
+ ) -> Result<(MemoryReservation, usize)> {
+ let reservation = MemoryConsumer::new("HashJoinInput").register(pool);
+ let mut counter = RecordBatchMemoryCounter::new();
+ for batch in batches {
+ reservation.try_grow(counter.count_batch(batch))?;
+ }
+ Ok((reservation, counter.memory_usage()))
+ }
+
+ #[test]
+ fn concat_build_batches_matches_concat_batches() -> Result<()> {
+ let batches = concat_test_batches(4, 100);
+ let schema = batches[0].schema();
+ let metrics = BuildProbeJoinMetrics::new(0,
&ExecutionPlanMetricsSet::new());
+
+ for reverse in [false, true] {
+ let pool: Arc<dyn MemoryPool> =
Arc::new(UnboundedMemoryPool::default());
+ let (mut reservation, inputs_reserved) = reserve_inputs(&batches,
&pool)?;
+ let expected = if reverse {
+ concat_batches(&schema, batches.iter().rev())?
+ } else {
+ concat_batches(&schema, batches.iter())?
+ };
+ let batch = concat_build_batches(
+ &schema,
+ batches.clone(),
+ reverse,
+ inputs_reserved,
+ &mut reservation,
+ &metrics,
+ )?;
+ assert_eq!(batch, expected);
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn concat_build_batches_reserves_copy() -> Result<()> {
+ let batches = concat_test_batches(4, 1000);
+ let schema = batches[0].schema();
+ let metrics = BuildProbeJoinMetrics::new(0,
&ExecutionPlanMetricsSet::new());
+ let inputs: usize =
batches.iter().map(get_record_batch_memory_size).sum();
+
+ // The inputs fit, but not together with their concatenated copy
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(inputs
* 5 / 4));
+ let (mut reservation, inputs_reserved) = reserve_inputs(&batches,
&pool)?;
+ assert_eq!(inputs_reserved, inputs);
+ let err = concat_build_batches(
+ &schema,
+ batches.clone(),
+ false,
+ inputs_reserved,
+ &mut reservation,
+ &metrics,
+ )
+ .unwrap_err();
+ assert_contains!(err.to_string(), "Resources exhausted");
+ drop(reservation);
+
+ // With room for the copy, the reservation ends up at what is retained
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(inputs
* 2));
+ let (mut reservation, inputs_reserved) = reserve_inputs(&batches,
&pool)?;
+ let batch = concat_build_batches(
+ &schema,
+ batches,
+ false,
+ inputs_reserved,
+ &mut reservation,
+ &metrics,
+ )?;
+ assert_eq!(reservation.size(), get_record_batch_memory_size(&batch));
+ assert_eq!(pool.reserved(), reservation.size());
+ Ok(())
+ }
+
+ /// Concatenating a single batch is zero-copy, so nothing more is reserved.
+ #[test]
+ fn concat_build_batches_single_batch_not_reserved_twice() -> Result<()> {
+ let batches = concat_test_batches(1, 1000);
+ let schema = batches[0].schema();
+ let metrics = BuildProbeJoinMetrics::new(0,
&ExecutionPlanMetricsSet::new());
+ let inputs = get_record_batch_memory_size(&batches[0]);
+
+ let pool: Arc<dyn MemoryPool> =
Arc::new(GreedyMemoryPool::new(inputs));
+ let (mut reservation, inputs_reserved) = reserve_inputs(&batches,
&pool)?;
+ let batch = concat_build_batches(
+ &schema,
+ batches,
+ true,
+ inputs_reserved,
+ &mut reservation,
+ &metrics,
+ )?;
+ assert_eq!(batch.num_rows(), 1000);
+ assert_eq!(reservation.size(), inputs);
+ Ok(())
+ }
+
+ /// Only the views of a view array are copied, its data buffers stay
shared.
+ #[test]
+ fn concat_build_batches_view_data_not_reserved_twice() -> Result<()> {
Review Comment:
Could be nice to add `BinaryView` coverage alongside `Utf8View` here.
`estimate_concat_allocation` handles both variants in the same match arm, so
this would give us some symmetry in the test coverage. Definitely not something
I'd block the PR on.
--
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]
