Skip to content

Commit 5ccf73e

Browse files
authored
Add support for empty projection in RecordBatch::project (apache#2691)
* Add support for empty projection in RecordBatch::project * Simplify
1 parent 5e2b4c7 commit 5ccf73e

File tree

1 file changed

+28
-1
lines changed

1 file changed

+28
-1
lines changed

arrow/src/record_batch.rs

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,14 @@ impl RecordBatch {
212212
})
213213
.collect::<Result<Vec<_>>>()?;
214214

215-
RecordBatch::try_new(SchemaRef::new(projected_schema), batch_fields)
215+
RecordBatch::try_new_with_options(
216+
SchemaRef::new(projected_schema),
217+
batch_fields,
218+
&RecordBatchOptions {
219+
match_field_names: true,
220+
row_count: Some(self.row_count),
221+
},
222+
)
216223
}
217224

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

875+
#[test]
876+
fn project_empty() {
877+
let c: ArrayRef = Arc::new(StringArray::from(vec!["d", "e", "f"]));
878+
879+
let record_batch =
880+
RecordBatch::try_from_iter(vec![("c", c.clone())]).expect("valid conversion");
881+
882+
let expected = RecordBatch::try_new_with_options(
883+
Arc::new(Schema::empty()),
884+
vec![],
885+
&RecordBatchOptions {
886+
match_field_names: true,
887+
row_count: Some(3),
888+
},
889+
)
890+
.expect("valid conversion");
891+
892+
assert_eq!(expected, record_batch.project(&[]).unwrap());
893+
}
894+
868895
#[test]
869896
fn test_no_column_record_batch() {
870897
let schema = Arc::new(Schema::new(vec![]));

0 commit comments

Comments
 (0)