jayzhan211 commented on code in PR #25496:
URL: https://github.com/apache/datafusion/pull/25496#discussion_r4061982350
##########
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:
Agreed, thanks for spelling it out. The keys are evaluated before they are
counted, so there is a window where the pool is behind RSS. Since there is no
generic upper bound for what an arbitrary expression or UDF allocates, and
`update_hash` and the ArrayMap null check have the same gap today, I'd rather
keep this PR to the retained buffers and leave reservation-backed expression
evaluation as a follow-up.
##########
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:
Good point, done in e10392689: `concat_test_batches` now has a nullable
`BinaryView` column as well, so the other concat tests cover it too, and
`concat_build_batches_view_data_not_reserved_twice` runs for both the
`Utf8View` and the `BinaryView` column.
--
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]
