feat: projeciton not work

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![],

src/mito2/src/read/compat.rs

@@ -606,7 +606,7 @@ fn may_compat_fields(
     actual: &RegionMetadata,
 ) -> Result<Option<CompatFields>> {
     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);
     }

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<ColumnId>,
+    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);
         }
 

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<ColumnId>,
+    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<PrimaryKeyProjectionMapper> {
         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<usize>,
+        projection_input: ProjectionInput,
         read_column_ids: Vec<ColumnId>,
     ) -> Result<PrimaryKeyProjectionMapper> {
+        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<Item = ColumnId> + '_ {
+        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!(

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<ColumnId>,
+    pub(crate) read_columns: ReadColumns,
     pub(crate) read_column_types: Vec<Option<ConcreteDataType>>,
     pub(crate) filters: Vec<String>,
     pub(crate) time_filters: Vec<String>,
@@ -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<ColumnId>,
+    /// 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<Option<ConcreteDataType>>,
@@ -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::<SharedScanRequestFingerprint>()
-            + self.inner.read_column_ids.capacity() * mem::size_of::<ColumnId>()
+            + self.inner.read_columns.columns().len()
+                * mem::size_of::<crate::read::read_columns::ReadColumn>()
             + self.inner.read_column_types.capacity() * mem::size_of::<Option<ConcreteDataType>>()
             + self.inner.filters.capacity() * mem::size_of::<String>()
             + 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());

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<ReadColumn>,
 }
@@ -80,12 +80,16 @@ impl ReadColumns {
         self.cols.iter().map(|column| column.column_id())
     }
 
+    pub fn column_ids(&self) -> Vec<ColumnId> {
+        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.