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

Removed clone requirement in StructArray -> RecordBatch #307

Merged
merged 1 commit into from
Aug 22, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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());
}