diff --git a/src/mito2/src/read/read_columns.rs b/src/mito2/src/read/read_columns.rs index ae5876237b..2178fa2538 100644 --- a/src/mito2/src/read/read_columns.rs +++ b/src/mito2/src/read/read_columns.rs @@ -170,7 +170,7 @@ fn normalize_nested_paths(nested_paths: Vec) -> Vec { normalized } -fn merge_nested_paths(merged: &mut Vec, incoming: Vec) { +pub(crate) fn merge_nested_paths(merged: &mut Vec, incoming: Vec) { for path in incoming { if merged .iter() diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 8894e970d2..3fe1aad6d7 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -61,7 +61,8 @@ use crate::read::flat_projection::FlatProjectionMapper; use crate::read::range::{FileRangeBuilder, MemRangeBuilder, RangeMeta, RowGroupIndex}; use crate::read::range_cache::{ScanRequestFingerprint, implied_time_range_from_exprs}; use crate::read::read_columns::{ - ReadColumns, merge, read_columns_from_predicate, read_columns_from_projection, + ReadColumns, merge, merge_nested_paths, read_columns_from_predicate, + read_columns_from_projection, }; use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; @@ -1340,7 +1341,7 @@ fn fill_json_nested_paths_from_hint( let mut paths = Vec::new(); let mut current = vec![column_name.clone()]; collect_json_nested_paths(json_type, &mut current, &mut paths); - read_column.nested_paths = paths; + merge_nested_paths(&mut read_column.nested_paths, paths) } } diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 53e16f86eb..39d8691a13 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -452,9 +452,22 @@ where .fields .into_iter() .zip(nested_paths) - .map(|(field, nested_paths)| { - let nested_paths = nested_paths.iter().map(Vec::as_slice).collect::>(); - prune_field_by_nested_paths(field, &nested_paths) + .map(|(field, paths)| { + if matches!(field.data_type(), ArrowDataType::Struct(_)) && !paths.is_empty() { + let child_paths = paths + .iter() + .map(|path| { + if path.first().is_some_and(|root| root == field.name()) { + &path[1..] + } else { + path + } + }) + .collect::>(); + prune_field_by_nested_paths(field, &child_paths) + } else { + field.clone() + } }) .collect::>(); schema.fields = fields.into() @@ -465,32 +478,14 @@ fn prune_field_by_nested_paths(field: &FieldRef, nested_paths: &[&[String]]) -> return field.clone(); }; - if nested_paths.is_empty() { - return field.clone(); - } - - let child_paths = nested_paths - .iter() - .map(|path| { - if path.first().is_some_and(|root| root == field.name()) { - &path[1..] - } else { - *path - } - }) - .collect::>(); - if child_paths.iter().any(|path| path.is_empty()) { - return field.clone(); - } - let pruned_fields = fields .iter() .filter_map(|field| { - let child_paths = child_paths + let child_paths = nested_paths .iter() .filter_map(|path| { path.first() - .is_some_and(|name| name.as_str() == field.name().as_str()) + .is_some_and(|name| name == field.name()) .then_some(&path[1..]) }) .collect::>(); diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index ba321f4383..45f2e9c905 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -905,12 +905,16 @@ impl CreateRequestBuilder { column_id += 1; } for i in 0..self.field_num { + let mut column_schema = + ColumnSchema::new(format!("field_{i}"), self.field_datatype.clone(), nullable); + if self.field_datatype.is_json() { + column_schema.mut_metadata().insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + JsonExtensionType::NAME.to_string(), + ); + } column_metadatas.push(ColumnMetadata { - column_schema: ColumnSchema::new( - format!("field_{i}"), - self.field_datatype.clone(), - nullable, - ), + column_schema, semantic_type: SemanticType::Field, column_id, });