diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 45f8399c3ad2..7a06724bd3a8 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -1180,6 +1180,7 @@ mod tests { use crate::read::range_cache::{ RangeScanCacheKey, RangeScanCacheValue, ScanRequestFingerprintBuilder, }; + use crate::read::read_columns::ReadColumns; use crate::sst::parquet::row_selection::RowGroupSelection; #[tokio::test] @@ -1353,7 +1354,7 @@ mod tests { region_id: RegionId::new(1, 1), row_groups: vec![(FileId::random(), 0)], scan: ScanRequestFingerprintBuilder { - read_column_ids: vec![], + read_columns: ReadColumns::from_column_ids(std::iter::empty()), read_column_types: vec![], filters: vec!["tag_0 = 1".to_string()], time_filters: vec![], diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index 75c60be5fa0d..92fd2674afd1 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -606,7 +606,7 @@ fn may_compat_fields( actual: &RegionMetadata, ) -> Result> { let expect_fields = mapper.batch_fields(); - let actual_fields = Batch::projected_fields(actual, mapper.column_ids()); + let actual_fields = Batch::projected_fields(actual, &mapper.read_columns().column_ids()); if expect_fields == actual_fields { return Ok(None); } diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index 5eb8e6c64783..787b1c601bd0 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -20,8 +20,10 @@ use api::v1::SemanticType; use common_error::ext::BoxedError; use common_recordbatch::error::{ArrowComputeSnafu, ExternalSnafu, NewDfRecordBatchSnafu}; use common_recordbatch::{DfRecordBatch, RecordBatch}; +use datafusion_common::cast_column; use datatypes::arrow::array::Array; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; +use datatypes::compute::CastOptions; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::schema::{Schema, SchemaRef}; use datatypes::value::Value; @@ -33,7 +35,7 @@ use store_api::storage::{ColumnId, ProjectionInput}; use crate::cache::CacheStrategy; use crate::error::{InvalidRequestSnafu, RecordBatchSnafu, Result}; use crate::read::projection::{read_column_ids_from_projection, repeated_vector_with_cache}; -use crate::read::read_columns::build_read_columns; +use crate::read::read_columns::{ReadColumns, build_read_columns}; use crate::sst::parquet::flat_format::sst_column_id_indices; use crate::sst::parquet::format::FormatProjection; use crate::sst::{ @@ -49,11 +51,11 @@ pub struct FlatProjectionMapper { metadata: RegionMetadataRef, /// Schema for converted [RecordBatch] to return. output_schema: SchemaRef, - /// Ids of columns to read from memtables and SSTs. + /// The columns to read from memtables and SSTs. /// The mapper won't deduplicate the column ids. /// /// Note that this doesn't contain the `__table_id` and `__tsid`. - read_column_ids: Vec, + read_cols: ReadColumns, /// Ids and DataTypes of columns of the expected batch. /// We can use this to check if the batch is compatible with the expected schema. /// @@ -121,7 +123,7 @@ impl FlatProjectionMapper { &id_to_index, // All columns with internal columns. metadata.column_metadatas.len() + 3, - read_cols, + &read_cols, ); let batch_schema = flat_projected_columns(metadata, &format_projection); @@ -168,7 +170,7 @@ impl FlatProjectionMapper { Ok(FlatProjectionMapper { metadata: metadata.clone(), output_schema, - read_column_ids, + read_cols, batch_schema, is_empty_projection, batch_indices, @@ -186,10 +188,8 @@ impl FlatProjectionMapper { &self.metadata } - /// Returns ids of projected columns that we need to read - /// from memtables and SSTs. - pub(crate) fn column_ids(&self) -> &[ColumnId] { - &self.read_column_ids + pub(crate) fn read_columns(&self) -> &ReadColumns { + &self.read_cols } /// Returns the field column start index in output batch. @@ -295,6 +295,27 @@ impl FlatProjectionMapper { array = casted; } } + // TODO(fys): remove this after we fix the schema mismatch issue. + else { + let target_schema = self.output_schema.arrow_schema(); + if *target_schema != batch.schema() { + common_telemetry::info!( + "Casting column {} from type {} to {} since the batch schema doesn't match the output schema", + self.output_schema.column_schemas()[output_idx].name, + batch.schema().field(*index).data_type(), + target_schema.field(output_idx).data_type() + ); + let source_col = batch.column(*index).clone(); + let target_field = self.output_schema.arrow_schema().field(output_idx); + array = + cast_column(&source_col, target_field, &CastOptions::default()).unwrap(); + } else { + common_telemetry::info!( + "Skip casting column {} since the batch schema matches the output schema", + self.output_schema.column_schemas()[output_idx].name + ); + } + } arrays.push(array); } diff --git a/src/mito2/src/read/projection.rs b/src/mito2/src/read/projection.rs index ec1321ef8b9b..3bdb3789a046 100644 --- a/src/mito2/src/read/projection.rs +++ b/src/mito2/src/read/projection.rs @@ -35,6 +35,7 @@ use crate::cache::CacheStrategy; use crate::error::{InvalidRequestSnafu, Result}; use crate::read::Batch; use crate::read::flat_projection::FlatProjectionMapper; +use crate::read::read_columns::{ReadColumns, build_read_columns}; /// Only cache vector when its length `<=` this value. pub(crate) const MAX_VECTOR_LENGTH_TO_CACHE: usize = 16384; @@ -94,12 +95,10 @@ impl ProjectionMapper { } } - /// Returns ids of projected columns that we need to read - /// from memtables and SSTs. - pub(crate) fn column_ids(&self) -> &[ColumnId] { + pub(crate) fn read_columns(&self) -> &ReadColumns { match self { - ProjectionMapper::PrimaryKey(m) => m.column_ids(), - ProjectionMapper::Flat(m) => m.column_ids(), + ProjectionMapper::PrimaryKey(m) => m.read_columns(), + ProjectionMapper::Flat(m) => m.read_columns(), } } @@ -151,7 +150,7 @@ pub struct PrimaryKeyProjectionMapper { /// Schema for converted [RecordBatch]. output_schema: SchemaRef, /// Ids of columns to read from memtables and SSTs. - read_column_ids: Vec, + read_cols: ReadColumns, /// Ids and DataTypes of field columns in the read [Batch]. batch_fields: Vec<(ColumnId, ConcreteDataType)>, /// `true` If the original projection is empty. @@ -170,15 +169,17 @@ impl PrimaryKeyProjectionMapper { ) -> Result { let projection: Vec<_> = projection.collect(); let read_column_ids = read_column_ids_from_projection(metadata, &projection)?; - Self::new_with_read_columns(metadata, projection, read_column_ids) + let projection_input = ProjectionInput::new().with_projection(projection); + Self::new_with_read_columns(metadata, projection_input, read_column_ids) } /// Returns a new mapper with output projection and explicit read columns. pub fn new_with_read_columns( metadata: &RegionMetadataRef, - projection: Vec, + projection_input: ProjectionInput, read_column_ids: Vec, ) -> Result { + let projection = projection_input.projection; // If the original projection is empty. let is_empty_projection = projection.is_empty(); @@ -199,6 +200,8 @@ impl PrimaryKeyProjectionMapper { } let codec = build_primary_key_codec(metadata); + let read_cols = + build_read_columns(metadata, &projection_input.nested_paths, &read_column_ids)?; // If projection is empty, we don't output any column. let output_schema = if is_empty_projection { Arc::new(Schema::new(vec![])) @@ -257,7 +260,7 @@ impl PrimaryKeyProjectionMapper { has_tags, codec, output_schema, - read_column_ids, + read_cols, batch_fields, is_empty_projection, }) @@ -280,8 +283,12 @@ impl PrimaryKeyProjectionMapper { /// Returns ids of projected columns that we need to read /// from memtables and SSTs. - pub(crate) fn column_ids(&self) -> &[ColumnId] { - &self.read_column_ids + pub(crate) fn column_ids_iter(&self) -> impl Iterator + '_ { + self.read_cols.column_ids_iter() + } + + pub(crate) fn read_columns(&self) -> &ReadColumns { + &self.read_cols } /// Returns ids of fields in [Batch]es the mapper expects to convert. @@ -585,7 +592,7 @@ mod tests { ); let cache = CacheStrategy::Disabled; let mapper = ProjectionMapper::all(&metadata).unwrap(); - assert_eq!([0, 1, 2, 3, 4], mapper.column_ids()); + assert_eq!(vec![0, 1, 2, 3, 4], mapper.read_columns().column_ids()); assert_eq!( [ (1, ConcreteDataType::int64_datatype()), @@ -621,7 +628,7 @@ mod tests { let cache = CacheStrategy::Disabled; // Columns v1, k0 let mapper = ProjectionMapper::new(&metadata, [4, 1].into_iter()).unwrap(); - assert_eq!([4, 1], mapper.column_ids()); + assert_eq!(vec![4, 1], mapper.read_columns().column_ids()); assert_eq!( [ (1, ConcreteDataType::int64_datatype()), @@ -658,7 +665,7 @@ mod tests { let mapper = ProjectionMapper::new_with_read_columns(&metadata, projection_input, vec![4, 1, 3]) .unwrap(); - assert_eq!([4, 1, 3], mapper.column_ids()); + assert_eq!(vec![4, 1, 3], mapper.read_columns().column_ids()); let batch = new_flat_batch(None, &[(1, 1)], &[(3, 3), (4, 4)], 3); let record_batch = mapper.as_flat().unwrap().convert(&batch, &cache).unwrap(); @@ -684,7 +691,7 @@ mod tests { let cache = CacheStrategy::Disabled; // Empty projection let mapper = ProjectionMapper::new(&metadata, [].into_iter()).unwrap(); - assert_eq!([0], mapper.column_ids()); // Should still read the time index column + assert_eq!(vec![0], mapper.read_columns().column_ids()); // Should still read the time index column assert!(mapper.output_schema().is_empty()); let flat_mapper = mapper.as_flat().unwrap(); assert_eq!( diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 2431a21f6a05..a3edd13bf832 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -25,10 +25,11 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use futures::TryStreamExt; use store_api::region_engine::PartitionRange; -use store_api::storage::{ColumnId, FileId, RegionId, TimeSeriesRowSelector}; +use store_api::storage::{FileId, RegionId, TimeSeriesRowSelector}; use crate::cache::CacheStrategy; use crate::read::BoxedRecordBatchStream; +use crate::read::read_columns::ReadColumns; use crate::read::scan_region::StreamContext; use crate::read::scan_util::PartitionMetrics; use crate::region::options::MergeMode; @@ -58,7 +59,7 @@ pub(crate) struct ScanRequestFingerprint { #[derive(Debug)] pub(crate) struct ScanRequestFingerprintBuilder { - pub(crate) read_column_ids: Vec, + pub(crate) read_columns: ReadColumns, pub(crate) read_column_types: Vec>, pub(crate) filters: Vec, pub(crate) time_filters: Vec, @@ -72,7 +73,7 @@ pub(crate) struct ScanRequestFingerprintBuilder { impl ScanRequestFingerprintBuilder { pub(crate) fn build(self) -> ScanRequestFingerprint { let Self { - read_column_ids, + read_columns, read_column_types, filters, time_filters, @@ -85,7 +86,7 @@ impl ScanRequestFingerprintBuilder { ScanRequestFingerprint { inner: Arc::new(SharedScanRequestFingerprint { - read_column_ids, + read_columns, read_column_types, filters, }), @@ -102,8 +103,8 @@ impl ScanRequestFingerprintBuilder { /// Non-copiable struct of the fingerprint. #[derive(Debug, PartialEq, Eq, Hash)] struct SharedScanRequestFingerprint { - /// Column ids of the projection. - read_column_ids: Vec, + /// Logical columns of the projection. + read_columns: ReadColumns, /// Column types of the projection. /// We keep this to ensure we won't reuse the fingerprint after a schema change. read_column_types: Vec>, @@ -113,8 +114,8 @@ struct SharedScanRequestFingerprint { impl ScanRequestFingerprint { #[cfg(test)] - pub(crate) fn read_column_ids(&self) -> &[ColumnId] { - &self.inner.read_column_ids + pub(crate) fn read_columns(&self) -> &ReadColumns { + &self.inner.read_columns } #[cfg(test)] @@ -149,7 +150,8 @@ impl ScanRequestFingerprint { pub(crate) fn estimated_size(&self) -> usize { mem::size_of::() - + self.inner.read_column_ids.capacity() * mem::size_of::() + + self.inner.read_columns.columns().len() + * mem::size_of::() + self.inner.read_column_types.capacity() * mem::size_of::>() + self.inner.filters.capacity() * mem::size_of::() + self @@ -453,7 +455,7 @@ pub fn bench_cache_flat_range_stream( let cache_strategy = CacheStrategy::EnableAll(cache_manager); let fingerprint = ScanRequestFingerprintBuilder { - read_column_ids: vec![], + read_columns: ReadColumns::from_column_ids(std::iter::empty()), read_column_types: vec![], filters: vec![], time_filters: vec![], @@ -644,7 +646,7 @@ mod tests { #[test] fn normalizes_and_clears_time_filters() { let normalized = ScanRequestFingerprintBuilder { - read_column_ids: vec![1, 2], + read_columns: ReadColumns::from_column_ids(vec![1, 2]), read_column_types: vec![None, None], filters: vec!["k0 = 'foo'".to_string()], time_filters: vec![], @@ -659,7 +661,7 @@ mod tests { assert!(normalized.time_filters().is_empty()); let fingerprint = ScanRequestFingerprintBuilder { - read_column_ids: vec![1, 2], + read_columns: ReadColumns::from_column_ids(vec![1, 2]), read_column_types: vec![None, None], filters: vec!["k0 = 'foo'".to_string()], time_filters: vec!["ts >= 1000".to_string()], @@ -673,7 +675,7 @@ mod tests { let reset = fingerprint.without_time_filters(); - assert_eq!(reset.read_column_ids(), fingerprint.read_column_ids()); + assert_eq!(reset.read_columns(), fingerprint.read_columns()); assert_eq!(reset.read_column_types(), fingerprint.read_column_types()); assert_eq!(reset.filters(), fingerprint.filters()); assert!(reset.time_filters().is_empty()); diff --git a/src/mito2/src/read/read_columns.rs b/src/mito2/src/read/read_columns.rs index c304daff87d9..a5c206656a29 100644 --- a/src/mito2/src/read/read_columns.rs +++ b/src/mito2/src/read/read_columns.rs @@ -55,7 +55,7 @@ use crate::error::{InvalidRequestSnafu, Result}; /// ``` /// /// If `nested_paths` is empty, the whole column will be read. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ReadColumns { cols: Vec, } @@ -80,12 +80,16 @@ impl ReadColumns { self.cols.iter().map(|column| column.column_id()) } + pub fn column_ids(&self) -> Vec { + self.column_ids_iter().collect() + } + pub fn columns(&self) -> &[ReadColumn] { &self.cols } } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ReadColumn { column_id: ColumnId, /// Nested filed paths under this column. diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 1b84c0f5291e..6f97694d5e60 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -56,6 +56,7 @@ use crate::read::compat::{self, CompatBatch, FlatCompatBatch, PrimaryKeyCompatBa use crate::read::projection::ProjectionMapper; use crate::read::range::{FileRangeBuilder, MemRangeBuilder, RangeMeta, RowGroupIndex}; use crate::read::range_cache::ScanRequestFingerprint; +use crate::read::read_columns::ReadColumns; use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; use crate::read::stream::ScanBatchStream; @@ -391,7 +392,26 @@ impl ScanRegion { let predicate = PredicateGroup::new(&self.version.metadata, &self.request.filters)?; let read_column_ids = match self.request.projection_indices() { - Some(p) => self.build_read_column_ids(p, &predicate)?, + Some(p) => + // FIXME(fys): `build_read_column_ids()` only adds filter-required root + // columns, but nested pruning later is driven by + // `projection_input.nested_paths` when we build `ReadColumns`. + // + // This can under-read nested data if the projection and predicate + // reference different subfields under the same root column. + // + // Example: + // SELECT j.a FROM t WHERE j.b = 1; + // + // In this case we add root column `j` because the predicate uses it, + // but if nested paths only contain `["j", "a"]`, the SST reader may + // prune `j` down to `j.a` and then evaluate `j.b = 1` on incomplete + // data. We need to either merge predicate-required nested paths into + // `ReadColumns`, or disable nested pruning for predicate-referenced + // root columns. + { + self.build_read_column_ids(p, &predicate)? + } None => self .version .metadata @@ -782,10 +802,10 @@ pub struct ScanInput { access_layer: AccessLayerRef, /// Maps projected Batches to RecordBatches. pub(crate) mapper: Arc, - /// Column ids to read from memtables and SSTs. + /// The columns to read from memtables and SSTs. /// Notice this is different from the columns in `mapper` which are projected columns. /// But this read columns might also include non-projected columns needed for filtering. - pub(crate) read_column_ids: Vec, + pub(crate) read_cols: ReadColumns, /// Time range filter for time index. pub(crate) time_range: Option, /// Predicate to push down. @@ -838,7 +858,7 @@ impl ScanInput { pub(crate) fn new(access_layer: AccessLayerRef, mapper: ProjectionMapper) -> ScanInput { ScanInput { access_layer, - read_column_ids: mapper.column_ids().to_vec(), + read_cols: mapper.read_columns().clone(), mapper: Arc::new(mapper), time_range: None, predicate: PredicateGroup::default(), @@ -1110,7 +1130,7 @@ impl ScanInput { .access_layer .read_sst(file.clone()) .predicate(predicate) - .projection(Some(self.read_column_ids.clone())) + .projection(Some(self.read_cols.clone())) .cache(self.cache_strategy.clone()) .inverted_index_appliers(self.inverted_index_appliers.clone()) .bloom_filter_index_appliers(self.bloom_filter_index_appliers.clone()) @@ -1450,12 +1470,11 @@ pub(crate) fn build_scan_fingerprint(input: &ScanInput) -> Option Option, sst_column_num: usize, - column_ids: ReadColumns, + column_ids: &ReadColumns, ) -> Self { let projected_columns = Self::collect_projected_columns(id_to_index, column_ids); @@ -833,7 +833,7 @@ impl FormatProjection { fn collect_projected_columns( id_to_index: &HashMap, - column_ids: ReadColumns, + column_ids: &ReadColumns, ) -> Vec<(ColumnId, usize, Vec>)> { let mut projected_columns: Vec<_> = column_ids .columns() diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 9cbcf69b3f64..9de1d909355d 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -122,11 +122,11 @@ pub struct ParquetReaderBuilder { object_store: ObjectStore, /// Predicate to push down. predicate: Option, - /// Metadata of columns to read. + /// The columns to read. /// /// `None` reads all columns. Due to schema change, the projection /// can contain columns not in the parquet file. - projection: Option>, + read_cols: Option, /// Strategy to cache SST data. cache_strategy: CacheStrategy, /// Index appliers. @@ -166,7 +166,7 @@ impl ParquetReaderBuilder { file_handle, object_store, predicate: None, - projection: None, + read_cols: None, cache_strategy: CacheStrategy::Disabled, inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], @@ -194,8 +194,8 @@ impl ParquetReaderBuilder { /// /// The reader only applies the projection to fields. #[must_use] - pub fn projection(mut self, projection: Option>) -> ParquetReaderBuilder { - self.projection = projection; + pub fn projection(mut self, read_cols: Option) -> ParquetReaderBuilder { + self.read_cols = read_cols; self } @@ -373,15 +373,17 @@ impl ParquetReaderBuilder { }; let expected_meta = self.expected_metadata.as_ref().unwrap_or(®ion_meta); - let column_ids: Vec<_> = self.projection.as_ref().cloned().unwrap_or_else(|| { + let read_cols = if let Some(read_cols) = &self.read_cols { + read_cols.clone() + } else { // Lists all column ids to read, we always use the expected metadata if possible. - expected_meta - .column_metadatas - .iter() - .map(|col| col.column_id) - .collect() - }); - let read_cols = ReadColumns::from_column_ids(column_ids); + ReadColumns::from_column_ids( + expected_meta + .column_metadatas + .iter() + .map(|col| col.column_id), + ) + }; let mut read_format = ReadFormat::new( region_meta.clone(), Some(read_cols),