fix

viirya · Nov 18, 2024 · 46bad5f · 46bad5f
1 parent e16ca2a
commit 46bad5f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 12 deletions.
diff --git 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<usize> =
-                        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<usize> =
+                    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
@@ -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()