This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 00ad7fca2f chore: extend record_batch macro to support variables and
expressions (#9522)
00ad7fca2f is described below
commit 00ad7fca2fc5e09c0da5f56f87edc3a454eec576
Author: Burak Şen <[email protected]>
AuthorDate: Wed Mar 18 16:25:34 2026 +0300
chore: extend record_batch macro to support variables and expressions
(#9522)
# Which issue does this PR close?
- Closes #9245.
# Rationale for this change
Currently record_batch! macro supports only literal values. In
datafusion repository there is also a record_batch! macro that supports
this.
https://github.com/apache/datafusion/issues/13037 can be closed after
Datafusion repository upgrades version
# What changes are included in this PR?
Extend record_batch! macro to support datafusion equivalent added in:
# Are these changes tested?
I've actually ported datafusion logic to here. I was not sure if it
makes sense to add unit tests for this macro but I can if requested
# Are there any user-facing changes?
No breaking changes to downstream since this only extends macro
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
arrow-array/src/record_batch.rs | 65 ++++++++++++++++++++++++++++++++++++-----
1 file changed, 58 insertions(+), 7 deletions(-)
diff --git a/arrow-array/src/record_batch.rs b/arrow-array/src/record_batch.rs
index cfec969165..780e14fd4f 100644
--- a/arrow-array/src/record_batch.rs
+++ b/arrow-array/src/record_batch.rs
@@ -135,6 +135,18 @@ macro_rules! create_array {
($ty: tt, [$($values: expr),*]) => {
std::sync::Arc::new(<$crate::create_array!(@from
$ty)>::from(vec![$($values),*]))
};
+
+ (Binary, $values: expr) => {
+ std::sync::Arc::new($crate::BinaryArray::from_vec($values))
+ };
+
+ (LargeBinary, $values: expr) => {
+ std::sync::Arc::new($crate::LargeBinaryArray::from_vec($values))
+ };
+
+ ($ty: tt, $values: expr) => {
+ std::sync::Arc::new(<$crate::create_array!(@from $ty)>::from($values))
+ };
}
/// Creates a record batch from literal slice of values, suitable for rapid
@@ -152,10 +164,22 @@ macro_rules! create_array {
/// ("c", Utf8, ["alpha", "beta", "gamma"])
/// );
/// ```
+///
+/// Variables and expressions are also supported:
+///
+/// ```rust
+/// use arrow_array::record_batch;
+///
+/// let values = vec![1, 2, 3];
+/// let batch = record_batch!(
+/// ("a", Int32, values),
+/// ("b", Float64, vec![Some(4.0), None, Some(5.0)])
+/// );
+/// ```
/// Due to limitation of [`create_array!`] macro, support for limited data
types is available.
#[macro_export]
macro_rules! record_batch {
- ($(($name: expr, $type: ident, [$($values: expr),*])),*) => {
+ ($(($name: expr, $type: ident, $($values: tt)+)),*) => {
{
let schema = std::sync::Arc::new(arrow_schema::Schema::new(vec![
$(
@@ -163,16 +187,14 @@ macro_rules! record_batch {
)*
]));
- let batch = $crate::RecordBatch::try_new(
+ $crate::RecordBatch::try_new(
schema,
vec![$(
- $crate::create_array!($type, [$($values),*]),
+ $crate::create_array!($type, $($values)+),
)*]
- );
-
- batch
+ )
}
- }
+ };
}
/// A two-dimensional batch of column-oriented data with a defined
@@ -981,6 +1003,35 @@ mod tests {
assert_eq!(5, record_batch.column(1).len());
}
+ #[test]
+ fn create_binary_record_batch_from_variables() {
+ let binary_values = vec![b"a".as_slice()];
+ let large_binary_values = vec![b"xxx".as_slice()];
+
+ let record_batch = record_batch!(
+ ("a", Binary, binary_values),
+ ("b", LargeBinary, large_binary_values)
+ )
+ .unwrap();
+
+ assert_eq!(1, record_batch.num_rows());
+ assert_eq!(2, record_batch.num_columns());
+ assert_eq!(
+ &DataType::Binary,
+ record_batch.schema().field(0).data_type()
+ );
+ assert_eq!(
+ &DataType::LargeBinary,
+ record_batch.schema().field(1).data_type()
+ );
+
+ let binary = record_batch.column(0).as_binary::<i32>();
+ assert_eq!(b"a", binary.value(0));
+
+ let large_binary = record_batch.column(1).as_binary::<i64>();
+ assert_eq!(b"xxx", large_binary.value(0));
+ }
+
#[test]
fn byte_size_should_not_regress() {
let schema = Schema::new(vec