Skip to content

Commit dd618b6

Browse files
committed
Use upstream arrow-rs record_batch! and create_array! macros
Removes DataFusion's custom `record_batch!` and `create_array!` macro implementations in favor of the upstream versions from arrow-rs added in apache/arrow-rs#6588. Changes: - Replace custom macro definitions with re-exports from arrow::array - Update syntax from vec![...] to array literal [...] across 67 usages - Add arrow_schema aliases in test modules for macro compatibility - Replace macro usage with manual RecordBatch construction where variables are used (macros only support literals) Closes #13037
1 parent 144f155 commit dd618b6

15 files changed

Lines changed: 166 additions & 207 deletions

File tree

datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,9 @@
1818
use std::sync::Arc;
1919

2020
use abi_stable::{export_root_module, prefix_type::PrefixTypeTrait};
21-
use arrow::array::RecordBatch;
21+
use arrow::array::{Float64Array, Int32Array, RecordBatch};
2222
use arrow::datatypes::{DataType, Field, Schema};
23-
use datafusion::{common::record_batch, datasource::MemTable};
23+
use datafusion::datasource::MemTable;
2424
use datafusion_ffi::table_provider::FFI_TableProvider;
2525
use ffi_module_interface::{TableProviderModule, TableProviderModuleRef};
2626

@@ -29,7 +29,19 @@ fn create_record_batch(start_value: i32, num_values: usize) -> RecordBatch {
2929
let a_vals: Vec<i32> = (start_value..end_value).collect();
3030
let b_vals: Vec<f64> = a_vals.iter().map(|v| *v as f64).collect();
3131

32-
record_batch!(("a", Int32, a_vals), ("b", Float64, b_vals)).unwrap()
32+
let schema = Arc::new(Schema::new(vec![
33+
Field::new("a", DataType::Int32, true),
34+
Field::new("b", DataType::Float64, true),
35+
]));
36+
37+
RecordBatch::try_new(
38+
schema,
39+
vec![
40+
Arc::new(Int32Array::from(a_vals)),
41+
Arc::new(Float64Array::from(b_vals)),
42+
],
43+
)
44+
.unwrap()
3345
}
3446

3547
/// Here we only wish to create a simple table provider as an example.

datafusion-examples/examples/memory_pool_execution_plan.rs

Lines changed: 34 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,8 @@
2424
//! - Handle memory pressure by spilling to disk
2525
//! - Release memory when done
2626
27-
use arrow::record_batch::RecordBatch;
27+
use arrow::array::{Int32Array, RecordBatch, StringArray};
28+
use arrow::datatypes::{DataType, Field, Schema};
2829
use arrow_schema::SchemaRef;
2930
use datafusion::common::record_batch;
3031
use datafusion::common::{exec_datafusion_err, internal_err};
@@ -59,27 +60,42 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
5960

6061
// Create smaller batches to ensure we get multiple RecordBatches from the scan
6162
// Make each batch smaller than the memory limit to force multiple batches
62-
let batch1 = record_batch!(
63-
("id", Int32, vec![1; 800]),
64-
("name", Utf8, vec!["Alice"; 800])
65-
)?;
63+
let schema = Arc::new(Schema::new(vec![
64+
Field::new("id", DataType::Int32, true),
65+
Field::new("name", DataType::Utf8, true),
66+
]));
6667

67-
let batch2 = record_batch!(
68-
("id", Int32, vec![2; 800]),
69-
("name", Utf8, vec!["Bob"; 800])
68+
let batch1 = RecordBatch::try_new(
69+
Arc::clone(&schema),
70+
vec![
71+
Arc::new(Int32Array::from(vec![1; 800])),
72+
Arc::new(StringArray::from(vec!["Alice"; 800])),
73+
],
7074
)?;
7175

72-
let batch3 = record_batch!(
73-
("id", Int32, vec![3; 800]),
74-
("name", Utf8, vec!["Charlie"; 800])
76+
let batch2 = RecordBatch::try_new(
77+
Arc::clone(&schema),
78+
vec![
79+
Arc::new(Int32Array::from(vec![2; 800])),
80+
Arc::new(StringArray::from(vec!["Bob"; 800])),
81+
],
7582
)?;
7683

77-
let batch4 = record_batch!(
78-
("id", Int32, vec![4; 800]),
79-
("name", Utf8, vec!["David"; 800])
84+
let batch3 = RecordBatch::try_new(
85+
Arc::clone(&schema),
86+
vec![
87+
Arc::new(Int32Array::from(vec![3; 800])),
88+
Arc::new(StringArray::from(vec!["Charlie"; 800])),
89+
],
8090
)?;
8191

82-
let schema = batch1.schema();
92+
let batch4 = RecordBatch::try_new(
93+
Arc::clone(&schema),
94+
vec![
95+
Arc::new(Int32Array::from(vec![4; 800])),
96+
Arc::new(StringArray::from(vec!["David"; 800])),
97+
],
98+
)?;
8399

84100
// Create a single MemTable with all batches in one partition to preserve order but ensure streaming
85101
let mem_table = Arc::new(MemTable::try_new(
@@ -286,10 +302,9 @@ impl ExecutionPlan for BufferingExecutionPlan {
286302

287303
// Since this is a simplified example, return an empty batch
288304
// In a real implementation, you would create a batch stream from the processed results
289-
record_batch!(("id", Int32, vec![5]), ("name", Utf8, vec!["Eve"]))
290-
.map_err(|e| {
291-
exec_datafusion_err!("Failed to create final RecordBatch: {e}")
292-
})
305+
record_batch!(("id", Int32, [5]), ("name", Utf8, ["Eve"])).map_err(|e| {
306+
exec_datafusion_err!("Failed to create final RecordBatch: {e}")
307+
})
293308
}),
294309
)))
295310
}

datafusion/common/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ pub mod utils;
6262

6363
/// Reexport arrow crate
6464
pub use arrow;
65+
/// Reexport arrow-rs macros for creating arrays and record batches
66+
pub use arrow::array::{create_array, record_batch};
6567
pub use column::Column;
6668
pub use dfschema::{
6769
qualified_name, DFSchema, DFSchemaRef, ExprSchema, SchemaExt, ToDFSchema,

0 commit comments

Comments
 (0)