Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Removed clone requirement in Struct -> RecordBatch. (#307)
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao authored Aug 22, 2021
1 parent 83256a8 commit 0573369
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 24 deletions.
14 changes: 14 additions & 0 deletions src/array/struct_.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,20 @@ impl StructArray {
}
}

pub fn into_data(self) -> (Vec<Field>, Vec<Arc<dyn Array>>, Option<Bitmap>) {
let Self {
data_type,
values,
validity,
} = self;
let fields = if let DataType::Struct(fields) = data_type {
fields
} else {
unreachable!()
};
(fields, values, validity)
}

pub fn slice(&self, offset: usize, length: usize) -> Self {
let validity = self.validity.clone().map(|x| x.slice(offset, length));
Self {
Expand Down
19 changes: 8 additions & 11 deletions src/record_batch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -339,17 +339,14 @@ impl Default for RecordBatchOptions {
}
}

impl From<&StructArray> for RecordBatch {
/// Create a record batch from struct array.
///
/// This currently does not flatten and nested struct types
fn from(struct_array: &StructArray) -> Self {
if let DataType::Struct(fields) = struct_array.data_type() {
let schema = Arc::new(Schema::new(fields.clone()));
let columns = struct_array.values().to_vec();
RecordBatch { schema, columns }
} else {
unreachable!("unable to get datatype as struct")
impl From<StructArray> for RecordBatch {
/// # Panics iff the null count of the array is not null.
fn from(array: StructArray) -> Self {
assert!(array.null_count() == 0);
let (fields, values, _) = array.into_data();
RecordBatch {
schema: Arc::new(Schema::new(fields)),
columns: values,
}
}
}
Expand Down
22 changes: 9 additions & 13 deletions tests/it/record_batch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,22 +91,18 @@ fn number_of_fields_mismatch() {
fn from_struct_array() {
let boolean = Arc::new(BooleanArray::from_slice(&[false, false, true, true])) as ArrayRef;
let int = Arc::new(Int32Array::from_slice(&[42, 28, 19, 31])) as ArrayRef;
let struct_array = StructArray::from_data(
vec![
Field::new("b", DataType::Boolean, false),
Field::new("c", DataType::Int32, false),
],
vec![boolean.clone(), int.clone()],
None,
);

let batch = RecordBatch::from(&struct_array);
let fields = vec![
Field::new("b", DataType::Boolean, false),
Field::new("c", DataType::Int32, false),
];

let array = StructArray::from_data(fields.clone(), vec![boolean.clone(), int.clone()], None);

let batch = RecordBatch::from(array);
assert_eq!(2, batch.num_columns());
assert_eq!(4, batch.num_rows());
assert_eq!(
struct_array.data_type(),
&DataType::Struct(batch.schema().fields().to_vec())
);
assert_eq!(&fields, batch.schema().fields());
assert_eq!(boolean.as_ref(), batch.column(0).as_ref());
assert_eq!(int.as_ref(), batch.column(1).as_ref());
}

0 comments on commit 0573369

Please sign in to comment.