From 3e67c67607c857b24df3d168398cc78e06e5d16b Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:04:19 +0800 Subject: [PATCH] refactor: simplify scan projection to root column indices (#8629) * refactor: simplify scan projection to root column indices Signed-off-by: fys * fix: sqlness test Signed-off-by: fys * chore: code style adjust Signed-off-by: fys * fix: cr Signed-off-by: fys * add comment and test Signed-off-by: fys * minor change Signed-off-by: fys * fix: unit test Signed-off-by: fys --------- Signed-off-by: fys --- src/catalog/src/system_schema.rs | 8 +- .../information_schema/region_info.rs | 4 +- .../system_schema/information_schema/ssts.rs | 4 +- src/cmd/src/datanode/scanbench.rs | 7 +- src/file-engine/src/query.rs | 2 +- src/metric-engine/src/engine/read.rs | 29 +- src/metric-engine/src/metadata_region.rs | 10 +- src/mito2/src/engine/partition_filter_test.rs | 8 +- src/mito2/src/engine/projection_test.rs | 8 +- src/mito2/src/engine/scan_test.rs | 79 +++- src/mito2/src/read/read_columns.rs | 429 +----------------- src/mito2/src/read/scan_region.rs | 290 ++++++------ src/query/src/dist_plan/analyzer/test.rs | 2 +- src/query/src/dummy_catalog.rs | 2 +- src/store-api/src/region_info.rs | 6 +- src/store-api/src/sst_entry.rs | 10 +- src/store-api/src/storage.rs | 2 +- src/store-api/src/storage/projection.rs | 57 --- src/store-api/src/storage/requests.rs | 41 +- src/table/src/table/adapter.rs | 3 +- src/table/src/table/numbers.rs | 2 +- src/table/src/test_util/memtable.rs | 5 +- .../information_schema/ssts_limit.result | 40 +- 23 files changed, 289 insertions(+), 759 deletions(-) diff --git a/src/catalog/src/system_schema.rs b/src/catalog/src/system_schema.rs index fe926d893c..6582704a88 100644 --- a/src/catalog/src/system_schema.rs +++ b/src/catalog/src/system_schema.rs @@ -144,12 +144,8 @@ impl DataSource for SystemTableDataSource { &self, request: ScanRequest, ) -> std::result::Result { - let projection = request - .projection_input - .as_ref() - .map(|input| input.projection.clone()); - - let projected_schema = match projection.as_ref() { + let projection = request.projection.clone(); + let projected_schema = match request.projection.as_ref() { Some(projection) => self.try_project(projection)?, None => self.table.schema(), }; diff --git a/src/catalog/src/system_schema/information_schema/region_info.rs b/src/catalog/src/system_schema/information_schema/region_info.rs index 56dc04b85c..3fad1c8440 100644 --- a/src/catalog/src/system_schema/information_schema/region_info.rs +++ b/src/catalog/src/system_schema/information_schema/region_info.rs @@ -60,7 +60,7 @@ impl InformationTable for InformationSchemaRegionInfo { } fn to_stream(&self, request: ScanRequest) -> Result { - let schema = if let Some(p) = request.projection_indices() { + let schema = if let Some(p) = request.projection.as_deref() { Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) } else { self.schema.clone() @@ -86,7 +86,7 @@ impl InformationTable for InformationSchemaRegionInfo { } fn scan_plan(&self, request: ScanRequest) -> Result>> { - let schema = if let Some(p) = request.projection_indices() { + let schema = if let Some(p) = request.projection.as_deref() { Arc::new(self.schema.try_project(p).context(ProjectSchemaSnafu)?) } else { self.schema.clone() diff --git a/src/catalog/src/system_schema/information_schema/ssts.rs b/src/catalog/src/system_schema/information_schema/ssts.rs index 19c44aee34..da5fcd2eed 100644 --- a/src/catalog/src/system_schema/information_schema/ssts.rs +++ b/src/catalog/src/system_schema/information_schema/ssts.rs @@ -36,7 +36,7 @@ use crate::information_schema::{ use crate::system_schema::utils; fn projected_schema(schema: &SchemaRef, request: &ScanRequest) -> Result { - if let Some(p) = request.projection_indices() { + if let Some(p) = request.projection.as_deref() { Ok(Arc::new(schema.try_project(p).context(ProjectSchemaSnafu)?)) } else { Ok(schema.clone()) @@ -331,7 +331,7 @@ mod tests { .unwrap(); let request = ScanRequest { - projection_input: Some(vec![0].into()), + projection: Some(vec![0]), ..Default::default() }; let plan = table.scan_to_plan(request).unwrap().unwrap(); diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index b2a715ad31..4909467f48 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -53,9 +53,7 @@ use store_api::metadata::RegionMetadata; use store_api::path_utils::WAL_DIR; use store_api::region_engine::{PrepareRequest, QueryScanContext, RegionEngine}; use store_api::region_request::{PathType, RegionOpenRequest, RegionRequest}; -use store_api::storage::{ - ProjectionInput, RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector, -}; +use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; use tokio::fs; use crate::datanode::objbench::{build_object_store, parse_config}; @@ -618,10 +616,9 @@ impl ScanbenchCommand { let mut total_rows_all = 0u64; let mut total_elapsed_all = std::time::Duration::ZERO; - let projection_input = projection.map(ProjectionInput::new); for iteration in 0..self.iterations { let request = ScanRequest { - projection_input: projection_input.clone(), + projection: projection.clone(), filters: filters.clone(), series_row_selector, distribution, diff --git a/src/file-engine/src/query.rs b/src/file-engine/src/query.rs index 0ac563606a..62ca3c7b66 100644 --- a/src/file-engine/src/query.rs +++ b/src/file-engine/src/query.rs @@ -44,7 +44,7 @@ impl FileRegion { pub fn query(&self, request: ScanRequest) -> Result { let store = build_backend(&self.url, &self.options).context(BuildBackendSnafu)?; - let projection = request.projection_indices(); + let projection = request.projection.as_deref(); let file_projection = self.projection_pushdown_to_file(projection)?; let file_filters = self.filters_pushdown_to_file(&request.filters)?; let file_schema = Arc::new(Schema::new(self.file_options.file_column_schemas.clone())); diff --git a/src/metric-engine/src/engine/read.rs b/src/metric-engine/src/engine/read.rs index d714ed4a23..53dd086d22 100644 --- a/src/metric-engine/src/engine/read.rs +++ b/src/metric-engine/src/engine/read.rs @@ -145,14 +145,10 @@ impl MetricEngineInner { mut request: ScanRequest, ) -> Result { // transform projection - let physical_projection = match request.projection_input.as_ref() { - Some(projection_input) => { - self.transform_projection( - physical_region_id, - logical_region_id, - &projection_input.projection, - ) - .await? + let physical_projection = match request.projection.as_ref() { + Some(projection) => { + self.transform_projection(physical_region_id, logical_region_id, projection) + .await? } None => { self.default_projection(physical_region_id, logical_region_id) @@ -160,10 +156,9 @@ impl MetricEngineInner { } }; - // Rewrite the top-level projection from logical-region schema indices to - // physical-region schema indices. `nested_paths` are left unchanged because - // they are expressed by column name rather than schema index. - request.projection_input.get_or_insert_default().projection = physical_projection; + // Rewrite the projection from logical-region schema indices to + // physical-region schema indices. + request.projection = Some(physical_projection); request .filters @@ -328,7 +323,7 @@ mod test { let logical_region_id = env.default_logical_region_id(); let invalid_index = usize::MAX; let request = ScanRequest { - projection_input: Some(vec![invalid_index].into()), + projection: Some(vec![invalid_index]), ..Default::default() }; @@ -373,9 +368,9 @@ mod test { .unwrap(); // check explicit projection - let projection_input = Some(vec![0, 1, 2, 3, 4, 5, 6].into()); + let projection = Some(vec![0, 1, 2, 3, 4, 5, 6]); let scan_req = ScanRequest { - projection_input, + projection, filters: vec![], ..Default::default() }; @@ -388,7 +383,7 @@ mod test { .unwrap(); assert_eq!( - scan_req.projection_indices().unwrap(), + scan_req.projection.as_deref().unwrap(), &[11, 10, 9, 8, 0, 1, 4] ); assert_eq!(scan_req.filters.len(), 1); @@ -407,7 +402,7 @@ mod test { .await .unwrap(); assert_eq!( - scan_req.projection_indices().unwrap(), + scan_req.projection.as_deref().unwrap(), &[11, 10, 9, 8, 0, 1, 4] ); } diff --git a/src/metric-engine/src/metadata_region.rs b/src/metric-engine/src/metadata_region.rs index 76abaacf10..e2a744b466 100644 --- a/src/metric-engine/src/metadata_region.rs +++ b/src/metric-engine/src/metadata_region.rs @@ -349,9 +349,8 @@ impl MetadataRegion { METADATA_SCHEMA_VALUE_COLUMN_INDEX, ] }; - let projection_input = Some(projection.into()); ScanRequest { - projection_input, + projection: Some(projection), filters: vec![filter_expr], ..Default::default() } @@ -362,9 +361,8 @@ impl MetadataRegion { METADATA_SCHEMA_KEY_COLUMN_INDEX, METADATA_SCHEMA_VALUE_COLUMN_INDEX, ]; - let projection_input = Some(projection.into()); ScanRequest { - projection_input, + projection: Some(projection), ..Default::default() } } @@ -681,9 +679,9 @@ impl MetadataRegion { let filter_expr = datafusion::prelude::col(METADATA_SCHEMA_KEY_COLUMN_NAME) .eq(datafusion::prelude::lit(key)); - let projection_input = Some(vec![METADATA_SCHEMA_VALUE_COLUMN_INDEX].into()); + let projection = Some(vec![METADATA_SCHEMA_VALUE_COLUMN_INDEX]); let scan_req = ScanRequest { - projection_input, + projection, filters: vec![filter_expr], ..Default::default() }; diff --git a/src/mito2/src/engine/partition_filter_test.rs b/src/mito2/src/engine/partition_filter_test.rs index 33f1d93aee..b71559c15e 100644 --- a/src/mito2/src/engine/partition_filter_test.rs +++ b/src/mito2/src/engine/partition_filter_test.rs @@ -126,9 +126,9 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) { .unwrap(); // Scan data in staging mode - should only see initial 5 rows (staging SST not visible) - let projection_input = Some(vec![1].into()); + let projection = Some(vec![1]); let request = ScanRequest { - projection_input, + projection, ..Default::default() }; let scanner = engine.scanner(region_id, request).await.unwrap(); @@ -150,9 +150,9 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) { // Scan after exiting staging - the old SST (tag_0 = "0".."4") should have // rows filtered by partition expr (tag_0 >= "5"), which means none of them pass. // But the staging SST (tag_0 = "5".."10") satisfies the partition expr. - let projection_input = Some(vec![1].into()); + let projection = Some(vec![1]); let request = ScanRequest { - projection_input, + projection, ..Default::default() }; let scanner = engine.scanner(region_id, request).await.unwrap(); diff --git a/src/mito2/src/engine/projection_test.rs b/src/mito2/src/engine/projection_test.rs index 3bc449d07b..5e6c38fa9d 100644 --- a/src/mito2/src/engine/projection_test.rs +++ b/src/mito2/src/engine/projection_test.rs @@ -106,9 +106,9 @@ async fn test_scan_projection_with_format(flat_format: bool) { put_rows(&engine, region_id, rows).await; // Scans tag_1, field_1, ts - let projection_input = Some(vec![1, 3, 4].into()); + let projection = Some(vec![1, 3, 4]); let request = ScanRequest { - projection_input, + projection, filters: Vec::new(), ..Default::default() }; @@ -184,9 +184,9 @@ async fn test_scan_projection_without_primary_key_with_format(flat_format: bool) put_rows(&engine, region_id, rows).await; // Scan with projection on field_0 and field_1, filter ts >= 2s - let projection_input = Some(vec![0, 1].into()); + let projection = Some(vec![0, 1]); let request = ScanRequest { - projection_input, // field_0 and field_1 (not ts) + projection, // field_0 and field_1 (not ts) filters: vec![col("ts").gt_eq(lit(ScalarValue::TimestampMillisecond(Some(2000), None)))], ..Default::default() }; diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index fc346412d0..b9eed37cb3 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -34,7 +34,7 @@ use futures::TryStreamExt; use serde_json::json; use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner}; use store_api::region_request::RegionRequest; -use store_api::storage::{ProjectionInput, RegionId, ScanRequest, TimeSeriesDistribution}; +use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution}; use crate::config::MitoConfig; use crate::error::Error; @@ -97,22 +97,48 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul test_util::put_rows(&engine, region_id, rows).await; test_util::flush_region(&engine, region_id, None).await; + // Without a type hint, the scanner reads the whole JSON2 root column. + let request = ScanRequest { + projection: Some(vec![1, 0]), + ..Default::default() + }; + let scanner = engine.scanner(region_id, request).await?; + let Scanner::Seq(seq_scan) = &scanner else { + unreachable!(); + }; + assert_eq!( + seq_scan.input().read_cols, + ReadColumns::from_deduped_column_ids([1, 0]) + ); + + 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 | ++------------------------------------------+-------+ +"#; + assert_eq!(batches.pretty_print()?, expected.trim()); + // Simulate a query expression like json_get(field_0, 'a.x'): the logical projection still // returns the JSON2 root column, while json_type_hint tells scan input construction which // nested physical path is needed. - - let request = ScanRequest { - projection_input: Some(ProjectionInput::new(vec![1, 0])), - json_type_hint: HashMap::from([( - "field_0".to_string(), + let json_type_hint = HashMap::from([( + "field_0".to_string(), + JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), JsonNativeType::Object(JsonObjectType::from([( - "a".to_string(), - JsonNativeType::Object(JsonObjectType::from([( - "x".to_string(), - JsonNativeType::i64(), - )])), + "x".to_string(), + JsonNativeType::i64(), )])), - )]), + )])), + )]); + let request = ScanRequest { + projection: Some(vec![1, 0]), + json_type_hint: json_type_hint.clone(), ..Default::default() }; let scanner = engine.scanner(region_id, request).await?; @@ -125,7 +151,6 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul seq_scan.input().read_cols, ReadColumns { cols: vec![ - ReadColumn::new(0, vec![]), ReadColumn::new( 1, vec![vec![ @@ -134,6 +159,7 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul "x".to_string() ]] ), + ReadColumn::new(0, vec![]), ] } ); @@ -150,6 +176,33 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul | {a: {x: 10}} | tag-1 | | {a: {x: 20}} | tag-2 | +--------------+-------+ +"#; + assert_eq!(batches.pretty_print()?, expected.trim()); + + // A type hint only narrows a projected JSON2 root; it must not add an unprojected root. + let request = ScanRequest { + projection: Some(vec![0]), + json_type_hint, + ..Default::default() + }; + let scanner = engine.scanner(region_id, request).await?; + let Scanner::Seq(seq_scan) = &scanner else { + unreachable!(); + }; + assert_eq!( + seq_scan.input().read_cols, + ReadColumns::from_deduped_column_ids([0]) + ); + + let stream = scanner.scan().await?; + let batches = RecordBatches::try_collect(stream).await?; + let expected = r#" ++-------+ +| tag_0 | ++-------+ +| tag-1 | +| tag-2 | ++-------+ "#; assert_eq!(batches.pretty_print()?, expected.trim()); Ok(()) diff --git a/src/mito2/src/read/read_columns.rs b/src/mito2/src/read/read_columns.rs index 2178fa2538..21699b8c43 100644 --- a/src/mito2/src/read/read_columns.rs +++ b/src/mito2/src/read/read_columns.rs @@ -12,17 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{BTreeMap, HashSet}; use std::mem; -use datafusion_common::HashMap; -use datafusion_expr::utils::expr_to_columns; -use snafu::OptionExt; -use store_api::metadata::RegionMetadataRef; -use store_api::storage::{ColumnId, NestedPath, ProjectionInput}; - -use crate::error::{InvalidRequestSnafu, Result}; -use crate::read::scan_region::PredicateGroup; +use store_api::storage::{ColumnId, NestedPath}; /// Logical columns to read from a region. /// @@ -136,422 +128,3 @@ impl ReadColumn { .sum::() } } - -pub fn merge(a: ReadColumns, b: ReadColumns) -> ReadColumns { - let mut merged = BTreeMap::>::new(); - - for col in a.cols.into_iter().chain(b.cols) { - if let Some(nested_paths) = merged.get_mut(&col.column_id) { - if nested_paths.is_empty() || col.nested_paths.is_empty() { - *nested_paths = vec![]; - } else { - merge_nested_paths(nested_paths, col.nested_paths); - } - continue; - } - - merged.insert(col.column_id, normalize_nested_paths(col.nested_paths)); - } - - ReadColumns { - cols: merged - .into_iter() - .map(|(column_id, nested_paths)| ReadColumn { - column_id, - nested_paths, - }) - .collect(), - } -} - -fn normalize_nested_paths(nested_paths: Vec) -> Vec { - let mut normalized = Vec::with_capacity(nested_paths.len()); - merge_nested_paths(&mut normalized, nested_paths); - normalized -} - -pub(crate) fn merge_nested_paths(merged: &mut Vec, incoming: Vec) { - for path in incoming { - if merged - .iter() - .any(|existing| path.starts_with(existing.as_slice())) - { - continue; - } - - merged.retain(|existing| !existing.starts_with(path.as_slice())); - merged.push(path); - } -} - -/// Build [`ReadColumns`] from [`ProjectionInput`]. -/// -/// Note: If `projection.projection` is empty, this function still reads the -/// time index column so the scan can preserve row counts for empty-output -/// queries such as `SELECT COUNT(*)`. -/// -/// Order: -/// - This function keeps the first-seen order from `projection.projection` -/// (duplicate indices are skipped). -/// - Keeping a stable order makes [`ReadColumns`] comparisons deterministic -/// (`Eq`/`Hash`) and avoids cache-key instability in upper layers. -pub fn read_columns_from_projection( - projection: ProjectionInput, - metadata: &RegionMetadataRef, -) -> Result { - let root_indices = if projection.projection.is_empty() { - vec![metadata.time_index_column_pos()] - } else { - projection.projection - }; - - let mut paths_by_col: HashMap> = - HashMap::with_capacity(projection.nested_paths.len()); - for path in projection.nested_paths { - let Some((root_name, _)) = path.split_first() else { - continue; - }; - paths_by_col - .entry(root_name.clone()) - .or_default() - .push(path); - } - - let mut read_cols = Vec::with_capacity(root_indices.len()); - let mut seen = HashSet::with_capacity(root_indices.len()); - for root_idx in root_indices { - if !seen.insert(root_idx) { - continue; - } - - let col = metadata - .column_metadatas - .get(root_idx) - .with_context(|| InvalidRequestSnafu { - region_id: metadata.region_id, - reason: format!("projection index {} is out of bounds", root_idx), - })?; - let col_id = col.column_id; - - let nested_paths = paths_by_col - .remove(&col.column_schema.name) - .unwrap_or_default(); - - read_cols.push(ReadColumn { - column_id: col_id, - nested_paths, - }); - } - - Ok(ReadColumns { cols: read_cols }) -} - -/// Build [`ReadColumns`] from [`PredicateGroup`]. -/// -/// Order: -/// - This function follows `metadata.column_metadatas` order when materializing -/// columns from predicate-referenced names. -/// - Using metadata order keeps the output deterministic for [`ReadColumns`] -/// equality/hash checks and for cache keys derived from read columns. -pub fn read_columns_from_predicate( - predicate: &PredicateGroup, - metadata: &RegionMetadataRef, -) -> ReadColumns { - let mut root_names = HashSet::new(); - let mut columns = HashSet::new(); - - if let Some(p) = predicate.predicate_without_region() { - for expr in p.exprs() { - columns.clear(); - if expr_to_columns(expr, &mut columns).is_err() { - continue; - } - root_names.extend(columns.drain().map(|column| column.name)); - } - } - - if let Some(expr) = predicate.region_partition_expr() { - expr.collect_column_names(&mut root_names); - } - - // TODO(fys): Parse nested paths from predicate expressions and attach them - // to read columns instead of always reading the whole root column. - let mut cols = Vec::with_capacity(root_names.len()); - for column in &metadata.column_metadatas { - if root_names.contains(&column.column_schema.name) { - cols.push(ReadColumn::new(column.column_id, vec![])); - } - } - - ReadColumns { cols } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use api::v1::SemanticType; - use datafusion_expr::{col, lit}; - use datatypes::prelude::ConcreteDataType; - use datatypes::schema::ColumnSchema; - use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; - use store_api::storage::RegionId; - - use super::*; - - #[test] - fn test_read_columns_from_empty_projection() { - let metadata = new_test_metadata(); - - let read_columns = - read_columns_from_projection(ProjectionInput::default(), &metadata).unwrap(); - - let expected = ReadColumns { - cols: vec![ReadColumn::new(2, vec![])], - }; - assert_eq!(expected, read_columns); - - let projection_input = - ProjectionInput::new(vec![]).with_nested_paths(vec![vec!["1".to_string()]]); - let read_columns = read_columns_from_projection(projection_input, &metadata).unwrap(); - - let expected = ReadColumns { - cols: vec![ReadColumn::new(2, vec![])], - }; - assert_eq!(expected, read_columns); - } - - #[test] - fn test_read_columns_from_projection_with_nested_paths() { - let metadata = new_test_metadata(); - let projection = ProjectionInput::new(vec![1, 0]).with_nested_paths(vec![ - nested_path(&["field_0", "a"]), - nested_path(&["field_0", "b", "c"]), - ]); - - let read_columns = read_columns_from_projection(projection, &metadata).unwrap(); - - let expected = ReadColumns { - cols: vec![ - ReadColumn::new( - 3, - vec![ - nested_path(&["field_0", "a"]), - nested_path(&["field_0", "b", "c"]), - ], - ), - ReadColumn::new(0, vec![]), - ], - }; - assert_eq!(expected, read_columns,); - } - - #[test] - fn test_read_columns_from_projection_dedups_duplicate_indices() { - let metadata = new_test_metadata(); - let projection = ProjectionInput::new(vec![1, 1, 0]).with_nested_paths(vec![ - nested_path(&["field_0", "a"]), - nested_path(&["field_0", "b", "c"]), - ]); - - let read_columns = read_columns_from_projection(projection, &metadata).unwrap(); - - let expected = ReadColumns { - cols: vec![ - ReadColumn::new( - 3, - vec![ - nested_path(&["field_0", "a"]), - nested_path(&["field_0", "b", "c"]), - ], - ), - ReadColumn::new(0, vec![]), - ], - }; - assert_eq!(expected, read_columns); - } - - #[test] - fn test_read_columns_from_projection_out_of_bound() { - let metadata = new_test_metadata(); - let projection = ProjectionInput::new(vec![3]); - - let err = read_columns_from_projection(projection, &metadata).unwrap_err(); - - assert!( - err.to_string() - .contains("projection index 3 is out of bound") - ); - } - - #[test] - fn test_read_columns_from_predicate_reads_root_columns_only() { - let metadata = new_test_metadata(); - let predicate = PredicateGroup::new( - metadata.as_ref(), - &[col("field_0").gt(lit(1)), col("tag_0").eq(lit("a"))], - ) - .unwrap(); - - let read_columns = read_columns_from_predicate(&predicate, &metadata); - - let expected = ReadColumns { - cols: vec![ReadColumn::new(0, vec![]), ReadColumn::new(3, vec![])], - }; - assert_eq!(expected, read_columns); - } - - #[test] - fn test_read_columns_from_predicate_empty() { - let metadata = new_test_metadata(); - let predicate = PredicateGroup::new(metadata.as_ref(), &[]).unwrap(); - - let read_columns = read_columns_from_predicate(&predicate, &metadata); - - assert!(read_columns.is_empty()); - } - - #[test] - fn test_merge_read_cols_with_only_root() { - let a = ReadColumns { - cols: vec![ReadColumn::new(3, vec![]), ReadColumn::new(1, vec![])], - }; - let b = ReadColumns { - cols: vec![ReadColumn::new(2, vec![])], - }; - - let merged = merge(a, b); - - assert_eq!( - merged, - ReadColumns { - cols: vec![ - ReadColumn::new(1, vec![]), - ReadColumn::new(2, vec![]), - ReadColumn::new(3, vec![]), - ], - } - ); - } - - #[test] - fn test_merge_read_cols_with_nested_paths() { - let a = ReadColumns { - cols: vec![ReadColumn::new(1, vec![nested_path(&["j", "a"])])], - }; - let b = ReadColumns { - cols: vec![ReadColumn::new( - 1, - vec![nested_path(&["j", "b"]), nested_path(&["j", "c"])], - )], - }; - - let merged = merge(a, b); - - assert_eq!( - merged, - ReadColumns { - cols: vec![ReadColumn::new( - 1, - vec![ - nested_path(&["j", "a"]), - nested_path(&["j", "b"]), - nested_path(&["j", "c"]), - ], - )], - } - ); - } - - #[test] - fn test_merge_read_cols_with_column_override() { - let a = ReadColumns { - cols: vec![ - ReadColumn::new(1, vec![nested_path(&["j", "a"])]), - ReadColumn::new(2, vec![nested_path(&["k", "b"])]), - ], - }; - let b = ReadColumns { - cols: vec![ - ReadColumn::new(1, vec![]), - ReadColumn::new(2, vec![nested_path(&["k", "b", "c"])]), - ], - }; - - let merged = merge(a, b); - - assert_eq!( - merged, - ReadColumns { - cols: vec![ - ReadColumn::new(1, vec![]), - ReadColumn::new(2, vec![nested_path(&["k", "b"])]) - ], - } - ); - } - - #[test] - fn test_merge_read_cols_dedups_redundant_nested_paths() { - let a = ReadColumns { - cols: vec![ReadColumn::new( - 1, - vec![ - nested_path(&["j", "a", "b"]), - nested_path(&["j", "a"]), - nested_path(&["j", "a", "b", "c"]), - ], - )], - }; - let b = ReadColumns { - cols: vec![ReadColumn::new(1, vec![nested_path(&["j", "a"])])], - }; - - let merged = merge(a, b); - - assert_eq!( - merged, - ReadColumns { - cols: vec![ReadColumn::new(1, vec![nested_path(&["j", "a"])])], - } - ); - } - - fn new_test_metadata() -> RegionMetadataRef { - let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); - builder - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - true, - ), - semantic_type: SemanticType::Tag, - column_id: 0, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "field_0".to_string(), - ConcreteDataType::string_datatype(), - true, - ), - semantic_type: SemanticType::Field, - column_id: 3, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - semantic_type: SemanticType::Timestamp, - column_id: 2, - }); - builder.primary_key(vec![0]); - Arc::new(builder.build().unwrap()) - } - - fn nested_path(parts: &[&str]) -> NestedPath { - parts.iter().map(|part| (*part).to_string()).collect() - } -} diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 763597b742..cc6d7b8d98 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::{HashMap, HashSet}; +use std::collections::HashSet; use std::fmt; use std::num::NonZeroU64; use std::sync::Arc; @@ -41,12 +41,12 @@ use futures::StreamExt; use itertools::Itertools; use partition::expr::PartitionExpr; use smallvec::SmallVec; -use snafu::ResultExt; +use snafu::{OptionExt, ResultExt}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::region_engine::{PartitionRange, RegionScannerRef}; use store_api::storage::{ - NestedPath, RegionId, ScanRequest, SequenceNumber, SequenceRange, TimeSeriesDistribution, - TimeSeriesRowSelector, + ColumnId, NestedPath, RegionId, ScanRequest, SequenceNumber, SequenceRange, + TimeSeriesDistribution, TimeSeriesRowSelector, }; use table::predicate::{Predicate, build_time_range_predicate, extract_time_range_from_expr}; use tokio::sync::{Semaphore, mpsc}; @@ -55,7 +55,7 @@ use tokio_stream::wrappers::ReceiverStream; use crate::access_layer::AccessLayerRef; use crate::cache::CacheStrategy; use crate::config::DEFAULT_MAX_CONCURRENT_SCAN_FILES; -use crate::error::{InvalidPartitionExprSnafu, Result}; +use crate::error::{InvalidPartitionExprSnafu, InvalidRequestSnafu, Result}; #[cfg(feature = "enterprise")] use crate::extension::{BoxedExtensionRange, BoxedExtensionRangeProvider}; use crate::memtable::{MemtableRange, RangesOptions}; @@ -64,10 +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::{ - ReadColumns, merge, merge_nested_paths, read_columns_from_predicate, - read_columns_from_projection, -}; +use crate::read::read_columns::{ReadColumn, ReadColumns}; use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; use crate::read::stream::ScanBatchStream; @@ -416,55 +413,37 @@ impl ScanRegion { /// Creates a scan input. #[tracing::instrument(skip_all, fields(region_id = %self.region_id()))] async fn scan_input(self) -> Result { + let metadata = &self.version.metadata; let sst_min_sequence = self.request.sst_min_sequence.and_then(NonZeroU64::new); let time_range = self.build_time_range_predicate(); - let predicate = PredicateGroup::new(&self.version.metadata, &self.request.filters)?; + let predicate = PredicateGroup::new(metadata, &self.request.filters)?; - let mut read_cols = match &self.request.projection_input { - Some(p) => { - // Read columns include the pushed-down projection and columns - // resolved from the predicate. - let metadata = &self.version.metadata; - let from_projection = read_columns_from_projection(p.clone(), metadata)?; - let from_predicate = read_columns_from_predicate(&predicate, metadata); - merge(from_projection, from_predicate) - } - None => { - let read_col_ids = self - .version - .metadata - .column_metadatas - .iter() - .map(|col| col.column_id); - ReadColumns::from_deduped_column_ids(read_col_ids) - } - }; - // Only narrow read columns and pass JSON type hints for structured JSON (JSON2) - // columns. Legacy JSONB columns have JSON extension metadata but their physical - // Arrow type is Binary, not Struct, so they must not enter structured JSON paths. - let has_structured_json = self - .version - .metadata + 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_structured_json_field); - if has_structured_json { - narrow_read_columns_by_json_type_hint( - &mut read_cols, - &self.request.json_type_hint, - &self.version.metadata, - ); - } - let read_col_ids = read_cols.column_ids(); + 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()) + }; // The mapper always computes projected column ids as the schema of SSTs may change. let projection = self .request - .projection_indices() - .map(|x| x.to_vec()) - .unwrap_or_else(|| (0..self.version.metadata.column_metadatas.len()).collect()); + .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| { @@ -477,7 +456,7 @@ impl ScanRegion { ); }); let mapper = FlatProjectionMapper::new_with_read_columns( - &self.version.metadata, + metadata, projection, read_cols, json_type_hint, @@ -630,6 +609,108 @@ impl ScanRegion { Ok(input) } + /// Builds the ordered root column ids required by the pushed-down projection and predicate. + fn build_read_col_ids( + &self, + projection: Option<&[usize]>, + predicate: &PredicateGroup, + ) -> Result> { + let metadata = &self.version.metadata; + let Some(projection) = projection else { + return Ok(metadata + .column_metadatas + .iter() + .map(|col| col.column_id) + .collect()); + }; + + let mut read_col_ids = Vec::new(); + let mut seen = HashSet::new(); + for idx in projection { + let col_id = metadata + .column_metadatas + .get(*idx) + .with_context(|| InvalidRequestSnafu { + region_id: metadata.region_id, + reason: format!("projection index {} is out of bounds", idx), + })? + .column_id; + let inserted = seen.insert(col_id); + // DataFusion's logical `OptimizeProjections` rule deduplicates pushed-down + // projection indices via `RequiredIndices::compact()`. + debug_assert!( + inserted, + "projection contains duplicate column id: {}", + col_id + ); + // Keep the projection order. + read_col_ids.push(col_id); + } + + if projection.is_empty() { + let time_index = metadata.time_index_column().column_id; + if seen.insert(time_index) { + read_col_ids.push(time_index); + } + } + + let mut extra_col_names = HashSet::new(); + let mut cols = HashSet::new(); + + if let Some(p) = predicate.predicate_without_region() { + for expr in p.exprs() { + cols.clear(); + if expr_to_columns(expr, &mut cols).is_err() { + continue; + } + extra_col_names.extend(cols.iter().map(|col| col.name.clone())); + } + } + + if let Some(expr) = predicate.region_partition_expr() { + expr.collect_column_names(&mut extra_col_names); + } + + if !extra_col_names.is_empty() { + for col in &metadata.column_metadatas { + if extra_col_names.remove(&col.column_schema.name) && !seen.contains(&col.column_id) + { + read_col_ids.push(col.column_id); + } + } + if !extra_col_names.is_empty() { + warn!( + "Some columns in filters are not found in region {}: {:?}", + metadata.region_id, extra_col_names + ); + } + } + 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 + .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 } + } + fn region_id(&self) -> RegionId { self.version.metadata.region_id } @@ -1500,29 +1581,11 @@ fn pre_filter_mode(append_mode: bool, merge_mode: MergeMode) -> PreFilterMode { } } -fn narrow_read_columns_by_json_type_hint( - read_columns: &mut ReadColumns, - json_type_hint: &HashMap, - metadata: &RegionMetadata, -) { - if json_type_hint.is_empty() { - return; - } - - for read_column in &mut read_columns.cols { - let Some(column) = metadata.column_by_id(read_column.column_id) else { - continue; - }; - let column_name = &column.column_schema.name; - let Some(json_type) = json_type_hint.get(column_name) else { - continue; - }; - - let mut paths = Vec::new(); - let mut current = vec![column_name.clone()]; - collect_json_nested_paths(json_type, &mut current, &mut paths); - merge_nested_paths(&mut read_column.nested_paths, paths) - } +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( @@ -2052,9 +2115,7 @@ mod tests { use super::*; use crate::cache::CacheManager; - use crate::error::InvalidMetadataSnafu; use crate::read::range_cache::ScanRequestFingerprintBuilder; - use crate::read::read_columns::ReadColumn; use crate::sst::file::FileMeta; use crate::test_util::memtable_util::metadata_with_primary_key; use crate::test_util::scheduler_util::SchedulerEnv; @@ -2164,85 +2225,24 @@ mod tests { #[test] fn test_fill_json_nested_paths_from_hint() -> Result<()> { - fn json_projection_test_metadata() -> Result { - let mut builder = RegionMetadataBuilder::new(RegionId::new(1024, 0)); - builder - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - true, - ), - semantic_type: SemanticType::Tag, - column_id: 0, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "j".to_string(), - ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), - true, - ), - semantic_type: SemanticType::Field, - column_id: 1, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - semantic_type: SemanticType::Timestamp, - column_id: 2, - }); - builder.primary_key(vec![0]); - builder.build().context(InvalidMetadataSnafu).map(Arc::new) - } - - let metadata = json_projection_test_metadata()?; - let hint = HashMap::from([( - "j".to_string(), - JsonNativeType::Object(JsonObjectType::from([ - ("a".to_string(), JsonNativeType::i64()), - ( - "b".to_string(), - JsonNativeType::Object(JsonObjectType::from([( - "c".to_string(), - JsonNativeType::String, - )])), - ), - ])), - )]); + 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() } - let mut read_columns = ReadColumns { - cols: vec![ReadColumn::new(1, vec![]), ReadColumn::new(0, vec![])], - }; - narrow_read_columns_by_json_type_hint(&mut read_columns, &hint, metadata.as_ref()); assert_eq!( - read_columns, - ReadColumns { - cols: vec![ - ReadColumn::new( - 1, - vec![nested_path(&["j", "a"]), nested_path(&["j", "b", "c"])] - ), - ReadColumn::new(0, vec![]) - ] - } - ); - - let mut read_columns = ReadColumns { - cols: vec![ReadColumn::new(0, vec![])], - }; - narrow_read_columns_by_json_type_hint(&mut read_columns, &hint, metadata.as_ref()); - assert_eq!( - read_columns, - ReadColumns { - cols: vec![ReadColumn::new(0, vec![])] - } + json_nested_paths("j", &hint), + vec![nested_path(&["j", "a"]), nested_path(&["j", "b", "c"])] ); Ok(()) } diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 9ed77eed6c..83ccc27379 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -251,7 +251,7 @@ impl TestDataSource { impl DataSource for TestDataSource { fn get_stream(&self, request: ScanRequest) -> Result { - let projected_schema = match request.projection_indices() { + let projected_schema = match request.projection.as_deref() { Some(projection) => Arc::new(self.schema.try_project(projection).unwrap()), None => self.schema.clone(), }; diff --git a/src/query/src/dummy_catalog.rs b/src/query/src/dummy_catalog.rs index aedc274cfa..e38813b3e6 100644 --- a/src/query/src/dummy_catalog.rs +++ b/src/query/src/dummy_catalog.rs @@ -188,7 +188,7 @@ impl TableProvider for DummyTableProvider { limit: Option, ) -> datafusion::error::Result> { let mut request = self.scan_request.lock().unwrap().clone(); - request.projection_input = projection.map(|p| p.clone().into()); + request.projection = projection.cloned(); request.filters = filters.to_vec(); request.limit = limit; if let Some(query_ctx) = &self.query_ctx { diff --git a/src/store-api/src/region_info.rs b/src/store-api/src/region_info.rs index 2d099d2734..164ea9d68a 100644 --- a/src/store-api/src/region_info.rs +++ b/src/store-api/src/region_info.rs @@ -133,7 +133,7 @@ impl RegionInfoEntry { pub fn build_plan(scan_request: ScanRequest) -> Result { let table_source = LogicalTableSource::new(Self::schema().arrow_schema().clone()); - let projection = scan_request.projection_input.map(|input| input.projection); + let projection = scan_request.projection; let mut builder = LogicalPlanBuilder::scan( Self::reserved_table_name_for_inspection(), Arc::new(table_source), @@ -315,9 +315,9 @@ mod tests { #[test] fn test_region_info_build_plan() { - let projection_input = Some(vec![0, 5, 7, 11].into()); + let projection = Some(vec![0, 5, 7, 11]); let request = ScanRequest { - projection_input, + projection, filters: vec![binary_expr(col("writable"), Operator::Eq, lit(true))], limit: Some(10), ..Default::default() diff --git a/src/store-api/src/sst_entry.rs b/src/store-api/src/sst_entry.rs index cca327fb4d..f1bb1f9c36 100644 --- a/src/store-api/src/sst_entry.rs +++ b/src/store-api/src/sst_entry.rs @@ -404,7 +404,7 @@ fn build_plan_helper( ) -> Result { let table_source = LogicalTableSource::new(schema.arrow_schema().clone()); - let projection = scan_request.projection_input.map(|input| input.projection); + let projection = scan_request.projection; let mut builder = LogicalPlanBuilder::scan(table_name, Arc::new(table_source), projection)?; for filter in scan_request.filters { @@ -947,9 +947,9 @@ mod tests { #[test] fn test_manifest_build_plan() { // Note: filter must reference a column in the projected schema - let projection_input = Some(vec![0, 1, 2].into()); + let projection = Some(vec![0, 1, 2]); let request = ScanRequest { - projection_input, + projection, filters: vec![binary_expr(col("table_id"), Operator::Gt, lit(0))], limit: Some(5), ..Default::default() @@ -979,9 +979,9 @@ mod tests { #[test] fn test_storage_build_plan() { - let projection_input = Some(vec![0, 2].into()); + let projection = Some(vec![0, 2]); let request = ScanRequest { - projection_input, + projection, filters: vec![binary_expr(col("file_path"), Operator::Eq, lit("/a"))], limit: Some(1), ..Default::default() diff --git a/src/store-api/src/storage.rs b/src/store-api/src/storage.rs index 33aa3a7b41..cae2439168 100644 --- a/src/store-api/src/storage.rs +++ b/src/store-api/src/storage.rs @@ -28,7 +28,7 @@ pub use datatypes::schema::{ pub use self::descriptors::*; pub use self::file::{FileId, FileRef, FileRefsManifest, GcReport, IndexVersion, ParseIdError}; -pub use self::projection::{NestedPath, ProjectionInput}; +pub use self::projection::NestedPath; pub use self::requests::{ ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector, VectorDistanceMetric, VectorIndexEngine, VectorIndexEngineType, VectorSearchMatches, VectorSearchRequest, diff --git a/src/store-api/src/storage/projection.rs b/src/store-api/src/storage/projection.rs index f8f27fe5a4..60e295ae80 100644 --- a/src/store-api/src/storage/projection.rs +++ b/src/store-api/src/storage/projection.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::fmt::{Display, Formatter}; - /// A nested field access path. /// /// Each path represents a field access on a nested column. @@ -21,58 +19,3 @@ use std::fmt::{Display, Formatter}; /// Example: /// - `j.a.b` -> `["j", "a", "b"]` pub type NestedPath = Vec; - -/// Projection information for a table scan. -#[derive(Default, Debug, Clone, PartialEq)] -pub struct ProjectionInput { - /// Top-level column projection. - /// - /// The indices are based on the schema exposed by the table scan input, - /// such as the schema passed to `TableProvider::scan`. - /// - /// Only the root columns with the specified schema indices are needed. - pub projection: Vec, - /// Nested field access paths used for sub-field projection. - /// - /// It extends and refines the top-level projection by specifying nested - /// field accesses inside complex columns such as JSON or struct columns. - /// - /// In other words: - /// - `projection` determines **which root columns are needed** - /// - `nested_paths` further determines **which sub-fields inside those - /// columns are required** - /// - /// Each path starts with the root column name and continues with - /// nested field names. - pub nested_paths: Vec, -} - -impl ProjectionInput { - pub fn new(projection: Vec) -> Self { - Self { - projection, - nested_paths: Vec::new(), - } - } - - pub fn with_nested_paths(mut self, nested_paths: Vec) -> Self { - self.nested_paths = nested_paths; - self - } -} - -impl From> for ProjectionInput { - fn from(projection: Vec) -> Self { - Self::new(projection) - } -} - -impl Display for ProjectionInput { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!( - f, - "ProjectionInput {{ projection: {:?}, nested_paths: {:?} }}", - self.projection, self.nested_paths - ) - } -} diff --git a/src/store-api/src/storage/requests.rs b/src/store-api/src/storage/requests.rs index ef694eaa01..7ae1bc5987 100644 --- a/src/store-api/src/storage/requests.rs +++ b/src/store-api/src/storage/requests.rs @@ -24,7 +24,7 @@ use datatypes::types::json_type::JsonNativeType; use itertools::Itertools; use strum::Display; -use crate::storage::{ColumnId, ProjectionInput, SequenceNumber}; +use crate::storage::{ColumnId, SequenceNumber}; /// A hint for KNN vector search. #[derive(Debug, Clone, PartialEq)] @@ -100,7 +100,7 @@ pub enum TimeSeriesDistribution { pub struct ScanRequest { /// Optional projection information for the scan. `None` reads all root /// columns. - pub projection_input: Option, + pub projection: Option>, /// Filters pushed down pub filters: Vec, /// Expected output ordering. This is only a hint and isn't guaranteed. @@ -140,15 +140,6 @@ pub struct ScanRequest { pub preserve_pk_dictionary_encoding: bool, } -impl ScanRequest { - /// Returns the top-level projected column indices. - pub fn projection_indices(&self) -> Option<&[usize]> { - self.projection_input - .as_ref() - .map(|projection_input| projection_input.projection.as_slice()) - } -} - impl Display for ScanRequest { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { enum Delimiter { @@ -171,7 +162,7 @@ impl Display for ScanRequest { let mut delimiter = Delimiter::None; write!(f, "ScanRequest {{ ")?; - if let Some(projection) = &self.projection_input { + if let Some(projection) = &self.projection { write!(f, "{}projection: {:?}", delimiter.as_str(), projection)?; } if !self.filters.is_empty() { @@ -280,9 +271,9 @@ mod tests { }; assert_eq!(request.to_string(), "ScanRequest { }"); - let projection_input = Some(vec![1, 2].into()); + let projection = Some(vec![1, 2]); let request = ScanRequest { - projection_input, + projection, filters: vec![ binary_expr(col("i"), Operator::Gt, lit(1)), binary_expr(col("s"), Operator::Eq, lit("x")), @@ -292,7 +283,7 @@ mod tests { }; assert_eq!( request.to_string(), - r#"ScanRequest { projection: ProjectionInput { projection: [1, 2], nested_paths: [] }, filters: [i > Int32(1), s = Utf8("x")], limit: 10 }"# + r#"ScanRequest { projection: [1, 2], filters: [i > Int32(1), s = Utf8("x")], limit: 10 }"# ); let request = ScanRequest { @@ -308,29 +299,15 @@ mod tests { r#"ScanRequest { filters: [i > Int32(1), s = Utf8("x")], limit: 10 }"# ); - let projection_input = Some(vec![1, 2].into()); + let projection = Some(vec![1, 2]); let request = ScanRequest { - projection_input, + projection, limit: Some(10), ..Default::default() }; assert_eq!( request.to_string(), - "ScanRequest { projection: ProjectionInput { projection: [1, 2], nested_paths: [] }, limit: 10 }" - ); - - let projection_input = Some(ProjectionInput::new(vec![1, 2]).with_nested_paths(vec![ - vec!["j".to_string(), "a".to_string(), "b".to_string()], - vec!["s".to_string(), "x".to_string()], - ])); - let request = ScanRequest { - projection_input, - limit: Some(10), - ..Default::default() - }; - assert_eq!( - request.to_string(), - r#"ScanRequest { projection: ProjectionInput { projection: [1, 2], nested_paths: [["j", "a", "b"], ["s", "x"]] }, limit: 10 }"# + "ScanRequest { projection: [1, 2], limit: 10 }" ); let request = ScanRequest { diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index 1c4a605e0e..0ae56ce1c3 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -152,11 +152,10 @@ impl TableProvider for DfTableProviderAdapter { limit: Option, ) -> DfResult> { let filters: Vec = filters.iter().map(Clone::clone).collect(); - let projection_input = projection.map(|p| p.clone().into()); let request = { let mut request = self.scan_req.lock().unwrap(); request.filters = filters; - request.projection_input = projection_input; + request.projection = projection.cloned(); request.limit = limit; request.clone() }; diff --git a/src/table/src/table/numbers.rs b/src/table/src/table/numbers.rs index cb4b8ce317..70c744b13a 100644 --- a/src/table/src/table/numbers.rs +++ b/src/table/src/table/numbers.rs @@ -110,7 +110,7 @@ impl NumbersDataSource { impl DataSource for NumbersDataSource { fn get_stream(&self, request: ScanRequest) -> Result { - let projection = request.projection_input.map(|input| input.projection); + let projection = request.projection; let projected_schema = match &projection { Some(projection) => Arc::new(self.schema.try_project(projection).unwrap()), None => self.schema.clone(), diff --git a/src/table/src/test_util/memtable.rs b/src/table/src/test_util/memtable.rs index ee73123aca..0ae751d35f 100644 --- a/src/table/src/test_util/memtable.rs +++ b/src/table/src/test_util/memtable.rs @@ -121,7 +121,7 @@ impl DataSource for MemtableDataSource { &self, request: ScanRequest, ) -> std::result::Result { - let df_recordbatch = if let Some(indices) = request.projection_indices() { + let df_recordbatch = if let Some(indices) = request.projection.as_deref() { self.recordbatch .df_record_batch() .project(indices) @@ -198,9 +198,8 @@ mod test { async fn test_scan_with_projection() { let table = build_testing_table(); - let projection_input = Some(vec![1].into()); let scan_req = ScanRequest { - projection_input, + projection: Some(vec![1]), ..Default::default() }; let stream = table.scan_to_stream(scan_req).await.unwrap(); diff --git a/tests/cases/distributed/information_schema/ssts_limit.result b/tests/cases/distributed/information_schema/ssts_limit.result index 0b64a3c0fa..248593410f 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.result +++ b/tests/cases/distributed/information_schema/ssts_limit.result @@ -72,27 +72,27 @@ FROM ( LIMIT 1 ); -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | Projection: count(Int64(1)) AS count(*) AS filtered_limited_rows | -| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | -| | Limit: skip=0, fetch=1 | -| | Projection: information_schema.ssts_manifest.region_id | -| | Filter: information_schema.ssts_manifest.table_id > UInt32(0) | -| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.table_id > UInt32(0)] | -| | ]] | -| physical_plan | ProjectionExec: expr=[count(Int64(1))@0 as filtered_limited_rows] | -| | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] | -| | CoalescePartitionsExec | -| | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] | -| | CoalescePartitionsExec: fetch=1 | -| | FilterExec: table_id@0 > 0, projection=[], fetch=1 | ++---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: count(Int64(1)) AS count(*) AS filtered_limited_rows | +| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | +| | Limit: skip=0, fetch=1 | +| | Projection: information_schema.ssts_manifest.region_id | +| | Filter: information_schema.ssts_manifest.table_id > UInt32(0) | +| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.table_id > UInt32(0)] | +| | ]] | +| physical_plan | ProjectionExec: expr=[count(Int64(1))@0 as filtered_limited_rows] | +| | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] | +| | CoalescePartitionsExec | +| | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] | +| | CoalescePartitionsExec: fetch=1 | +| | FilterExec: table_id@0 > 0, projection=[], fetch=1 | | | RepartitionExec: REDACTED -| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: ProjectionInput { projection: [2], nested_paths: [] }, filters: [table_id > UInt32(0)] }, schema=Schema { fields: [Field { name: "table_id", data_type: UInt32 }], metadata: {"greptime:version": "0"} } | -| | | -+---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: [2], filters: [table_id > UInt32(0)] }, schema=Schema { fields: [Field { name: "table_id", data_type: UInt32 }], metadata: {"greptime:version": "0"} } | +| | | ++---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ DROP TABLE ssts_limit_case;