Skip to content

Commit 022c4a5

Browse files
committed
feat: enhance schema normalization in RecordBatch
- Updated the documentation for schema normalization to clarify the conditions under which the batch is returned unchanged. - Added a test to verify that schema normalization correctly strips metadata from the RecordBatch when the expected schema does not contain it. - Provided additional assertions to ensure zero-copy behavior in schema normalization. - Improved code comments for better clarity on the usage of the `normalize_batch_schema` function across various modules.
1 parent 9945e7e commit 022c4a5

3 files changed

Lines changed: 39 additions & 2 deletions

File tree

‎datafusion/physical-plan/src/common.rs‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ pub(crate) type SharedMemoryReservation = Arc<Mutex<MemoryReservation>>;
5151
///
5252
/// | Case | Result |
5353
/// |---|---|
54-
/// | `batch.schema() == expected_schema` | Returns the batch unchanged (fast path, zero copies). |
54+
/// | `batch` schema is structurally equal to `expected_schema` | Returns the batch unchanged (fast path, zero copies). |
5555
/// | Column count and data types match positionally | Returns a new `RecordBatch` that shares the original data buffers (zero-copy) but carries `expected_schema`. |
5656
/// | Column count or data types are incompatible | Returns `Err(internal_err!(...))`. |
5757
///
@@ -420,6 +420,9 @@ mod tests {
420420
let expected_schema =
421421
Arc::new(Schema::new(vec![Field::new("val", DataType::Int32, false)]));
422422

423+
// Capture the column Arc before the batch is moved.
424+
let batch_col_before = Arc::clone(batch.column(0));
425+
423426
let result = normalize_batch_schema(batch, &expected_schema)?;
424427

425428
assert_eq!(result.schema(), expected_schema);
@@ -431,6 +434,8 @@ mod tests {
431434
.downcast_ref::<Int32Array>()
432435
.unwrap();
433436
assert_eq!(col.values(), &[10, 20, 30]);
437+
// Verify the rename was truly zero-copy: same underlying column buffer.
438+
assert!(Arc::ptr_eq(result.column(0), &batch_col_before));
434439
Ok(())
435440
}
436441

@@ -452,6 +457,37 @@ mod tests {
452457
Ok(())
453458
}
454459

460+
#[test]
461+
fn test_normalize_batch_schema_strips_metadata() -> Result<()> {
462+
// When the batch schema carries Arrow schema-level metadata but
463+
// `expected_schema` does not, the helper must replace the schema
464+
// (conforming to the operator contract) and therefore strip the extra
465+
// metadata. This test confirms the behaviour is intentional.
466+
use std::collections::HashMap;
467+
let meta: HashMap<String, String> =
468+
[("key".to_string(), "value".to_string())].into_iter().collect();
469+
let src_schema = Arc::new(
470+
Schema::new(vec![Field::new("child_col", DataType::Int32, false)])
471+
.with_metadata(meta),
472+
);
473+
let expected_schema = Arc::new(Schema::new(vec![Field::new(
474+
"col",
475+
DataType::Int32,
476+
false,
477+
)]));
478+
let batch = RecordBatch::try_new(
479+
Arc::clone(&src_schema),
480+
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
481+
)?;
482+
483+
let result = normalize_batch_schema(batch, &expected_schema)?;
484+
485+
// Schema metadata from the source batch is gone — replaced by expected_schema.
486+
assert_eq!(result.schema(), expected_schema);
487+
assert!(result.schema_ref().metadata().is_empty());
488+
Ok(())
489+
}
490+
455491
#[test]
456492
fn test_normalize_batch_schema_error_column_count_mismatch() -> Result<()> {
457493
// Batch has 1 column; expected schema has 2 — must error.

‎datafusion/physical-plan/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ pub use crate::stream::EmptyRecordBatchStream;
5353
pub use crate::topk::TopK;
5454
pub use crate::visitor::{ExecutionPlanVisitor, accept, visit_execution_plan};
5555
pub use crate::work_table::WorkTable;
56+
pub use crate::common::normalize_batch_schema;
5657
pub use spill::spill_manager::SpillManager;
5758

5859
mod ordering;

‎datafusion/physical-plan/src/recursive_query.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -320,7 +320,7 @@ impl RecursiveQueryStream {
320320
let baseline_metrics = self.baseline_metrics.clone();
321321

322322
// Rebind to the declared output schema via the shared helper.
323-
// See `normalize_batch_schema` in `common` for full semantics.
323+
// See [`crate::common::normalize_batch_schema`] for full semantics.
324324
batch = normalize_batch_schema(batch, &self.schema)?;
325325

326326
if let Some(deduplicator) = &mut self.distinct_deduplicator {

0 commit comments

Comments
 (0)