diff --git a/src/cmd/src/datanode/parquetbench.rs b/src/cmd/src/datanode/parquetbench.rs index 301c6440ec..2b46a775f7 100644 --- a/src/cmd/src/datanode/parquetbench.rs +++ b/src/cmd/src/datanode/parquetbench.rs @@ -554,7 +554,7 @@ async fn run_flat_prune_iteration( let reader_builder = ParquetReaderBuilder::new(table_dir, path_type, file_handle, object_store) .expected_metadata(Some(region_meta)) .cache(CacheStrategy::Disabled) - .projection(projection.map(ReadColumns::from_deduped_column_ids)); + .projection(projection.map(ReadColumns::new)); let mut reader_metrics = ReaderMetrics::default(); let start = Instant::now(); let mut stats = IterationStats::default(); diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 27db7daa5a..ac1f5b6691 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -49,7 +49,7 @@ use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef} use smallvec::SmallVec; use snafu::{OptionExt, ResultExt}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; -use store_api::storage::{ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; +use store_api::storage::{ColumnId, ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; use crate::cache::cache_size::parquet_meta_size; use crate::cache::file_cache::{FileType, IndexKey}; @@ -65,6 +65,7 @@ use crate::memtable::record_batch_estimated_size; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS}; use crate::read::Batch; use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue}; +use crate::read::read_columns::JsonTargetTypes; use crate::sst::file::{RegionFileId, RegionIndexId}; use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::sst::parquet::read_columns::ParquetReadColumns; @@ -2025,6 +2026,11 @@ pub struct SelectorResultValue { pub result: SelectorResult, /// The read columns of rows. pub read_cols: ParquetReadColumns, + /// JSON2 target types used by flat-format reads. + /// + /// JSON2 projection is query-driven; the same parquet columns can produce + /// different cached batches under different type hints. + pub json_target_types: JsonTargetTypes, } impl SelectorResultValue { @@ -2033,6 +2039,7 @@ impl SelectorResultValue { SelectorResultValue { result: SelectorResult::PrimaryKey(result), read_cols, + json_target_types: Arc::default(), } } @@ -2040,21 +2047,26 @@ impl SelectorResultValue { pub fn new_flat( result: Vec, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, ) -> SelectorResultValue { SelectorResultValue { result: SelectorResult::Flat(result), read_cols, + json_target_types, } } /// Returns memory used by the value (estimated). fn estimated_size(&self) -> usize { - match &self.result { + let result_size: usize = match &self.result { SelectorResult::PrimaryKey(batches) => { batches.iter().map(|batch| batch.memory_size()).sum() } SelectorResult::Flat(batches) => batches.iter().map(record_batch_estimated_size).sum(), - } + }; + result_size + + self.json_target_types.len() + * (mem::size_of::() + mem::size_of::()) } } @@ -2752,7 +2764,7 @@ mod tests { region_id: RegionId::new(1, 1), row_groups: vec![(FileId::random(), 0)], scan: ScanRequestFingerprintBuilder { - read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()), + read_columns: ReadColumns::new(std::iter::empty()), read_column_types: vec![], filters: vec!["tag_0 = 1".to_string()], time_filters: vec![], diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index 89b70f89c0..f862d11e1a 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use common_time::Timestamp; @@ -118,15 +118,28 @@ impl CompactionSstReaderBuilder<'_> { }; let projection = (0..self.metadata.column_metadatas.len()).collect(); - let read_columns = ReadColumns::from_deduped_column_ids( - self.metadata.column_metadatas.iter().map(|x| x.column_id), - ); - let mapper = FlatProjectionMapper::new_with_read_columns( - &self.metadata, - projection, - read_columns, - json_type_hint.as_ref(), - )?; + let read_column_ids = self + .metadata + .column_metadatas + .iter() + .map(|x| x.column_id) + .collect::>(); + let json_target_types = json_type_hint + .as_ref() + .map(|hint| { + hint.iter() + .filter_map(|(col_name, json_type)| { + self.metadata + .column_by_name(col_name) + .map(|col| (col.column_id, json_type.clone())) + }) + .collect::>() + }) + .unwrap_or_default(); + let read_columns = + ReadColumns::new(read_column_ids).with_json_target_types(json_target_types); + let mapper = + FlatProjectionMapper::new_with_read_columns(&self.metadata, projection, read_columns)?; let mut scan_input = ScanInput::new(self.sst_layer, mapper) .with_files(self.inputs.to_vec()) diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index b9eed37cb3..4baa707f4e 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -38,7 +38,7 @@ use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution}; use crate::config::MitoConfig; use crate::error::Error; -use crate::read::read_columns::{ReadColumn, ReadColumns}; +use crate::read::read_columns::ReadColumns; use crate::read::scan_region::Scanner; use crate::test_util; use crate::test_util::{CreateRequestBuilder, TestEnv}; @@ -108,18 +108,19 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul }; assert_eq!( seq_scan.input().read_cols, - ReadColumns::from_deduped_column_ids([1, 0]) + ReadColumns::new([1, 0]) + .with_json_target_types(BTreeMap::from([(1, JsonNativeType::Variant)])) ); let stream = scanner.scan().await?; let batches = RecordBatches::try_collect(stream).await?; let expected = r#" -+------------------------------------------+-------+ -| field_0 | tag_0 | -+------------------------------------------+-------+ -| {a: {x: 10, y: ignored-a}, b: ignored-b} | tag-1 | -| {a: {x: 20, y: ignored-c}, b: ignored-d} | tag-2 | -+------------------------------------------+-------+ ++------------------------------------------------+-------+ +| field_0 | tag_0 | ++------------------------------------------------+-------+ +| {"a":{"x":10,"y":"ignored-a"},"b":"ignored-b"} | tag-1 | +| {"a":{"x":20,"y":"ignored-c"},"b":"ignored-d"} | tag-2 | ++------------------------------------------------+-------+ "#; assert_eq!(batches.pretty_print()?, expected.trim()); @@ -149,19 +150,8 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul // whole JSON2 struct. tag_0 is still read as a normal root column. assert_eq!( seq_scan.input().read_cols, - ReadColumns { - cols: vec![ - ReadColumn::new( - 1, - vec![vec![ - "field_0".to_string(), - "a".to_string(), - "x".to_string() - ]] - ), - ReadColumn::new(0, vec![]), - ] - } + ReadColumns::new([1, 0]) + .with_json_target_types(BTreeMap::from([(1, json_type_hint["field_0"].clone())])) ); // The scanner should still return a valid RecordBatch in the requested logical projection. @@ -189,10 +179,7 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul let Scanner::Seq(seq_scan) = &scanner else { unreachable!(); }; - assert_eq!( - seq_scan.input().read_cols, - ReadColumns::from_deduped_column_ids([0]) - ); + assert_eq!(seq_scan.input().read_cols, ReadColumns::new([0])); let stream = scanner.scan().await?; let batches = RecordBatches::try_collect(stream).await?; @@ -205,6 +192,19 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul +-------+ "#; assert_eq!(batches.pretty_print()?, expected.trim()); + + // JSON type hints are derived from the whole query plan, so they may include + // columns unrelated to this scan. Mito2 should ignore hints for unknown + // columns (`other_json`) instead of failing. + let request = ScanRequest { + projection: Some(vec![0]), + json_type_hint: HashMap::from([("other_json".to_string(), JsonNativeType::String)]), + ..Default::default() + }; + let scanner = engine.scanner(region_id, request).await?; + let stream = scanner.scan().await?; + let batches = RecordBatches::try_collect(stream).await?; + assert_eq!(batches.pretty_print()?, expected.trim()); Ok(()) } diff --git a/src/mito2/src/memtable/bulk/context.rs b/src/mito2/src/memtable/bulk/context.rs index af38ef66c1..994d3a027d 100644 --- a/src/mito2/src/memtable/bulk/context.rs +++ b/src/mito2/src/memtable/bulk/context.rs @@ -72,9 +72,9 @@ impl BulkIterContext { let codec = build_primary_key_codec(®ion_metadata); let read_cols = if let Some(col_ids) = projection { - ReadColumns::from_deduped_column_ids(col_ids.iter().copied()) + ReadColumns::new(col_ids.iter().copied()) } else { - ReadColumns::from_deduped_column_ids( + ReadColumns::new( region_metadata .column_metadatas .iter() diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index af0bdc58c5..3fa625694c 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -25,11 +25,9 @@ use datatypes::arrow::compute::{TakeOptions, take}; use datatypes::arrow::datatypes::{FieldRef, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; -use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::DataType; use datatypes::value::Value; use datatypes::vectors::VectorRef; -use datatypes::vectors::json::array::JsonArray; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec, build_primary_key_codec_with_fields, @@ -40,8 +38,8 @@ use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::ColumnId; use crate::error::{ - CompatReaderSnafu, ComputeArrowSnafu, ConvertValueSnafu, CreateDefaultSnafu, DecodeSnafu, - EncodeSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu, + CompatReaderSnafu, ComputeArrowSnafu, CreateDefaultSnafu, DecodeSnafu, EncodeSnafu, + NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu, }; use crate::read::flat_projection::{FlatProjectionMapper, flat_projected_columns}; use crate::sst::parquet::flat_format::{FlatReadFormat, primary_key_column_index}; @@ -99,20 +97,12 @@ impl FlatCompatBatch { let actual = read_format.metadata(); let format_projection = read_format.format_projection(); let mut actual_schema = flat_projected_columns(actual, format_projection); - if read_format - .arrow_schema() - .fields() - .iter() - .any(is_json2_extension_type) - { - for field in read_format.arrow_schema().fields() { - if is_json2_extension_type(field) - && let Some(column_id) = - actual.column_by_name(field.name()).map(|x| x.column_id) - && let Some(i) = actual_schema.iter().position(|x| x.0 == column_id) - { - actual_schema[i].1 = ConcreteDataType::from_arrow_type(field.data_type()); - } + for (column_id, target_type) in read_format.json_target_types().iter() { + if let Some(i) = actual_schema + .iter() + .position(|(actual_column_id, _)| actual_column_id == column_id) + { + actual_schema[i].1 = ConcreteDataType::json2(target_type.clone()); } } @@ -270,16 +260,9 @@ impl FlatCompatBatch { let old_column = batch.column(*pos); if let Some(ty) = cast_type { - let casted = if let Some(json_type) = ty.as_json() - && json_type.is_json2() - { - JsonArray::from(old_column) - .project_to(&json_type.as_arrow_type()) - .context(ConvertValueSnafu)? - } else { + let casted = datatypes::arrow::compute::cast(old_column, &ty.as_arrow_type()) - .context(ComputeArrowSnafu)? - }; + .context(ComputeArrowSnafu)?; Ok(casted) } else { Ok(old_column.clone()) @@ -777,7 +760,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 1, 2, 3]), + ReadColumns::new([0, 1, 2, 3]), None, "test", false, @@ -864,13 +847,12 @@ mod tests { let mapper = FlatProjectionMapper::new_with_read_columns( &expected_metadata, vec![1, 2], - ReadColumns::from_deduped_column_ids([1, 2, 3]), - None, + ReadColumns::new([1, 2, 3]), ) .unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([1, 2, 3]), + ReadColumns::new([1, 2, 3]), None, "test", false, @@ -959,7 +941,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 1, 2, 3]), + ReadColumns::new([0, 1, 2, 3]), None, "test", false, @@ -1041,7 +1023,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata, - ReadColumns::from_deduped_column_ids([0, 1, 2]), + ReadColumns::new([0, 1, 2]), None, "test", false, @@ -1109,7 +1091,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 2, 3]), + ReadColumns::new([0, 2, 3]), None, "test", true, diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index fbd6088f67..e631cf45db 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -14,7 +14,6 @@ //! Utilities for projection on flat format. -use std::collections::HashMap; use std::sync::Arc; use api::v1::SemanticType; @@ -27,7 +26,7 @@ use datatypes::arrow::array::Array; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{ConcreteDataType, DataType}; -use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; +use datatypes::schema::{Schema, SchemaRef}; use datatypes::types::JsonType; use datatypes::types::json_type::JsonNativeType; use datatypes::value::Value; @@ -86,8 +85,8 @@ impl FlatProjectionMapper { ) -> Result { let projection: Vec<_> = projection.into_iter().collect(); let read_column_ids = read_column_ids_from_projection(metadata, &projection)?; - let read_cols = ReadColumns::from_deduped_column_ids(read_column_ids); - Self::new_with_read_columns(metadata, projection, read_cols, None) + let read_cols = ReadColumns::new(read_column_ids); + Self::new_with_read_columns(metadata, projection, read_cols) } /// Returns a new mapper with output projection and explicit read columns. @@ -95,7 +94,6 @@ impl FlatProjectionMapper { metadata: &RegionMetadataRef, projection: Vec, read_cols: ReadColumns, - json_type_hint: Option<&HashMap>, ) -> Result { // If the original projection is empty. let is_empty_projection = projection.is_empty(); @@ -115,7 +113,11 @@ impl FlatProjectionMapper { output_col_ids.push(col.column_id); let mut schema = col.column_schema.clone(); - maybe_concretize_json2_datatype(&mut schema, json_type_hint); + if let Some(data_type) = + json2_read_datatype(col.column_id, &schema.data_type, &read_cols) + { + schema.data_type = data_type; + } col_schemas.push(schema); } @@ -124,6 +126,7 @@ impl FlatProjectionMapper { // TODO(yingwen): Support different flat schema options. let format_projection = FormatProjection::compute_format_projection( + metadata, &id_to_index, // All columns with internal columns. metadata.column_metadatas.len() + 3, @@ -133,21 +136,8 @@ impl FlatProjectionMapper { let mut batch_schema = flat_projected_columns(metadata, &format_projection); for (column_id, data_type) in batch_schema.iter_mut() { - if let Some(json_type) = data_type.as_json() - && json_type.is_json2() - { - if let Some(concretized) = metadata - .column_by_id(*column_id) - .and_then(|metadata| { - json_type_hint.and_then(|x| x.get(&metadata.column_schema.name).cloned()) - }) - .map(ConcreteDataType::json2) - { - *data_type = concretized; - } else if is_empty_json2_type(json_type) { - // see `merge_scan::maybe_amend_json2_field` - *data_type = ConcreteDataType::json2(JsonNativeType::Variant); - } + if let Some(updated) = json2_read_datatype(*column_id, data_type, &read_cols) { + *data_type = updated; } } @@ -401,24 +391,25 @@ impl FlatProjectionMapper { } } -fn maybe_concretize_json2_datatype( - schema: &mut ColumnSchema, - json_type_hint: Option<&HashMap>, -) { - if let Some(json_type) = schema.data_type.as_json() - && json_type.is_json2() - { - if let Some(concretized) = json_type_hint - .and_then(|x| x.get(&schema.name)) - .cloned() - .map(ConcreteDataType::json2) - { - schema.data_type = concretized; - } else if is_empty_json2_type(json_type) { - // see `merge_scan::maybe_amend_json2_field` - schema.data_type = ConcreteDataType::json2(JsonNativeType::Variant); - } +fn json2_read_datatype( + column_id: ColumnId, + data_type: &ConcreteDataType, + read_cols: &ReadColumns, +) -> Option { + let json_type = data_type.as_json()?; + if !json_type.is_json2() { + return None; } + + if let Some(concretized) = read_cols.json_target_type(column_id).cloned() { + return Some(ConcreteDataType::json2(concretized)); + } + + if is_empty_json2_type(json_type) { + return Some(ConcreteDataType::json2(JsonNativeType::Variant)); + } + + None } fn is_empty_json2_type(json_type: &JsonType) -> bool { @@ -547,9 +538,8 @@ impl CompactionProjectionMapper { .collect::>(); let read_col_ids = metadata.column_metadatas.iter().map(|col| col.column_id); - let read_cols = ReadColumns::from_deduped_column_ids(read_col_ids); - let mapper = - FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols, None)?; + let read_cols = ReadColumns::new(read_col_ids); + let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?; let assembler = DfBatchAssembler::new(mapper.output_schema()); Ok(Self { mapper, assembler }) @@ -617,7 +607,6 @@ impl DfBatchAssembler { #[cfg(test)] mod tests { use datatypes::schema::ColumnSchema; - use datatypes::types::json_type::JsonObjectType; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; @@ -687,18 +676,10 @@ mod tests { #[test] fn test_json_type_hint_does_not_concretize_legacy_json() { let metadata = metadata_with_legacy_json(); - let hint = HashMap::from([( - "j".to_string(), - JsonNativeType::Object(JsonObjectType::from([( - "a".to_string(), - JsonNativeType::i64(), - )])), - )]); let mapper = FlatProjectionMapper::new_with_read_columns( &metadata, vec![0, 1], - ReadColumns::from_deduped_column_ids([0, 1]), - Some(&hint), + ReadColumns::new([0, 1]), ) .unwrap(); diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index 6e1a62d806..c2316525f7 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -29,6 +29,7 @@ use crate::cache::{ selector_result_cache_hit, selector_result_cache_miss, }; use crate::error::{ComputeArrowSnafu, Result}; +use crate::read::read_columns::JsonTargetTypes; use crate::read::{ Batch, BatchReader, BoxedBatchReader, BoxedRecordBatchStream, timestamp_array_to_i64_slice, }; @@ -137,6 +138,7 @@ impl FlatRowGroupLastRowCachedReader { row_group_idx: usize, cache_strategy: CacheStrategy, read_cols: &ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, ) -> Self { let key = SelectorResultKey { @@ -148,13 +150,14 @@ impl FlatRowGroupLastRowCachedReader { if let Some(value) = cache_strategy.get_selector_result(&key) { let is_flat = matches!(&value.result, SelectorResult::Flat(_)); let schema_matches = value.read_cols == *read_cols; - if is_flat && schema_matches { + let json_target_types_matches = value.json_target_types == json_target_types; + if is_flat && schema_matches && json_target_types_matches { Self::new_hit(value) } else { - Self::new_miss(key, read_cols, reader, cache_strategy) + Self::new_miss(key, read_cols, json_target_types, reader, cache_strategy) } } else { - Self::new_miss(key, read_cols, reader, cache_strategy) + Self::new_miss(key, read_cols, json_target_types, reader, cache_strategy) } } @@ -174,6 +177,7 @@ impl FlatRowGroupLastRowCachedReader { fn new_miss( key: SelectorResultKey, read_cols: &ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, cache_strategy: CacheStrategy, ) -> Self { @@ -181,6 +185,7 @@ impl FlatRowGroupLastRowCachedReader { Self::Miss(FlatRowGroupLastRowReader::new( key, read_cols.clone(), + json_target_types, reader, cache_strategy, )) @@ -260,6 +265,7 @@ pub(crate) struct FlatRowGroupLastRowReader { yielded_batches: Vec, cache_strategy: CacheStrategy, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, /// Accumulates small selector-output batches before concatenating. pending: BatchBuffer, } @@ -268,6 +274,7 @@ impl FlatRowGroupLastRowReader { fn new( key: SelectorResultKey, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, cache_strategy: CacheStrategy, ) -> Self { @@ -278,6 +285,7 @@ impl FlatRowGroupLastRowReader { yielded_batches: vec![], cache_strategy, read_cols, + json_target_types, pending: BatchBuffer::new(), } } @@ -326,6 +334,7 @@ impl FlatRowGroupLastRowReader { let value = Arc::new(SelectorResultValue::new_flat( batches, self.read_cols.clone(), + self.json_target_types.clone(), )); self.cache_strategy.put_selector_result(self.key, value); } diff --git a/src/mito2/src/read/projection.rs b/src/mito2/src/read/projection.rs index 27fe76345b..c6d4baee3a 100644 --- a/src/mito2/src/read/projection.rs +++ b/src/mito2/src/read/projection.rs @@ -295,8 +295,7 @@ mod tests { let mapper = FlatProjectionMapper::new_with_read_columns( &metadata, vec![4, 1], - ReadColumns::from_deduped_column_ids([4, 1, 3]), - None, + ReadColumns::new([4, 1, 3]), ) .unwrap(); assert_eq!(&[4, 1, 3], mapper.read_columns().column_ids().as_slice()); diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 036a574c1f..ee0e421f86 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -852,7 +852,7 @@ pub fn bench_cache_flat_range_stream( let cache_strategy = CacheStrategy::EnableAll(cache_manager); let fingerprint = ScanRequestFingerprintBuilder { - read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()), + read_columns: ReadColumns::new(std::iter::empty()), read_column_types: vec![], filters: vec![], time_filters: vec![], @@ -916,7 +916,7 @@ mod tests { filter_deleted: bool, partition_expr_version: u64, ) -> ScanRequestFingerprint { - let read_columns = ReadColumns::from_deduped_column_ids([1, 2]); + let read_columns = ReadColumns::new([1, 2]); ScanRequestFingerprintBuilder { read_columns, read_column_types: vec![None, None], diff --git a/src/mito2/src/read/read_columns.rs b/src/mito2/src/read/read_columns.rs index 21699b8c43..42a9f71ddc 100644 --- a/src/mito2/src/read/read_columns.rs +++ b/src/mito2/src/read/read_columns.rs @@ -12,119 +12,72 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; use std::mem; +use std::sync::Arc; -use store_api::storage::{ColumnId, NestedPath}; +use datatypes::types::json_type::JsonNativeType; +use store_api::storage::ColumnId; + +pub(crate) type JsonTargetTypes = Arc>; /// Logical columns to read from a region. /// -/// Read columns describe which logical columns and nested fields should be read -/// from storage. Each read column is identified by its [`ColumnId`], -/// which represents the root column in the storage schema. -/// -/// Nested fields under the column are specified by [`NestedPath`] entries. -/// Each path includes the root column name as its first element. -/// -/// For example, assume column id `9` corresponds to a root column named `j` -/// with nested fields: -/// -/// ```text -/// j -/// ├── a -/// └── b -/// └── c -/// ``` -/// -/// The following SQL: -/// -/// SELECT j.a, j.b.c FROM t -/// -/// may produce read columns like: -/// -/// ```text -/// ReadColumn { -/// column_id: 9, -/// nested_paths: [ -/// ["j", "a"], -/// ["j", "b", "c"], -/// ] -/// } -/// ``` -/// -/// If `nested_paths` is empty, the whole column will be read. +/// Read columns describe which logical root columns should be read from storage. +/// JSON2 columns can carry query-time target types that are later translated to +/// physical nested parquet paths by the parquet reader. #[derive(Debug, Default, Clone, PartialEq, Eq, Hash)] pub struct ReadColumns { - pub cols: Vec, + pub col_ids: Vec, + json_target_types: JsonTargetTypes, } impl ReadColumns { - pub fn from_deduped_column_ids(column_ids: I) -> Self + /// Creates read columns from logical column ids. + /// + /// This preserves the input order and duplicate ids. + pub fn new(col_ids: I) -> Self where I: IntoIterator, { - let cols = column_ids - .into_iter() - .map(|col_id| ReadColumn::new(col_id, vec![])) - .collect(); - ReadColumns { cols } + Self { + col_ids: col_ids.into_iter().collect(), + json_target_types: Arc::default(), + } + } + + pub fn with_json_target_types( + mut self, + json_target_types: BTreeMap, + ) -> Self { + self.json_target_types = Arc::new(json_target_types); + self } pub fn is_empty(&self) -> bool { - self.cols.is_empty() + self.col_ids.is_empty() } pub fn column_ids_iter(&self) -> impl Iterator + '_ { - self.cols.iter().map(|column| column.column_id) + self.col_ids.iter().copied() } pub fn column_ids(&self) -> Vec { self.column_ids_iter().collect() } - pub fn columns(&self) -> &[ReadColumn] { - &self.cols + pub fn json_target_types(&self) -> &JsonTargetTypes { + &self.json_target_types + } + + pub fn json_target_type(&self, column_id: ColumnId) -> Option<&JsonNativeType> { + self.json_target_types.get(&column_id) } pub fn estimated_size(&self) -> usize { - self.cols.capacity() * mem::size_of::() - + self - .cols - .iter() - .map(ReadColumn::estimated_size) - .sum::() - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct ReadColumn { - pub column_id: ColumnId, - /// Nested field paths under this column. - /// Empty means reading the whole column. - pub nested_paths: Vec, -} - -impl ReadColumn { - pub fn new(column_id: ColumnId, nested_paths: Vec) -> Self { - Self { - column_id, - nested_paths, - } - } - - pub fn nested_paths(&self) -> &[NestedPath] { - &self.nested_paths - } - - pub fn estimated_size(&self) -> usize { - mem::size_of::() - + self.nested_paths.capacity() * mem::size_of::() - + self - .nested_paths - .iter() - .map(|path| { - path.capacity() * mem::size_of::() - + path.iter().map(|node| node.capacity()).sum::() - }) - .sum::() + self.col_ids.capacity() * mem::size_of::() + + self.col_ids.len() * mem::size_of::() + + self.json_target_types.len() + * (mem::size_of::() + mem::size_of::()) } } diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 250f0d7fc2..c0340464ec 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -14,7 +14,7 @@ //! Scans a region according to the scan request. -use std::collections::HashSet; +use std::collections::{BTreeMap, HashSet}; use std::fmt; use std::num::NonZeroU64; use std::sync::Arc; @@ -35,18 +35,18 @@ use datafusion_expr::Expr; use datafusion_expr::utils::expr_to_columns; use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use datatypes::extension::json::is_json2_extension_type; +use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::JsonNativeType; use datatypes::value::timestamp_to_scalar_value; use futures::StreamExt; -use itertools::Itertools; use partition::expr::PartitionExpr; use smallvec::SmallVec; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::region_engine::{PartitionRange, RegionScannerRef}; use store_api::storage::{ - ColumnId, NestedPath, RegionId, ScanRequest, SequenceNumber, SequenceRange, - TimeSeriesDistribution, TimeSeriesRowSelector, + ColumnId, RegionId, ScanRequest, SequenceNumber, SequenceRange, TimeSeriesDistribution, + TimeSeriesRowSelector, }; use table::predicate::{Predicate, build_time_range_predicate, extract_time_range_from_expr}; use tokio::sync::{Semaphore, mpsc}; @@ -64,7 +64,7 @@ use crate::read::compat::{self, FlatCompatBatch}; 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::{ReadColumn, ReadColumns}; +use crate::read::read_columns::ReadColumns; use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; use crate::read::stream::ScanBatchStream; @@ -420,23 +420,7 @@ impl ScanRegion { let read_col_ids = self.build_read_col_ids(self.request.projection.as_deref(), &predicate)?; - - // Narrow JSON2 columns to avoid reading unnecessary nested fields. - // - // `read_col_ids` selects the root columns required by the projection and predicates, - // while nested projection is currently only applied to JSON2 columns, whose type hints - // further narrow them to the requested nested fields. - let has_structured_json = metadata - .schema - .arrow_schema() - .fields() - .iter() - .any(is_json2_extension_type); - let read_cols = if has_structured_json { - self.read_columns_with_json_type_hint(&read_col_ids) - } else { - ReadColumns::from_deduped_column_ids(read_col_ids.iter().copied()) - }; + let read_cols = self.build_read_columns(&read_col_ids)?; // The mapper always computes projected column ids as the schema of SSTs may change. let projection = self @@ -444,23 +428,7 @@ impl ScanRegion { .projection .clone() .unwrap_or_else(|| (0..metadata.column_metadatas.len()).collect()); - let json_type_hint = has_structured_json - .then_some(&self.request.json_type_hint) - .inspect(|json_type_hint| { - debug!( - "Concretized JSON type: {{{}}}", - json_type_hint - .iter() - .map(|(k, v)| format!("{}: {}", k, v)) - .join(", ") - ); - }); - let mapper = FlatProjectionMapper::new_with_read_columns( - metadata, - projection, - read_cols, - json_type_hint, - )?; + let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?; let mapper = if self.request.preserve_pk_dictionary_encoding { mapper.with_pk_dictionary_encoding() } else { @@ -609,7 +577,8 @@ impl ScanRegion { Ok(input) } - /// Builds the ordered root column ids required by the pushed-down projection and predicate. + /// Builds the deduplicated root column ids required by the projection and + /// predicate. fn build_read_col_ids( &self, projection: Option<&[usize]>, @@ -688,27 +657,50 @@ impl ScanRegion { Ok(read_col_ids) } - /// Builds read columns with nested paths derived from JSON type hints. - fn read_columns_with_json_type_hint(&self, col_ids: &[ColumnId]) -> ReadColumns { - let cols = col_ids + /// Builds logical read columns and attaches JSON2 target types when needed. + /// + /// The behavior about JSON2 is as follows: + /// - JSON2 columns without hints use Variant to read the whole column. + /// - Hints targeting non-JSON2 read columns are rejected. + fn build_read_columns(&self, col_ids: &[ColumnId]) -> Result { + let metadata = &self.version.metadata; + let json_type_hint = &self.request.json_type_hint; + + let has_json2 = metadata + .schema + .arrow_schema() + .fields() .iter() - .map(|&col_id| { - let nested_paths = self - .version - .metadata - .column_by_id(col_id) - .and_then(|column| { - let col_name = &column.column_schema.name; - self.request - .json_type_hint - .get(col_name) - .map(|json_type| json_nested_paths(col_name, json_type)) - }) - .unwrap_or_default(); - ReadColumn::new(col_id, nested_paths) - }) - .collect(); - ReadColumns { cols } + .any(is_json2_extension_type); + + if !has_json2 && json_type_hint.is_empty() { + return Ok(ReadColumns::new(col_ids.iter().copied())); + } + + let mut json_target_types = BTreeMap::new(); + for &col_id in col_ids { + let Some(col) = metadata.column_by_id(col_id) else { + continue; + }; + let col_name = &col.column_schema.name; + let hint = json_type_hint.get(col_name); + if !col.column_schema.data_type.is_json2() { + ensure!( + hint.is_none(), + InvalidRequestSnafu { + region_id: metadata.region_id, + reason: format!( + "JSON type hint targets non-JSON2 column {} (id: {}, type: {})", + col_name, col_id, col.column_schema.data_type + ), + } + ); + continue; + } + let target_type = hint.cloned().unwrap_or(JsonNativeType::Variant); + json_target_types.insert(col_id, target_type); + } + Ok(ReadColumns::new(col_ids.iter().copied()).with_json_target_types(json_target_types)) } fn region_id(&self) -> RegionId { @@ -1583,30 +1575,6 @@ fn pre_filter_mode(append_mode: bool, merge_mode: MergeMode) -> PreFilterMode { } } -fn json_nested_paths(column_name: &str, json_type: &JsonNativeType) -> Vec { - let mut paths = Vec::new(); - let mut current = vec![column_name.to_string()]; - collect_json_nested_paths(json_type, &mut current, &mut paths); - paths -} - -fn collect_json_nested_paths( - json_type: &JsonNativeType, - current: &mut NestedPath, - paths: &mut Vec, -) { - match json_type { - JsonNativeType::Object(fields) if !fields.is_empty() => { - for (field, child) in fields { - current.push(field.clone()); - collect_json_nested_paths(child, current, paths); - current.pop(); - } - } - _ => paths.push(current.clone()), - } -} - /// Output of [build_scan_fingerprint]: the cache fingerprint plus the derived /// implied time range used to decide whether the cache key can drop the time /// predicates for a given partition (see `build_range_cache_key`). @@ -1703,9 +1671,15 @@ pub(crate) fn build_scan_fingerprint(input: &ScanInput) -> Option Result<()> { - let hint = JsonNativeType::Object(JsonObjectType::from([ - ("a".to_string(), JsonNativeType::i64()), - ( - "b".to_string(), - JsonNativeType::Object(JsonObjectType::from([( - "c".to_string(), - JsonNativeType::String, - )])), - ), - ])); - - fn nested_path(parts: &[&str]) -> NestedPath { - parts.iter().map(|part| part.to_string()).collect() - } - - assert_eq!( - json_nested_paths("j", &hint), - vec![nested_path(&["j", "a"]), nested_path(&["j", "b", "c"])] - ); - Ok(()) - } - #[tokio::test] async fn test_build_scan_fingerprint_for_eligible_scan() { let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); @@ -2370,6 +2321,158 @@ mod tests { assert_ne!(0, metadata.partition_expr_version); } + #[tokio::test] + async fn test_build_scan_fingerprint_uses_json_target_types() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "k0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "j".to_string(), + ConcreteDataType::json2(JsonNativeType::Variant), + true, + ), + semantic_type: SemanticType::Field, + column_id: 2, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + + let make_input = |target_type| async { + let env = SchedulerEnv::new().await; + let read_cols = ReadColumns::new([0, 1, 2]) + .with_json_target_types(BTreeMap::from([(2, target_type)])); + let mapper = + FlatProjectionMapper::new_with_read_columns(&metadata, vec![0, 1, 2], read_cols) + .unwrap(); + let predicate = + PredicateGroup::new(metadata.as_ref(), &[col("k0").eq(lit("foo"))]).unwrap(); + let file = FileHandle::new( + FileMeta::default(), + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + ScanInput::new(env.access_layer.clone(), mapper) + .with_predicate(predicate) + .with_cache(CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(1024) + .build(), + ))) + .with_files(vec![file]) + }; + + let int_target = JsonNativeType::i64(); + let string_target = JsonNativeType::String; + let int_fingerprint = build_scan_fingerprint(&make_input(int_target.clone()).await) + .unwrap() + .fingerprint; + let string_fingerprint = build_scan_fingerprint(&make_input(string_target).await) + .unwrap() + .fingerprint; + + assert_ne!(int_fingerprint, string_fingerprint); + assert_eq!( + Some(&Some(ConcreteDataType::json2(int_target))), + int_fingerprint.read_column_types().get(2) + ); + } + + #[tokio::test] + async fn test_scan_input_rejects_json_type_hint_for_non_json2_column() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "k0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "j".to_string(), + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), + JsonNativeType::i64(), + )]))), + true, + ), + semantic_type: SemanticType::Field, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "v0".to_string(), + ConcreteDataType::int64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 3, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + metadata.clone(), + Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), + 0, + None, + )); + let version = Arc::new( + crate::region::version::VersionBuilder::new(metadata.clone(), mutable).build(), + ); + let env = SchedulerEnv::new().await; + let request = ScanRequest { + projection: Some(vec![0, 1, 2, 3]), + json_type_hint: std::collections::HashMap::from([( + "v0".to_string(), + JsonNativeType::i64(), + )]), + ..Default::default() + }; + + let err = ScanRegion::new( + version, + env.access_layer.clone(), + request, + CacheStrategy::Disabled, + ) + .scan_input() + .await; + let Err(err) = err else { + panic!("scan input should reject JSON type hint for non-JSON2 column"); + }; + + assert!(err.to_string().contains("non-JSON2 column v0")); + } + #[test] fn test_update_dyn_filters_with_empty_base_predicates() { let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index 2b9e800fb2..a3a94ef25d 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -235,6 +235,7 @@ impl FileRange { self.row_group_idx, cache_strategy, self.context.read_format().parquet_read_columns(), + self.context.read_format().json_target_types().clone(), flat_row_group_reader, ); FlatPruneReader::new_with_last_row_reader(self.context.clone(), reader, skip_fields) @@ -827,9 +828,7 @@ mod tests { let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -895,9 +894,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 237b7c0f2c..e381ea76f4 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -51,12 +51,11 @@ use crate::error::{ ComputeArrowSnafu, DecodeSnafu, InvalidParquetSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, }; -use crate::read::read_columns::ReadColumns; +use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; use crate::sst::parquet::format::{ FIXED_POS_COLUMN_NUM, FormatProjection, INTERNAL_COLUMN_NUM, PrimaryKeyArray, PrimaryKeyReadFormat, StatValues, column_null_counts, column_values, }; -use crate::sst::parquet::json_align::align_schema_by_nested_paths; use crate::sst::parquet::read_columns::ParquetReadColumns; use crate::sst::{ FlatSchemaOptions, flat_sst_arrow_schema_column_num, tag_maybe_to_dictionary_field, @@ -181,6 +180,8 @@ pub(crate) fn field_column_start(metadata: &RegionMetadata, num_columns: usize) pub struct FlatReadFormat { /// Sequence number to override the sequence read from the SST. override_sequence: Option, + /// Logical columns requested by this read. + read_cols: ReadColumns, /// Parquet format adapter. parquet_adapter: ParquetAdapter, /// Output schema to wrap binary `__primary_key` back to a dictionary; `None` disables wrapping. @@ -210,22 +211,25 @@ impl FlatReadFormat { // Only skip auto convert when the primary key encoding is sparse. ParquetAdapter::PrimaryKeyToFlat(ParquetPrimaryKeyToFlat::new( metadata, - read_cols, + read_cols.clone(), skip_auto_convert, )) } else { ParquetAdapter::PrimaryKeyToFlat(ParquetPrimaryKeyToFlat::new( - metadata, read_cols, false, + metadata, + read_cols.clone(), + false, )) } } else { let file_schema = file_schema .unwrap_or_else(|| to_flat_sst_arrow_schema(&metadata, &Default::default())); - ParquetAdapter::Flat(ParquetFlat::new(metadata, read_cols, file_schema)) + ParquetAdapter::Flat(ParquetFlat::new(metadata, read_cols.clone(), file_schema)) }; Ok(FlatReadFormat { override_sequence: None, + read_cols, parquet_adapter, pk_dict_wrap_schema: None, }) @@ -294,22 +298,51 @@ impl FlatReadFormat { } } - /// Gets the projected output schema produced by parquet reading. + /// Gets the projected output schema expected by the scan. pub(crate) fn output_arrow_schema(&self) -> Result { - let read_columns = self.parquet_read_columns(); - let projection = read_columns.root_indices(); + let projection = self.parquet_read_columns().root_indices(); let mut schema = self .arrow_schema() .project(projection) .context(ComputeArrowSnafu)?; - if read_columns.has_nested() { - debug_assert_eq!(schema.fields().len(), read_columns.columns().len()); - let nested_paths = read_columns.columns().iter().map(|x| x.nested_paths()); - align_schema_by_nested_paths(&mut schema, nested_paths); + let mut fields = schema.fields().iter().cloned().collect::>(); + for (column_id, target_type) in self.json_target_types().iter() { + let Some(index) = self.parquet_projected_index_by_id(*column_id) else { + continue; + }; + let Some(field) = schema.fields().get(index) else { + continue; + }; + fields[index] = Arc::new( + field + .as_ref() + .clone() + .with_data_type(ConcreteDataType::json2(target_type.clone()).as_arrow_type()), + ); } + schema.fields = fields.into(); Ok(Arc::new(schema)) } + /// Index of a column in the projected schema produced directly by parquet + /// reading, before any primary-key-to-flat conversion. + fn parquet_projected_index_by_id(&self, column_id: ColumnId) -> Option { + match &self.parquet_adapter { + ParquetAdapter::Flat(p) => p + .format_projection + .column_id_to_projected_index + .get(&column_id) + .copied(), + // `format_projection` addresses the post-conversion flat batch here. + // This helper needs the raw primary-key projection used by parquet reading. + ParquetAdapter::PrimaryKeyToFlat(p) => p + .format + .field_id_to_projected_index() + .get(&column_id) + .copied(), + } + } + /// Gets the metadata of the SST. pub(crate) fn metadata(&self) -> &RegionMetadataRef { match &self.parquet_adapter { @@ -326,6 +359,11 @@ impl FlatReadFormat { } } + /// Gets JSON2 target types keyed by column id. + pub(crate) fn json_target_types(&self) -> &JsonTargetTypes { + self.read_cols.json_target_types() + } + /// Gets the projection in the flat format. /// /// When `skip_auto_convert` is enabled (primary-key format read), this returns the @@ -491,6 +529,7 @@ impl ParquetPrimaryKeyToFlat { } else { // Computes the format projection for the new format. let format_projection = FormatProjection::compute_format_projection( + &metadata, &id_to_index, sst_column_num, read_cols.clone(), @@ -540,8 +579,12 @@ impl ParquetFlat { let id_to_index = sst_column_id_indices(&metadata); let sst_column_num = flat_sst_arrow_schema_column_num(&metadata, &FlatSchemaOptions::default()); - let format_projection = - FormatProjection::compute_format_projection(&id_to_index, sst_column_num, read_cols); + let format_projection = FormatProjection::compute_format_projection( + &metadata, + &id_to_index, + sst_column_num, + read_cols, + ); Self { metadata, @@ -841,9 +884,7 @@ impl FlatReadFormat { pub fn new_with_all_columns(metadata: RegionMetadataRef) -> FlatReadFormat { Self::new( Arc::clone(&metadata), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -965,7 +1006,7 @@ mod tests { .collect(); let mut read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), None, "test", false, @@ -1045,7 +1086,7 @@ mod tests { let metadata = Arc::new(build_metadata(1, 2, PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids([0_u32, 2_u32]), + ReadColumns::new([0_u32, 2_u32]), None, "test", false, diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index a6338c1a1c..f3f0a34a68 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -39,6 +39,7 @@ use datatypes::arrow::array::{ use datatypes::arrow::datatypes::{SchemaRef, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::DataType; +use datatypes::types::json_type::JsonNativeType; use datatypes::vectors::Helper; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields, @@ -47,12 +48,12 @@ use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; use parquet::file::statistics::Statistics; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::{ColumnMetadata, RegionMetadataRef}; -use store_api::storage::{ColumnId, SequenceNumber}; +use store_api::storage::{ColumnId, NestedPath, SequenceNumber}; use crate::error::{ ConvertVectorSnafu, DecodeSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, }; -use crate::read::read_columns::ReadColumns; +use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; use crate::read::{Batch, BatchBuilder, BatchColumn}; use crate::sst::file::{FileMeta, FileTimeRange}; use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns}; @@ -233,6 +234,7 @@ impl PrimaryKeyReadFormat { let arrow_schema = to_sst_arrow_schema(&metadata); let format_projection = FormatProjection::compute_format_projection( + &metadata, &field_id_to_index, arrow_schema.fields.len(), read_cols, @@ -592,18 +594,21 @@ impl FormatProjection { /// /// `id_to_index` is a mapping from column id to the index of the column in the SST. pub(crate) fn compute_format_projection( + metadata: &RegionMetadataRef, id_to_index: &HashMap, sst_column_num: usize, cols: ReadColumns, ) -> Self { + let json_target_types = cols.json_target_types().clone(); let mut projected_columns: Vec<_> = cols - .cols + .col_ids .into_iter() - .filter_map(|col| { - id_to_index - .get(&col.column_id) - .copied() - .map(|index_of_sst| (col.column_id, index_of_sst, col.nested_paths)) + .filter_map(|col_id| { + id_to_index.get(&col_id).copied().map(|index_of_sst| { + let nested_paths = + json_target_nested_paths(metadata, &json_target_types, col_id); + (col_id, index_of_sst, nested_paths) + }) }) .collect(); // Sorts columns by their indices in the SST. SST uses a bitmap for projection. @@ -681,6 +686,45 @@ impl FormatProjection { } } +fn json_target_nested_paths( + metadata: &RegionMetadataRef, + json_target_types: &JsonTargetTypes, + column_id: ColumnId, +) -> Vec { + let Some(target_type) = json_target_types.get(&column_id) else { + return Vec::new(); + }; + let Some(column) = metadata.column_by_id(column_id) else { + return Vec::new(); + }; + + json_nested_paths(&column.column_schema.name, target_type) +} + +fn json_nested_paths(column_name: &str, json_type: &JsonNativeType) -> Vec { + let mut paths = Vec::new(); + let mut current = vec![column_name.to_string()]; + collect_json_nested_paths(json_type, &mut current, &mut paths); + paths +} + +fn collect_json_nested_paths( + json_type: &JsonNativeType, + current: &mut NestedPath, + paths: &mut Vec, +) { + match json_type { + JsonNativeType::Object(fields) if !fields.is_empty() => { + for (field, child) in fields { + current.push(field.clone()); + collect_json_nested_paths(child, current, paths); + current.pop(); + } + } + _ => paths.push(current.clone()), + } +} + /// Values of column statistics of the SST. /// /// It also distinguishes the case that a column is not found and @@ -710,9 +754,7 @@ impl PrimaryKeyReadFormat { pub fn new_with_all_columns(metadata: RegionMetadataRef) -> PrimaryKeyReadFormat { Self::new( Arc::clone(&metadata), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), ) } } @@ -818,6 +860,7 @@ pub(crate) fn need_override_sequence(parquet_meta: &ParquetMetaData) -> bool { #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::sync::Arc; use api::v1::OpType; @@ -840,7 +883,6 @@ mod tests { use super::*; use crate::error::InvalidMetadataSnafu; - use crate::read::read_columns::ReadColumn; use crate::sst::parquet::flat_format::{ FlatReadFormat, FlatWriteFormat, sequence_column_index, sst_column_id_indices, }; @@ -986,29 +1028,25 @@ mod tests { fn test_projection_indices() { let metadata = build_test_region_metadata(); // Only read tag1 - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([3])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([3])); assert_eq!( &[2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Only read field1 - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([4])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([4])); assert_eq!( &[0, 2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Only read ts - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([5])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([5])); assert_eq!( &[2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Read field0, tag0, ts - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids([2, 1, 5])); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new([2, 1, 5])); assert_eq!( &[1, 2, 3, 4, 5], read_format.parquet_read_columns().root_indices() @@ -1049,14 +1087,16 @@ mod tests { let metadata = Arc::new(builder.build().context(InvalidMetadataSnafu)?); let column_id_to_parquet_index = sst_column_id_indices(&metadata); let projection = FormatProjection::compute_format_projection( + &metadata, &column_id_to_parquet_index, metadata.column_metadatas.len() + FIXED_POS_COLUMN_NUM, - ReadColumns { - cols: vec![ReadColumn::new( - 4, - vec![vec!["j".to_string(), "a".to_string()]], - )], - }, + ReadColumns::new([4]).with_json_target_types(BTreeMap::from([( + 4, + JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), + JsonNativeType::i64(), + )])), + )])), ); let columns = projection.parquet_read_cols.columns(); @@ -1111,8 +1151,7 @@ mod tests { .iter() .map(|col| col.column_id) .collect(); - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids(column_ids)); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new(column_ids)); assert_eq!(arrow_schema, *read_format.arrow_schema()); let record_batch = RecordBatch::new_empty(arrow_schema); @@ -1131,8 +1170,7 @@ mod tests { .iter() .map(|col| col.column_id) .collect(); - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids(column_ids)); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new(column_ids)); let columns: Vec = vec![ Arc::new(Int64Array::from(vec![1, 1, 10, 10])), // field1 @@ -1160,9 +1198,7 @@ mod tests { let metadata = build_test_region_metadata(); let read_format = PrimaryKeyReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), ); let columns: Vec = vec![ @@ -1330,56 +1366,36 @@ mod tests { // The projection includes all "fixed position" columns: ts(4), __primary_key(5), __sequence(6), __op_type(7) // Only read tag1 (column_id=3, index=1) + fixed columns - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([3]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([3]), None, "test", false) + .unwrap(); assert_eq!( &[1, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Only read field1 (column_id=4, index=2) + fixed columns - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([4]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([4]), None, "test", false) + .unwrap(); assert_eq!( &[2, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Only read ts (column_id=5, index=4) + fixed columns (ts is already included in fixed) - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([5]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([5]), None, "test", false) + .unwrap(); assert_eq!( &[4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Read field0(column_id=2, index=3), tag0(column_id=1, index=0), ts(column_id=5, index=4) + fixed columns - let read_format = FlatReadFormat::new( - metadata, - ReadColumns::from_deduped_column_ids([2, 1, 5]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata, ReadColumns::new([2, 1, 5]), None, "test", false) + .unwrap(); assert_eq!( &[0, 3, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() @@ -1391,7 +1407,7 @@ mod tests { let metadata = build_test_region_metadata(); let mut format = FlatReadFormat::new( metadata, - ReadColumns::from_deduped_column_ids(std::iter::once(1)), // Just read tag0 + ReadColumns::new(std::iter::once(1)), // Just read tag0 Some(build_test_flat_sst_schema()), "test", false, @@ -1608,7 +1624,7 @@ mod tests { .collect(); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), Some(build_test_arrow_schema()), "test", false, @@ -1674,7 +1690,7 @@ mod tests { .collect(); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids.clone()), + ReadColumns::new(column_ids.clone()), None, "test", false, @@ -1744,7 +1760,7 @@ mod tests { let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), None, "test", true, diff --git a/src/mito2/src/sst/parquet/json_align/mod.rs b/src/mito2/src/sst/parquet/json_align/mod.rs index 89b993d3d0..54c97bdbb8 100644 --- a/src/mito2/src/sst/parquet/json_align/mod.rs +++ b/src/mito2/src/sst/parquet/json_align/mod.rs @@ -15,10 +15,8 @@ use datatypes::arrow::record_batch::RecordBatch; use futures::stream::BoxStream; -mod schema; mod stream; -pub(crate) use schema::align_schema_by_nested_paths; pub(crate) use stream::NestedSchemaAligner; use crate::error::Result; diff --git a/src/mito2/src/sst/parquet/json_align/schema.rs b/src/mito2/src/sst/parquet/json_align/schema.rs deleted file mode 100644 index 86bd4d4d0d..0000000000 --- a/src/mito2/src/sst/parquet/json_align/schema.rs +++ /dev/null @@ -1,250 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::BTreeMap; -use std::sync::Arc; - -use arrow_schema::{DataType as ArrowDataType, FieldRef}; -use datatypes::arrow::datatypes::Schema; -use datatypes::extension::json::is_json2_extension_type; -use store_api::storage::NestedPath; - -/// Aligns nested struct fields according to the requested nested paths. -/// -/// For each root field: -/// - An empty path list keeps the whole field unchanged. -/// - Non-JSON root fields ignore nested paths and keep the whole field unchanged. -/// - JSON2 root fields are rebuilt from `nested_paths`. -/// - Existing schema fields are preserved only when they are the requested leaf. -/// - Requested paths missing from the schema are synthesized with JSONB (`Binary`) -/// leaves. -/// -/// For example, if the schema has `j: struct>` -/// and `nested_paths` requests `j.a.x` and `j.a.z`, the result is -/// `j: struct>`. -pub(crate) fn align_schema_by_nested_paths<'a, I>(schema: &mut Schema, nested_paths: I) -where - I: IntoIterator, -{ - let fields = schema - .fields - .into_iter() - .zip(nested_paths) - .map(|(field, paths)| { - if !paths.is_empty() && is_json2_extension_type(field) { - let child_paths = paths - .iter() - .map(|path| { - if path.first().is_some_and(|root| root == field.name()) { - &path[1..] - } else { - path - } - }) - .collect::>(); - rebuild_field_by_nested_paths(field, &child_paths) - } else { - field.clone() - } - }) - .collect::>(); - schema.fields = fields.into() -} - -fn rebuild_field_by_nested_paths(field: &FieldRef, nested_paths: &[&[String]]) -> FieldRef { - if nested_paths.iter().any(|path| path.is_empty()) { - return field.clone(); - }; - - let fields = group_child_paths(nested_paths) - .into_iter() - .map(|(name, paths)| { - let existing = find_struct_child(field, &name); - build_field_from_nested_paths(&name, existing, &paths) - }) - .collect::>(); - - Arc::new( - field - .as_ref() - .clone() - .with_data_type(ArrowDataType::Struct(fields.into())), - ) -} - -fn group_child_paths<'a>(nested_paths: &[&'a [String]]) -> BTreeMap> { - let mut child_paths = BTreeMap::>::new(); - for path in nested_paths { - let Some((name, remaining)) = path.split_first() else { - continue; - }; - child_paths.entry(name.clone()).or_default().push(remaining); - } - child_paths -} - -fn find_struct_child<'a>(field: &'a FieldRef, name: &str) -> Option<&'a FieldRef> { - let ArrowDataType::Struct(fields) = field.data_type() else { - return None; - }; - fields.iter().find(|field| field.name() == name) -} - -fn build_field_from_nested_paths( - name: &str, - existing: Option<&FieldRef>, - nested_paths: &[&[String]], -) -> FieldRef { - if nested_paths.iter().any(|path| path.is_empty()) { - return existing.cloned().unwrap_or_else(|| new_jsonb_field(name)); - } - - let fields = group_child_paths(nested_paths) - .into_iter() - .map(|(name, paths)| { - let existing_child = existing.and_then(|field| find_struct_child(field, &name)); - build_field_from_nested_paths(&name, existing_child, &paths) - }) - .collect::>(); - - let field = existing - .map(|field| field.as_ref().clone()) - .unwrap_or_else(|| arrow_schema::Field::new(name, ArrowDataType::Binary, true)); - Arc::new(field.with_data_type(ArrowDataType::Struct(fields.into()))) -} - -fn new_jsonb_field(name: &str) -> FieldRef { - Arc::new(arrow_schema::Field::new(name, ArrowDataType::Binary, true)) -} - -#[cfg(test)] -mod tests { - use arrow_schema::Field; - use datatypes::extension::json::Json2ExtensionType; - - use super::*; - - #[test] - fn test_align_schema_by_nested_paths() { - fn new_field(name: &str, data_type: ArrowDataType) -> FieldRef { - Arc::new(Field::new(name, data_type, true)) - } - - fn struct_field(name: &str, fields: impl IntoIterator) -> FieldRef { - new_field(name, ArrowDataType::Struct(fields.into_iter().collect())) - } - - fn json_struct_field(name: &str, fields: impl IntoIterator) -> FieldRef { - Arc::new( - Field::new( - name, - ArrowDataType::Struct(fields.into_iter().collect()), - true, - ) - .with_extension_type(Json2ExtensionType::default()), - ) - } - - let mut schema = Schema::new([ - json_struct_field( - "j", - [ - struct_field( - "a", - [ - new_field("x", ArrowDataType::Int64), - new_field("y", ArrowDataType::Utf8), - struct_field( - "z", - [ - new_field("q", ArrowDataType::Boolean), - new_field("r", ArrowDataType::Float64), - ], - ), - ], - ), - new_field("b", ArrowDataType::Utf8), - struct_field( - "c", - vec![ - new_field("d", ArrowDataType::Int64), - new_field("e", ArrowDataType::Utf8), - ], - ), - ], - ), - new_field("tag", ArrowDataType::Utf8), - struct_field( - "k", - [ - new_field("k_0", ArrowDataType::Int64), - new_field("k_1", ArrowDataType::Utf8), - ], - ), - ]); - - let nested_paths = [ - vec![ - ["j", "a", "x"].iter().map(|x| x.to_string()).collect(), - ["j", "a", "z", "q"].iter().map(|x| x.to_string()).collect(), - ["j", "a", "m"].iter().map(|x| x.to_string()).collect(), - ["j", "b", "x"].iter().map(|x| x.to_string()).collect(), - ["j", "c"].iter().map(|x| x.to_string()).collect(), - ["j", "d", "e"].iter().map(|x| x.to_string()).collect(), - ], - vec![["tag", "ignored"].iter().map(|x| x.to_string()).collect()], - vec![], - ]; - - align_schema_by_nested_paths( - &mut schema, - nested_paths.iter().map(|paths| paths.as_slice()), - ); - - let expected = Schema::new([ - json_struct_field( - "j", - [ - struct_field( - "a", - [ - new_field("m", ArrowDataType::Binary), - new_field("x", ArrowDataType::Int64), - struct_field("z", vec![new_field("q", ArrowDataType::Boolean)]), - ], - ), - struct_field("b", [new_field("x", ArrowDataType::Binary)]), - struct_field( - "c", - [ - new_field("d", ArrowDataType::Int64), - new_field("e", ArrowDataType::Utf8), - ], - ), - struct_field("d", [new_field("e", ArrowDataType::Binary)]), - ], - ), - new_field("tag", ArrowDataType::Utf8), - struct_field( - "k", - [ - new_field("k_0", ArrowDataType::Int64), - new_field("k_1", ArrowDataType::Utf8), - ], - ), - ]); - - assert_eq!(schema, expected); - } -} diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index 5b08a0e495..fcaaae4a54 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -1457,9 +1457,7 @@ mod tests { Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -1495,9 +1493,7 @@ mod tests { )); let legacy_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "memtable", false, @@ -1524,9 +1520,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let raw_pk_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "memtable", true, @@ -1560,9 +1554,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let full_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -1594,7 +1586,7 @@ mod tests { let ts = metric_metadata.time_index_column().column_id; let projected_read_format = FlatReadFormat::new( metric_metadata.clone(), - ReadColumns::from_deduped_column_ids([field_0, ts]), + ReadColumns::new([field_0, ts]), None, "test", true, @@ -1646,9 +1638,7 @@ mod tests { )); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -1686,9 +1676,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -1734,9 +1722,7 @@ mod tests { Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 97a4abfac0..4f890cb83e 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -451,7 +451,7 @@ impl ParquetReaderBuilder { } else { let expected_meta = self.expected_metadata.as_ref().unwrap_or(®ion_meta); // Lists all column ids to read, we always use the expected metadata if possible. - ReadColumns::from_deduped_column_ids( + ReadColumns::new( expected_meta .column_metadatas .iter() @@ -2472,9 +2472,7 @@ mod tests { let metadata = Arc::new(sst_region_metadata()); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -2679,7 +2677,7 @@ mod tests { let region_metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( region_metadata.clone(), - ReadColumns::from_deduped_column_ids( + ReadColumns::new( region_metadata .column_metadatas .iter() @@ -2946,9 +2944,7 @@ mod tests { let expected_metadata = expected_metadata_with_reused_tag_name(metadata.as_ref()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -2970,9 +2966,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, diff --git a/src/query/src/optimizer/json_type_concretize.rs b/src/query/src/optimizer/json_type_concretize.rs index 4210433eb3..4625b61978 100644 --- a/src/query/src/optimizer/json_type_concretize.rs +++ b/src/query/src/optimizer/json_type_concretize.rs @@ -16,12 +16,14 @@ use std::collections::HashMap; use arrow_schema::DataType; use common_function::scalars::json::json_get::JsonGetWithType; -use datafusion::datasource::DefaultTableSource; +use datafusion::datasource::{DefaultTableSource, TableProvider}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{Result, plan_datafusion_err, plan_err}; use datafusion_expr::{Expr, LogicalPlan}; use datafusion_optimizer::{OptimizerConfig, OptimizerRule}; +use datatypes::extension::json::is_json2_extension_type; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; +use table::table::adapter::DfTableProviderAdapter; use crate::dummy_catalog::DummyTableProvider; @@ -56,22 +58,56 @@ impl OptimizerRule for JsonTypeConcretizeRule { return Ok(Transformed::no(plan)); }; - let Some(adapter) = source - .table_provider - .as_any() - .downcast_ref::() - else { - return Ok(Transformed::no(plan)); - }; - - adapter.with_json_type_hint(json_types.clone()); - Ok(Transformed::yes(plan)) + if apply_json_type_hint(source.table_provider.as_ref(), &json_types) { + Ok(Transformed::yes(plan)) + } else { + Ok(Transformed::no(plan)) + } } _ => Ok(Transformed::no(plan)), }) } } +// FIXME: `json_types` is keyed only by unqualified column name. In joins with +// same-named JSON2 columns, a hint deduced from one scan can be applied to +// another scan. Carry the originating relation/scan when deducing hints. +/// Applies JSON type hints to providers that can carry scan request hints. +/// +/// Returns `true` if at least one JSON2 hint is retained and written to the provider. +fn apply_json_type_hint( + provider: &dyn TableProvider, + json_types: &HashMap, +) -> bool { + let schema = provider.schema(); + let json_types = json_types + .iter() + .filter(|(column, _)| { + schema + .fields() + .iter() + .any(|field| field.name() == *column && is_json2_extension_type(field)) + }) + .map(|(column, json_type)| (column.clone(), json_type.clone())) + .collect::>(); + + if json_types.is_empty() { + return false; + } + + if let Some(adapter) = provider.as_any().downcast_ref::() { + adapter.with_json_type_hint(json_types); + return true; + } + + if let Some(adapter) = provider.as_any().downcast_ref::() { + adapter.with_json_type_hint(json_types); + return true; + } + + false +} + fn deduce_json_types(plan: &LogicalPlan) -> Result> { let mut json_types = HashMap::::new(); @@ -228,11 +264,11 @@ mod tests { #[test] fn test_json_type_concretize_rule_rewrite() -> Result<()> { let exprs = vec![ - json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"), - json_get_expr(col("k0"), path_expr("a.c"), None)?.alias("ac"), - json_get_expr(col("k0"), path_expr("d"), Some(DataType::Boolean))?.alias("d"), + json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"), + json_get_expr(col("j"), path_expr("a.c"), None)?.alias("ac"), + json_get_expr(col("j"), path_expr("d"), Some(DataType::Boolean))?.alias("d"), ]; - let (provider, plan) = build_plan(exprs)?; + let (provider, plan) = build_json2_plan(exprs)?; assert!( JsonTypeConcretizeRule @@ -253,17 +289,17 @@ mod tests { let request = provider.scan_request(); assert_eq!(1, request.json_type_hint.len()); - assert_eq!(Some(&expected), request.json_type_hint.get("k0")); + assert_eq!(Some(&expected), request.json_type_hint.get("j")); Ok(()) } #[test] fn test_json_type_concretize_rule_conflict_to_variant() -> Result<()> { let exprs = vec![ - json_get_expr(col("k0"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"), - json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"), + json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"), + json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"), ]; - let (provider, plan) = build_plan(exprs)?; + let (provider, plan) = build_json2_plan(exprs)?; assert!( JsonTypeConcretizeRule @@ -277,11 +313,26 @@ mod tests { )])); assert_eq!( Some(&expected), - provider.scan_request().json_type_hint.get("k0") + provider.scan_request().json_type_hint.get("j") ); Ok(()) } + #[test] + fn test_json_type_concretize_rule_ignores_non_json2_columns() -> Result<()> { + let exprs = + vec![json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab")]; + let (provider, plan) = build_plan(exprs)?; + + assert!( + !JsonTypeConcretizeRule + .rewrite(plan, &OptimizerContext::default())? + .transformed + ); + assert!(provider.scan_request().json_type_hint.is_empty()); + Ok(()) + } + #[test] fn test_json_type_concretize_rule_no_json_get() -> Result<()> { let (provider, plan) = build_plan(vec![col("k0"), col("v0")])?; diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index 0ae56ce1c3..ac71b448e9 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -28,6 +28,7 @@ use datafusion_expr::TableProviderFilterPushDown as DfTableProviderFilterPushDow use datafusion_expr::expr::Expr; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; +use datatypes::types::json_type::JsonNativeType; use store_api::storage::{ScanRequest, VectorSearchRequest}; use crate::table::{TableRef, TableType}; @@ -63,6 +64,10 @@ impl DfTableProviderAdapter { self.scan_req.lock().unwrap().vector_search = Some(hint); } + pub fn with_json_type_hint(&self, hint: std::collections::HashMap) { + self.scan_req.lock().unwrap().json_type_hint = hint; + } + pub fn get_vector_search_hint(&self) -> Option { self.scan_req.lock().unwrap().vector_search.clone() } diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 3cb462a8a7..035c4ed07b 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -77,6 +77,75 @@ drop table json2_whole_and_path_read; Affected Rows: 0 +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +Affected Rows: 1 + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +Affected Rows: 1 + +admin flush_table('json2_join_same_name_left'); + ++------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_left') | ++------------------------------------------------+ +| 0 | ++------------------------------------------------+ + +admin flush_table('json2_join_same_name_right'); + ++-------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_right') | ++-------------------------------------------------+ +| 0 | ++-------------------------------------------------+ + +-- FIXME: This should return `right` and `1`. The current NULL values are caused +-- by JSON type hints losing the table qualifier in joins. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + ++---------------------------------------------+---------------------------------------------------+ +| json_get(r.j,Utf8("a")) | arrow_cast(json_get(l.j,Utf8("a")),Utf8("Int64")) | ++---------------------------------------------+---------------------------------------------------+ +| {"a":"right","right_only":"should be kept"} | | ++---------------------------------------------+---------------------------------------------------+ + +drop table json2_join_same_name_left; + +Affected Rows: 0 + +drop table json2_join_same_name_right; + +Affected Rows: 0 + create table json2_without_append_mode ( ts timestamp time index, j json2 diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 4917b9690e..986ed58eb2 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -44,6 +44,45 @@ select j from json2_whole_and_path_read where j.a.b = 1; drop table json2_whole_and_path_read; +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +admin flush_table('json2_join_same_name_left'); + +admin flush_table('json2_join_same_name_right'); + +-- FIXME: This should return `right` and `1`. The current NULL values are caused +-- by JSON type hints losing the table qualifier in joins. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + +drop table json2_join_same_name_left; + +drop table json2_join_same_name_right; + create table json2_without_append_mode ( ts timestamp time index, j json2