Skip to content

Commit

Permalink
Executing ProjectionExec with no column should not return an Err (#4912)
Browse files Browse the repository at this point in the history
* Fix projection with no column

* Fix clippy
  • Loading branch information
viirya authored Jan 15, 2023
1 parent 2801c8c commit fa8ad40
Showing 1 changed file with 25 additions and 2 deletions.
27 changes: 25 additions & 2 deletions datafusion/core/src/physical_plan/projection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use crate::physical_plan::{
};
use arrow::datatypes::{Field, Schema, SchemaRef};
use arrow::error::Result as ArrowResult;
use arrow::record_batch::RecordBatch;
use arrow::record_batch::{RecordBatch, RecordBatchOptions};
use log::debug;

use super::expressions::{Column, PhysicalSortExpr};
Expand Down Expand Up @@ -328,7 +328,13 @@ impl ProjectionStream {
.map(|r| r.map(|v| v.into_array(batch.num_rows())))
.collect::<Result<Vec<_>>>()?;

RecordBatch::try_new(self.schema.clone(), arrays)
if arrays.is_empty() {
let options =
RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
RecordBatch::try_new_with_options(self.schema.clone(), arrays, &options)
} else {
RecordBatch::try_new(self.schema.clone(), arrays)
}
}
}

Expand Down Expand Up @@ -372,6 +378,7 @@ impl RecordBatchStream for ProjectionStream {
mod tests {

use super::*;
use crate::physical_plan::common::collect;
use crate::physical_plan::expressions::{self, col};
use crate::prelude::SessionContext;
use crate::scalar::ScalarValue;
Expand Down Expand Up @@ -418,6 +425,22 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn project_no_column() -> Result<()> {
let session_ctx = SessionContext::new();
let task_ctx = session_ctx.task_ctx();

let csv = test::scan_partitioned_csv(1)?;
let expected = collect(csv.execute(0, task_ctx.clone())?).await.unwrap();

let projection = ProjectionExec::try_new(vec![], csv)?;
let stream = projection.execute(0, task_ctx.clone())?;
let output = collect(stream).await.unwrap();
assert_eq!(output.len(), expected.len());

Ok(())
}

#[tokio::test]
async fn test_stats_projection_columns_only() {
let source = Statistics {
Expand Down

0 comments on commit fa8ad40

Please sign in to comment.