Skip to content

Commit

Permalink
Add support for empty projection in RecordBatch::project (apache#2691)
Browse files Browse the repository at this point in the history
* Add support for empty projection in RecordBatch::project

* Simplify
  • Loading branch information
Dandandan committed Sep 9, 2022
1 parent 5e2b4c7 commit 5ccf73e
Showing 1 changed file with 28 additions and 1 deletion.
29 changes: 28 additions & 1 deletion arrow/src/record_batch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -212,7 +212,14 @@ impl RecordBatch {
})
.collect::<Result<Vec<_>>>()?;

RecordBatch::try_new(SchemaRef::new(projected_schema), batch_fields)
RecordBatch::try_new_with_options(
SchemaRef::new(projected_schema),
batch_fields,
&RecordBatchOptions {
match_field_names: true,
row_count: Some(self.row_count),
},
)
}

/// Returns the number of columns in the record batch.
Expand Down Expand Up @@ -865,6 +872,26 @@ mod tests {
assert_eq!(expected, record_batch.project(&[0, 2]).unwrap());
}

#[test]
fn project_empty() {
let c: ArrayRef = Arc::new(StringArray::from(vec!["d", "e", "f"]));

let record_batch =
RecordBatch::try_from_iter(vec![("c", c.clone())]).expect("valid conversion");

let expected = RecordBatch::try_new_with_options(
Arc::new(Schema::empty()),
vec![],
&RecordBatchOptions {
match_field_names: true,
row_count: Some(3),
},
)
.expect("valid conversion");

assert_eq!(expected, record_batch.project(&[]).unwrap());
}

#[test]
fn test_no_column_record_batch() {
let schema = Arc::new(Schema::new(vec![]));
Expand Down

0 comments on commit 5ccf73e

Please sign in to comment.