diff --git a/native/core/src/execution/datafusion/planner.rs b/native/core/src/execution/datafusion/planner.rs index cd79e8e0b..8cd301612 100644 --- a/native/core/src/execution/datafusion/planner.rs +++ b/native/core/src/execution/datafusion/planner.rs @@ -1023,17 +1023,15 @@ impl PhysicalPlanner { .with_file_groups(file_groups); // Check for projection, if so generate the vector and add to FileScanConfig. - if !required_schema_arrow.fields.is_empty() { - let mut projection_vector: Vec = - Vec::with_capacity(required_schema_arrow.fields.len()); - // TODO: could be faster with a hashmap rather than iterating over data_schema_arrow with index_of. - required_schema_arrow.fields.iter().for_each(|field| { - projection_vector.push(data_schema_arrow.index_of(field.name()).unwrap()); - }); + let mut projection_vector: Vec = + Vec::with_capacity(required_schema_arrow.fields.len()); + // TODO: could be faster with a hashmap rather than iterating over data_schema_arrow with index_of. + required_schema_arrow.fields.iter().for_each(|field| { + projection_vector.push(data_schema_arrow.index_of(field.name()).unwrap()); + }); - assert_eq!(projection_vector.len(), required_schema_arrow.fields.len()); - file_scan_config = file_scan_config.with_projection(Some(projection_vector)); - } + assert_eq!(projection_vector.len(), required_schema_arrow.fields.len()); + file_scan_config = file_scan_config.with_projection(Some(projection_vector)); let mut table_parquet_options = TableParquetOptions::new(); // TODO: Maybe these are configs? diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index e3b71f959..1dec5173b 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -287,7 +287,7 @@ fn prepare_output( let results = output_batch.columns(); let num_rows = output_batch.num_rows(); - if results.len() < num_cols { + if results.len() != num_cols { return Err(CometError::Internal(format!( "Output column count mismatch: expected {num_cols}, got {}", results.len() @@ -305,7 +305,7 @@ fn prepare_output( } let mut i = 0; - while i < num_cols { + while i < results.len() { let array_ref = results.get(i).ok_or(CometError::IndexOutOfBounds(i))?; array_ref .to_data()