Skip to content

Commit d955609

Browse files
author
Tanvir Alam
committed
Fix index out-of-bounds panic in AbsoluteRowIdOptimizer statistics projection
Signed-off-by: Tanvir Alam <tanvralm@amazon.com>
1 parent c496725 commit d955609

File tree

1 file changed

+7
-5
lines changed

1 file changed

+7
-5
lines changed

plugins/engine-datafusion/jni/src/absolute_row_id_optimizer.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ use datafusion::{
2626
physical_optimizer::PhysicalOptimizerRule,
2727
physical_plan::{projection::ProjectionExec, ExecutionPlan},
2828
};
29+
use datafusion::physical_plan::ColumnStatistics;
2930
use datafusion_datasource::TableSchema;
3031

3132
#[derive(Debug)]
@@ -109,21 +110,22 @@ impl AbsoluteRowIdOptimizer {
109110
) -> Result<ProjectionExec, DataFusionError> {
110111
let (new_schema, new_projections) =
111112
self.build_updated_file_source_schema(datasource, data_source_exec_schema.clone());
112-
113+
113114
let table_partition_cols = datasource.table_partition_cols().clone();
114115
let new_table_schema = TableSchema::new(new_schema.clone(), table_partition_cols);
115-
116+
116117
use datafusion::datasource::physical_plan::ParquetSource;
117118
let new_file_source = Arc::new(ParquetSource::new(new_table_schema));
118-
119+
119120
let file_scan_config = FileScanConfigBuilder::from(datasource.clone())
120121
.with_source(new_file_source)
121122
.with_projection_indices(Some(new_projections))
122123
.expect("Failed to set projection indices")
124+
.with_statistics(datasource.statistics().add_column_statistics(ColumnStatistics::new_unknown()))
123125
.build();
124126

125127
let new_datasource = DataSourceExec::from_data_source(file_scan_config);
126-
128+
127129
let projection_exprs = self
128130
.build_projection_exprs(&new_datasource.schema())
129131
.expect("Failed to build projection expressions");
@@ -153,7 +155,7 @@ impl PhysicalOptimizerRule for AbsoluteRowIdOptimizer {
153155
return Ok(Transformed::new(Arc::new(projection), true, TreeNodeRecursion::Continue));
154156

155157
}
156-
158+
157159
Ok(Transformed::no(node))
158160
})?;
159161

0 commit comments

Comments
 (0)