kosiew commented on code in PR #25716:
URL: https://github.com/apache/datafusion/pull/25716#discussion_r4183118630
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -7607,6 +7729,181 @@ mod tests {
Ok(())
}
+ #[test]
+ fn concat_build_batches_deduplicates_view_buffers() -> Result<()> {
+ use arrow::array::StringViewBuilder;
+
+ let mut builder = StringViewBuilder::new();
+ builder.append_value("this is a long string that exceeds inline size
12");
+ builder.append_value("another long string that exceeds inline size
12");
+ let base_array: StringViewArray = builder.finish();
+ let schema =
+ Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8View,
true)]));
+
+ let batch1 = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(base_array.clone())],
+ )?;
+ let batch2 = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(base_array.clone())],
+ )?;
+ let batch3 = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(base_array.clone())],
+ )?;
+
+ // Before deduplication, concat_batches puts 3 duplicate buffer
references in data_buffers
+ let concatenated_raw =
+ concat_batches(&schema, &[batch1.clone(), batch2.clone(),
batch3.clone()])?;
+ let raw_view_arr = concatenated_raw
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .unwrap();
+ assert_eq!(raw_view_arr.data_buffers().len(), 3);
+
+ // After concat_build_batches, buffer references are deduplicated down
to 1
+ let metrics = BuildProbeJoinMetrics::new(0,
&ExecutionPlanMetricsSet::new());
+ let pool: Arc<dyn MemoryPool> =
Arc::new(UnboundedMemoryPool::default());
+ let batches = vec![batch1, batch2, batch3];
+ let (mut reservation, inputs_reserved) = reserve_inputs(&batches,
&pool)?;
+
+ let batch = concat_build_batches(
+ &schema,
+ batches,
+ false,
+ inputs_reserved,
+ &mut reservation,
+ &metrics,
+ )?;
+
+ let view_arr = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringViewArray>()
+ .unwrap();
+ assert_eq!(view_arr.data_buffers().len(), 1);
+ assert_eq!(batch.num_rows(), 6);
+ Ok(())
+ }
+
+ #[test]
+ fn concat_build_batches_deduplicates_binary_view_buffers() -> Result<()> {
+ use arrow::array::BinaryViewBuilder;
+
+ let mut builder = BinaryViewBuilder::new();
+ builder.append_value(b"this is a long binary that exceeds inline size
12");
+ builder.append_value(b"another long binary that exceeds inline size
12");
+ let base_array: BinaryViewArray = builder.finish();
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "b",
+ DataType::BinaryView,
+ true,
+ )]));
+
+ let batch1 = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(base_array.clone())],
+ )?;
+ let batch2 = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(base_array.clone())],
+ )?;
+ let batch3 = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(base_array.clone())],
+ )?;
+
+ let metrics = BuildProbeJoinMetrics::new(0,
&ExecutionPlanMetricsSet::new());
+ let pool: Arc<dyn MemoryPool> =
Arc::new(UnboundedMemoryPool::default());
+ let batches = vec![batch1, batch2, batch3];
+ let (mut reservation, inputs_reserved) = reserve_inputs(&batches,
&pool)?;
+
+ let batch = concat_build_batches(
+ &schema,
+ batches,
+ false,
+ inputs_reserved,
+ &mut reservation,
+ &metrics,
+ )?;
+
+ let view_arr = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<BinaryViewArray>()
+ .unwrap();
+ assert_eq!(view_arr.data_buffers().len(), 1);
+ assert_eq!(batch.num_rows(), 6);
+ Ok(())
+ }
+
+ #[test]
+ fn concat_build_batches_deduplicates_slices_regression() -> Result<()> {
+ use arrow::array::StringViewBuilder;
+
+ let mut builder = StringViewBuilder::new();
+ builder.append_value("this is a long string that exceeds inline size
12");
+ builder.append_value("another long string that exceeds inline size
12");
+ let base_array: StringViewArray = builder.finish();
+ let schema =
+ Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8View,
true)]));
+
+ // Create a shorter slice that starts at the same address
+ let short_slice = base_array.slice(0, 1);
Review Comment:
This test does not actually cover buffers with the same pointer but
different lengths. `GenericByteViewArray::slice` only slices the views, so
please construct two arrays using differently sized
`Buffer::slice_with_length(0, ...)` ranges with the shorter one first, then
assert that values from the longer range still survive concatenation.
--
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]
