diff --git a/src/mito-codec/src/error.rs b/src/mito-codec/src/error.rs index 78a656415a..3d3b0efcc0 100644 --- a/src/mito-codec/src/error.rs +++ b/src/mito-codec/src/error.rs @@ -67,6 +67,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid sparse primary key: {}", reason))] + InvalidSparsePrimaryKey { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Encode null value"))] IndexEncodeNull { #[snafu(implicit)] @@ -94,6 +101,7 @@ impl ErrorExt for Error { StatusCode::InvalidArguments } NotSupportedField { .. } | UnsupportedOperation { .. } => StatusCode::Unsupported, + InvalidSparsePrimaryKey { .. } => StatusCode::InvalidArguments, EvaluateFilter { source, .. } => source.status_code(), } } diff --git a/src/mito-codec/src/row_converter/sparse.rs b/src/mito-codec/src/row_converter/sparse.rs index 37e1813470..c70293886f 100644 --- a/src/mito-codec/src/row_converter/sparse.rs +++ b/src/mito-codec/src/row_converter/sparse.rs @@ -21,13 +21,16 @@ use datatypes::prelude::ConcreteDataType; use datatypes::value::{Value, ValueRef}; use memcomparable::{Deserializer, Serializer}; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; +use snafu::{ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; use store_api::storage::ColumnId; use store_api::storage::consts::ReservedColumnId; -use crate::error::{DeserializeFieldSnafu, Result, SerializeFieldSnafu, UnsupportedOperationSnafu}; +use crate::error::{ + DeserializeFieldSnafu, InvalidSparsePrimaryKeySnafu, Result, SerializeFieldSnafu, + UnsupportedOperationSnafu, +}; use crate::key_values::KeyValue; use crate::primary_key_filter::SparsePrimaryKeyFilter; use crate::row_converter::dense::SortField; @@ -346,6 +349,57 @@ impl SparsePrimaryKeyCodec { Ok(()) } + /// Decodes the reserved `(table_id, tsid)` prefix from a sparse primary key. + pub fn decode_ids(&self, bytes: &[u8]) -> Result<(u32, u64)> { + // Two column IDs, two non-null markers, a u32 table ID, and a u64 TSID. + const INTERNAL_PREFIX_LEN: usize = 4 + 1 + 4 + 4 + 1 + 8; + ensure!( + bytes.len() >= INTERNAL_PREFIX_LEN, + InvalidSparsePrimaryKeySnafu { + reason: format!( + "internal prefix requires at least {INTERNAL_PREFIX_LEN} bytes, got {}", + bytes.len() + ), + } + ); + + let mut deserializer = Deserializer::new(bytes); + + let table_id_column = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?; + ensure!( + table_id_column == RESERVED_COLUMN_ID_TABLE_ID, + InvalidSparsePrimaryKeySnafu { + reason: format!( + "expected table id column {}, got {}", + RESERVED_COLUMN_ID_TABLE_ID, table_id_column + ), + } + ); + let table_id = self.inner.table_id_field.deserialize(&mut deserializer)?; + + let tsid_column = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?; + ensure!( + tsid_column == RESERVED_COLUMN_ID_TSID, + InvalidSparsePrimaryKeySnafu { + reason: format!( + "expected tsid column {}, got {}", + RESERVED_COLUMN_ID_TSID, tsid_column + ), + } + ); + let tsid = self.inner.tsid_field.deserialize(&mut deserializer)?; + + match (table_id, tsid) { + (Value::UInt32(table_id), Value::UInt64(tsid)) => Ok((table_id, tsid)), + (table_id, tsid) => InvalidSparsePrimaryKeySnafu { + reason: format!( + "expected UInt32 table id and UInt64 tsid, got {table_id:?} and {tsid:?}" + ), + } + .fail(), + } + } + /// Decodes the given bytes into a [`SparseValues`]. fn decode_sparse(&self, bytes: &[u8]) -> Result { let mut deserializer = Deserializer::new(bytes); @@ -805,6 +859,21 @@ mod tests { assert_eq!(result, Value::UInt32(42)); } + #[test] + fn test_decode_ids() { + let region_metadata = test_region_metadata(); + let codec = SparsePrimaryKeyCodec::new(®ion_metadata); + let mut buffer = Vec::new(); + codec.encode_internal(42, 100, &mut buffer).unwrap(); + + assert_eq!((42, 100), codec.decode_ids(&buffer).unwrap()); + + let mut invalid = buffer.clone(); + invalid[0..4].copy_from_slice(&1_u32.to_be_bytes()); + assert!(codec.decode_ids(&invalid).is_err()); + assert!(codec.decode_ids(&buffer[..buffer.len() - 1]).is_err()); + } + #[test] fn test_has_column() { let region_metadata = test_region_metadata(); diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 5e93f011fa..f25b7475fd 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -465,6 +465,14 @@ pub enum Error { error: datafusion::error::DataFusionError, }, + #[snafu(display("Failed to merge candidate series"))] + MergeCandidateSeries { + #[snafu(implicit)] + location: Location, + #[snafu(source)] + error: datafusion::error::DataFusionError, + }, + #[snafu(display("Failed to compute vector"))] ComputeVector { #[snafu(implicit)] @@ -1402,6 +1410,7 @@ impl ErrorExt for Error { | DecodeWal { .. } | ComputeArrow { .. } | EvalPartitionFilter { .. } + | MergeCandidateSeries { .. } | BiErrors { .. } | StopScheduler { .. } | ComputeVector { .. } diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs index 192bc92b16..1f162fc439 100644 --- a/src/mito2/src/read.rs +++ b/src/mito2/src/read.rs @@ -33,6 +33,7 @@ pub mod read_columns; pub mod scan_region; pub mod scan_util; pub(crate) mod seq_scan; +pub(crate) mod series_candidate; pub mod series_scan; pub mod stream; pub(crate) mod unordered_scan; diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index 48f641b1ef..3f58b3bd43 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -41,7 +41,7 @@ use store_api::storage::ColumnId; use crate::error::{ CompatReaderSnafu, ComputeArrowSnafu, ConvertValueSnafu, CreateDefaultSnafu, DecodeSnafu, - EncodeSnafu, NewRecordBatchSnafu, Result, UnsupportedOperationSnafu, + 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}; @@ -117,7 +117,10 @@ impl FlatCompatBatch { } let expect_schema = mapper.batch_schema(); - if expect_schema == actual_schema { + if expect_schema == actual_schema + && actual.primary_key == mapper.metadata().primary_key + && actual.primary_key_encoding == mapper.metadata().primary_key_encoding + { // Although the SST has a different schema, but the schema after projection is the same // as expected schema. return Ok(None); @@ -295,11 +298,16 @@ impl FlatCompatBatch { ) .collect::>>()?; - let compat_batch = RecordBatch::try_new(self.arrow_schema.clone(), columns) - .context(NewRecordBatchSnafu)?; + let mut columns = columns; + let primary_key_index = primary_key_column_index(columns.len()); + columns[primary_key_index] = self.compat_primary_key(&columns[primary_key_index])?; - // Handles primary keys. - self.compat_pk.compat(compat_batch) + RecordBatch::try_new(self.arrow_schema.clone(), columns).context(NewRecordBatchSnafu) + } + + /// Makes an encoded primary-key array compatible with the expected metadata. + pub(crate) fn compat_primary_key(&self, primary_key: &ArrayRef) -> Result { + self.compat_pk.compat(primary_key) } } @@ -408,18 +416,44 @@ impl FlatRewritePrimaryKey { fn rewrite_key( &self, append_values: &[(ColumnId, Value)], - batch: RecordBatch, - ) -> Result { - let old_pk_dict_array = batch - .column(primary_key_column_index(batch.num_columns())) - .as_any() - .downcast_ref::() - .unwrap(); - let old_pk_values_array = old_pk_dict_array - .values() - .as_any() - .downcast_ref::() - .unwrap(); + primary_key: &ArrayRef, + ) -> Result { + if let Some(old_pk_dict_array) = primary_key.as_any().downcast_ref::() { + let old_pk_values_array = old_pk_dict_array + .values() + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: "Primary-key dictionary values are not binary", + })?; + let new_pk_values_array = + Arc::new(self.rewrite_values(append_values, old_pk_values_array)?); + return Ok(Arc::new(PrimaryKeyArray::new( + old_pk_dict_array.keys().clone(), + new_pk_values_array, + ))); + } + + let old_pk_values_array = + primary_key + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: format!( + "Primary-key column is neither binary nor dictionary, got {:?}", + primary_key.data_type() + ), + })?; + Ok(Arc::new( + self.rewrite_values(append_values, old_pk_values_array)?, + )) + } + + fn rewrite_values( + &self, + append_values: &[(ColumnId, Value)], + old_pk_values_array: &BinaryArray, + ) -> Result { let mut builder = BinaryBuilder::with_capacity( old_pk_values_array.len(), old_pk_values_array.value_data().len(), @@ -461,14 +495,7 @@ impl FlatRewritePrimaryKey { } builder.append_value(&buffer); } - let new_pk_values_array = Arc::new(builder.finish()); - let new_pk_dict_array = - PrimaryKeyArray::new(old_pk_dict_array.keys().clone(), new_pk_values_array); - - let mut columns = batch.columns().to_vec(); - columns[primary_key_column_index(batch.num_columns())] = Arc::new(new_pk_dict_array); - - RecordBatch::try_new(batch.schema(), columns).context(NewRecordBatchSnafu) + Ok(builder.finish()) } } @@ -538,34 +565,58 @@ impl FlatCompatPrimaryKey { }) } - /// Makes primary key of the `batch` compatible. - /// - /// Callers must ensure other columns except the `__primary_key` column is compatible. - fn compat(&self, batch: RecordBatch) -> Result { + /// Makes an encoded primary-key array compatible. + fn compat(&self, primary_key: &ArrayRef) -> Result { if let Some(rewriter) = &self.rewriter { // If we have different encoding, rewrite the whole primary key. - return rewriter.rewrite_key(&self.values, batch); + return rewriter.rewrite_key(&self.values, primary_key); } - self.append_key(batch) + self.append_key(primary_key) } - /// Appends values to the primary key of the `batch`. - fn append_key(&self, batch: RecordBatch) -> Result { + /// Appends values to the primary key array. + fn append_key(&self, primary_key: &ArrayRef) -> Result { let Some(converter) = &self.converter else { - return Ok(batch); + return Ok(primary_key.clone()); }; - let old_pk_dict_array = batch - .column(primary_key_column_index(batch.num_columns())) - .as_any() - .downcast_ref::() - .unwrap(); - let old_pk_values_array = old_pk_dict_array - .values() - .as_any() - .downcast_ref::() - .unwrap(); + if let Some(old_pk_dict_array) = primary_key.as_any().downcast_ref::() { + let old_pk_values_array = old_pk_dict_array + .values() + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: "Primary-key dictionary values are not binary", + })?; + let new_pk_values_array = + Arc::new(self.append_values(old_pk_values_array, converter.as_ref())?); + return Ok(Arc::new(PrimaryKeyArray::new( + old_pk_dict_array.keys().clone(), + new_pk_values_array, + ))); + } + + let old_pk_values_array = + primary_key + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: format!( + "Primary-key column is neither binary nor dictionary, got {:?}", + primary_key.data_type() + ), + })?; + Ok(Arc::new( + self.append_values(old_pk_values_array, converter.as_ref())?, + )) + } + + fn append_values( + &self, + old_pk_values_array: &BinaryArray, + converter: &dyn PrimaryKeyCodec, + ) -> Result { let mut builder = BinaryBuilder::with_capacity( old_pk_values_array.len(), old_pk_values_array.value_data().len() @@ -594,15 +645,7 @@ impl FlatCompatPrimaryKey { builder.append_value(&buffer); } - let new_pk_values_array = Arc::new(builder.finish()); - let new_pk_dict_array = - PrimaryKeyArray::new(old_pk_dict_array.keys().clone(), new_pk_values_array); - - // Overrides the primary key column. - let mut columns = batch.columns().to_vec(); - columns[primary_key_column_index(batch.num_columns())] = Arc::new(new_pk_dict_array); - - RecordBatch::try_new(batch.schema(), columns).context(NewRecordBatchSnafu) + Ok(builder.finish()) } } @@ -612,7 +655,7 @@ mod tests { use api::v1::{OpType, SemanticType}; use datatypes::arrow::array::{ - ArrayRef, BinaryDictionaryBuilder, Int64Array, StringDictionaryBuilder, + ArrayRef, BinaryArray, BinaryDictionaryBuilder, Int64Array, StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, }; use datatypes::arrow::datatypes::UInt32Type; @@ -974,6 +1017,64 @@ mod tests { assert_eq!(expected_batch, result); } + #[test] + fn test_compat_primary_key_with_different_encoding_only() { + let mut actual_metadata = new_metadata( + &[ + ( + 0, + SemanticType::Timestamp, + ConcreteDataType::timestamp_millisecond_datatype(), + ), + (1, SemanticType::Tag, ConcreteDataType::string_datatype()), + (2, SemanticType::Field, ConcreteDataType::int64_datatype()), + ], + &[1], + ); + actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Dense; + let actual_metadata = Arc::new(actual_metadata); + + let mut expected_metadata = (*actual_metadata).clone(); + expected_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse; + let expected_metadata = Arc::new(expected_metadata); + + let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); + let read_format = FlatReadFormat::new( + actual_metadata, + ReadColumns::from_deduped_column_ids([0, 1, 2]), + None, + "test", + false, + ) + .unwrap(); + let compat = FlatCompatBatch::try_new(&mapper, &read_format, false) + .unwrap() + .unwrap(); + + let dense_key = encode_key(&[Some("tag1")]); + let sparse_key = encode_sparse_key(&[(1, Some("tag1"))]); + + let dictionary_key = build_flat_test_pk_array(&[&dense_key, &dense_key]); + let result = compat.compat_primary_key(&dictionary_key).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + let values = result + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(result.keys().value(0) as usize), sparse_key); + assert_eq!(values.value(result.keys().value(1) as usize), sparse_key); + + let binary_key: ArrayRef = Arc::new(BinaryArray::from(vec![ + Some(dense_key.as_slice()), + Some(dense_key.as_slice()), + ])); + let result = compat.compat_primary_key(&binary_key).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.value(0), sparse_key); + assert_eq!(result.value(1), sparse_key); + } + #[test] fn test_flat_compat_batch_compact_sparse() { let mut actual_metadata = new_metadata( diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index d305caf23b..9f11ccac99 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -62,11 +62,18 @@ pub(crate) struct ScanRequestFingerprint { append_mode: bool, filter_deleted: bool, merge_mode: MergeMode, + stage: RangeScanStage, /// We keep the partition expr version to ensure we won't reuse the fingerprint after we change the partition expr. /// We store the version instead of the whole partition expr or partition expr filters. partition_expr_version: u64, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum RangeScanStage { + Data, + CandidateSeries, +} + #[derive(Debug)] pub(crate) struct ScanRequestFingerprintBuilder { pub(crate) read_columns: ReadColumns, @@ -105,6 +112,7 @@ impl ScanRequestFingerprintBuilder { append_mode, filter_deleted, merge_mode, + stage: RangeScanStage::Data, partition_expr_version, } } @@ -154,6 +162,20 @@ impl ScanRequestFingerprint { append_mode: self.append_mode, filter_deleted: self.filter_deleted, merge_mode: self.merge_mode, + stage: self.stage, + partition_expr_version: self.partition_expr_version, + } + } + + fn for_candidate_series(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + time_filters: self.time_filters.clone(), + series_row_selector: self.series_row_selector, + append_mode: self.append_mode, + filter_deleted: self.filter_deleted, + merge_mode: self.merge_mode, + stage: RangeScanStage::CandidateSeries, partition_expr_version: self.partition_expr_version, } } @@ -447,6 +469,22 @@ fn implied_from_between( pub(crate) fn build_range_cache_key( stream_ctx: &StreamContext, part_range: &PartitionRange, +) -> Option { + build_range_cache_key_inner(stream_ctx, part_range, false) +} + +/// Builds a cache key for a candidate-series partition-range result. +pub(crate) fn build_candidate_range_cache_key( + stream_ctx: &StreamContext, + part_range: &PartitionRange, +) -> Option { + build_range_cache_key_inner(stream_ctx, part_range, true) +} + +fn build_range_cache_key_inner( + stream_ctx: &StreamContext, + part_range: &PartitionRange, + candidate_series: bool, ) -> Option { if !stream_ctx.input.cache_strategy.has_range_result_cache() { return None; @@ -499,6 +537,11 @@ pub(crate) fn build_range_cache_key( } else { fingerprint.clone() }; + let scan = if candidate_series { + scan.for_candidate_series() + } else { + scan + }; Some(RangeScanCacheKey { region_id: stream_ctx.input.region_metadata().region_id, @@ -1090,6 +1133,22 @@ mod tests { assert!(key_a.scan.time_filters().is_empty()); } + #[tokio::test] + async fn candidate_series_cache_key_is_separate_from_data() { + let partition_range = ( + Timestamp::new_millisecond(1000), + Timestamp::new_millisecond(2000), + ); + let (ctx, part_range) = + new_stream_context(vec![col("k0").eq(lit("foo"))], None, partition_range).await; + + let data_key = build_range_cache_key(&ctx, &part_range).unwrap(); + let candidate_key = build_candidate_range_cache_key(&ctx, &part_range).unwrap(); + + assert_ne!(data_key.scan, candidate_key.scan); + assert_eq!(data_key.row_groups, candidate_key.row_groups); + } + #[tokio::test] async fn disables_optimization_on_or_clause() { let partition_range = ( diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs new file mode 100644 index 0000000000..16ef944c4a --- /dev/null +++ b/src/mito2/src/read/series_candidate.rs @@ -0,0 +1,737 @@ +// 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. + +//! Candidate metric-series discovery for the two-stage series scan. + +use std::sync::Arc; +use std::time::Instant; + +use async_stream::try_stream; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool}; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_plan::expressions::Column; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; +use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion_common::DataFusionError; +use datatypes::arrow::array::{Array, BinaryArray, BinaryBuilder}; +use datatypes::arrow::compute::SortOptions; +use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; +use mito_codec::row_converter::{PrimaryKeyFilter, SparsePrimaryKeyCodec}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::region_engine::PartitionRange; +use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, ReservedColumnId}; +use tokio::sync::Semaphore; + +use crate::error::{ + InvalidRequestSnafu, JoinSnafu, MergeCandidateSeriesSnafu, NewRecordBatchSnafu, Result, + UnexpectedSnafu, +}; +use crate::read::BoxedRecordBatchStream; +use crate::read::pruner::{PartitionPruner, Pruner}; +use crate::read::range::RowGroupIndex; +use crate::read::range_cache::{ + build_candidate_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, +}; +use crate::read::scan_region::StreamContext; +use crate::read::scan_util::{PartitionMetrics, new_filter_metrics, scan_flat_mem_ranges}; +use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; +use crate::sst::parquet::format::PrimaryKeyArray; +use crate::sst::parquet::prefilter::{ + CachedPrimaryKeyFilter, build_primary_key_filter, prefilter_flat_batch_by_primary_key, +}; +use crate::sst::parquet::reader::ReaderMetrics; +use crate::sst::parquet::row_group::ParquetFetchMetrics; + +const CANDIDATE_SERIES_BATCH_SIZE: usize = 500; + +/// Identifies one series in a physical metric region. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct MetricSeriesId { + pub(crate) table_id: u32, + pub(crate) tsid: u64, +} + +pub(crate) type MetricSeriesIdStream = BoxStream<'static, Result>>; + +/// Builds candidate metric series from the ranges assigned to a [`SeriesScan`](super::series_scan::SeriesScan). +#[allow(dead_code)] +pub(crate) struct SeriesCandidateScanner { + stream_ctx: Arc, + partitions: Vec>, + pruner: Arc, + range_semaphore: Arc, + memory_pool: Arc, + metrics_set: ExecutionPlanMetricsSet, + part_metrics: PartitionMetrics, +} + +#[allow(dead_code)] +impl SeriesCandidateScanner { + /// Creates a candidate-series scanner for native memtable and SST ranges. + /// + /// Callers must fall back to the legacy series-scan path when the scan contains + /// extension ranges. Candidate-series discovery does not support other range types. + pub(crate) fn try_new( + stream_ctx: Arc, + partitions: Vec>, + pruner: Arc, + range_semaphore: Arc, + memory_pool: Arc, + metrics_set: ExecutionPlanMetricsSet, + part_metrics: PartitionMetrics, + ) -> Result { + validate_metric_metadata(&stream_ctx)?; + #[cfg(feature = "enterprise")] + ensure!( + stream_ctx.input.extension_ranges().is_empty(), + InvalidRequestSnafu { + region_id: stream_ctx.input.region_metadata().region_id, + reason: "candidate-series scan does not support extension ranges; use the legacy series-scan path", + } + ); + Ok(Self { + stream_ctx, + partitions, + pruner, + range_semaphore, + memory_pool, + metrics_set, + part_metrics, + }) + } + + /// Builds a globally sorted stream of candidate metric-series IDs. + pub(crate) async fn build_stream(&self) -> Result { + let all_ranges = self + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + self.pruner.add_partition_ranges(&all_ranges); + let partition_pruner = Arc::new(PartitionPruner::new(self.pruner.clone(), &all_ranges)); + + let range_builder = SeriesCandidateRangeBuilder { + stream_ctx: self.stream_ctx.clone(), + range_semaphore: self.range_semaphore.clone(), + memory_pool: self.memory_pool.clone(), + metrics_set: self.metrics_set.clone(), + part_metrics: self.part_metrics.clone(), + }; + let mut tasks = Vec::with_capacity(all_ranges.len()); + for (range_idx, part_range) in all_ranges.into_iter().enumerate() { + let range_builder = range_builder.clone(); + let partition_pruner = partition_pruner.clone(); + tasks.push(common_runtime::spawn_query(async move { + let _permit = range_builder + .range_semaphore + .clone() + .acquire_owned() + .await + .map_err(|error| { + UnexpectedSnafu { + reason: format!("failed to acquire candidate range permit: {error}"), + } + .build() + })?; + range_builder + .build_range_stream(part_range, partition_pruner, range_idx) + .await + })); + } + + let mut range_streams = Vec::with_capacity(tasks.len()); + for task in tasks { + range_streams.push(task.await.context(JoinSnafu)??); + } + + // Keep scanner-level merge metrics in the same synthetic partition as + // SeriesDistributor. Output partitions occupy 0..self.partitions.len(). + let merged = merge_primary_key_streams( + range_streams, + self.memory_pool.clone(), + &self.metrics_set, + self.partitions.len(), + "SeriesCandidateScanner::final_merge", + )?; + decode_metric_series(merged, self.stream_ctx.input.region_metadata().clone()) + } +} + +#[derive(Clone)] +struct SeriesCandidateRangeBuilder { + stream_ctx: Arc, + range_semaphore: Arc, + memory_pool: Arc, + metrics_set: ExecutionPlanMetricsSet, + part_metrics: PartitionMetrics, +} + +impl SeriesCandidateRangeBuilder { + async fn build_range_stream( + &self, + part_range: PartitionRange, + partition_pruner: Arc, + merge_partition: usize, + ) -> Result { + let cache_key = build_candidate_range_cache_key(&self.stream_ctx, &part_range); + if let Some(key) = cache_key.as_ref() { + if let Some(value) = self.stream_ctx.input.cache_strategy.get_range_result(key) { + self.part_metrics.inc_range_cache_hit(); + return Ok(cached_flat_range_stream(value)); + } + self.part_metrics.inc_range_cache_miss(); + } + + let range_meta = &self.stream_ctx.ranges[part_range.identifier]; + let mut sources = Vec::with_capacity(range_meta.row_group_indices.len()); + for index in &range_meta.row_group_indices { + let source = self + .build_source(*index, range_meta.time_range, partition_pruner.clone()) + .await?; + if let Some(source) = source { + sources.push(source); + } + } + + let sources = self.stream_ctx.input.create_parallel_flat_sources( + sources, + self.range_semaphore.clone(), + 2, + )?; + let stream = merge_primary_key_streams( + sources, + self.memory_pool.clone(), + &self.metrics_set, + merge_partition, + "SeriesCandidateScanner::range_merge", + )?; + + Ok(match cache_key { + Some(key) => cache_flat_range_stream( + stream, + self.stream_ctx.input.cache_strategy.clone(), + key, + self.part_metrics.clone(), + ), + None => stream, + }) + } + + async fn build_source( + &self, + index: RowGroupIndex, + time_range: crate::sst::file::FileTimeRange, + partition_pruner: Arc, + ) -> Result> { + let metadata = self.stream_ctx.input.region_metadata().clone(); + if self.stream_ctx.is_mem_range_index(index) { + let raw = scan_flat_mem_ranges( + self.stream_ctx.clone(), + self.part_metrics.clone(), + index, + time_range, + ); + let filter = build_primary_key_filter( + &metadata, + self.stream_ctx.input.predicate_group().predicate(), + ); + return Ok(Some(candidate_primary_key_stream(Box::pin(raw), filter))); + } + + if self.stream_ctx.is_file_range_index(index) { + if partition_pruner.try_skip_manifest_pruned_file_range(index, &self.part_metrics) { + return Ok(None); + } + + let mut reader_metrics = ReaderMetrics { + filter_metrics: new_filter_metrics(self.part_metrics.explain_verbose()), + ..Default::default() + }; + let ranges = partition_pruner + .build_file_ranges(index, &self.part_metrics, &mut reader_metrics) + .await?; + self.part_metrics.inc_num_file_ranges(ranges.len()); + self.part_metrics + .merge_reader_metrics(&reader_metrics, None); + + // Reuse the exact encoded-PK filter selected by the file's reader plan, but + // execute it in `candidate_primary_key_stream` after the PK-only read. + let filter = ranges.first().and_then(|range| range.primary_key_filter()); + let part_metrics = self.part_metrics.clone(); + let raw = Box::pin(try_stream! { + let fetch_metrics = part_metrics + .explain_verbose() + .then(|| Arc::new(ParquetFetchMetrics::default())); + let mut reader_metrics = ReaderMetrics { + fetch_metrics: fetch_metrics.clone(), + ..Default::default() + }; + for range in ranges { + let build_start = Instant::now(); + let Some(mut reader) = range + .primary_key_reader(fetch_metrics.as_deref()) + .await? + else { + continue; + }; + reader_metrics.build_cost += build_start.elapsed(); + + let scan_start = Instant::now(); + while let Some(batch) = reader.try_next().await? { + reader_metrics.num_record_batches += 1; + reader_metrics.num_batches += 1; + reader_metrics.num_rows += batch.num_rows(); + yield batch; + } + reader_metrics.scan_cost += scan_start.elapsed(); + } + reader_metrics.observe_rows("candidate_series"); + part_metrics.merge_reader_metrics(&reader_metrics, None); + }); + return Ok(Some(candidate_primary_key_stream(raw, filter))); + } + + UnexpectedSnafu { + reason: format!( + "candidate-series scan received unsupported range index {}", + index.index + ), + } + .fail() + } +} + +fn validate_metric_metadata(stream_ctx: &StreamContext) -> Result<()> { + let metadata = stream_ctx.input.region_metadata(); + let valid_prefix = metadata + .primary_key + .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]); + let valid_types = metadata + .column_by_id(ReservedColumnId::table_id()) + .zip(metadata.column_by_id(ReservedColumnId::tsid())) + .is_some_and(|(table_id, tsid)| { + table_id.column_schema.data_type == ConcreteDataType::uint32_datatype() + && tsid.column_schema.data_type == ConcreteDataType::uint64_datatype() + }); + ensure!( + metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse && valid_prefix && valid_types, + InvalidRequestSnafu { + region_id: metadata.region_id, + reason: "candidate-series scan requires sparse (__table_id, __tsid) primary keys", + } + ); + Ok(()) +} + +fn primary_key_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new( + PRIMARY_KEY_COLUMN_NAME, + DataType::Binary, + false, + )])) +} + +/// Filters a source by encoded-primary-key predicates and emits one binary row per local key. +fn candidate_primary_key_stream( + mut input: BoxedRecordBatchStream, + mut filter: Option, +) -> BoxedRecordBatchStream { + Box::pin(try_stream! { + let mut last_primary_key = Vec::new(); + let mut has_last = false; + while let Some(batch) = input.try_next().await? { + if let Some(batch) = normalize_candidate_batch( + batch, + filter.as_mut(), + &mut last_primary_key, + &mut has_last, + )? { + yield batch; + } + } + }) +} + +fn normalize_candidate_batch( + mut batch: RecordBatch, + filter: Option<&mut CachedPrimaryKeyFilter>, + last_primary_key: &mut Vec, + has_last: &mut bool, +) -> Result> { + let pk_idx = batch + .schema() + .column_with_name(PRIMARY_KEY_COLUMN_NAME) + .map(|(idx, _)| idx) + .context(UnexpectedSnafu { + reason: "candidate source does not contain __primary_key", + })?; + if let Some(filter) = filter { + let Some(filtered) = prefilter_flat_batch_by_primary_key( + batch, + pk_idx, + filter as &mut dyn PrimaryKeyFilter, + )? + else { + return Ok(None); + }; + batch = filtered; + } + + let pk_column = batch.column(pk_idx); + let mut builder = BinaryBuilder::new(); + if let Some(array) = pk_column.as_any().downcast_ref::() { + let values = array + .values() + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: "dictionary primary-key values are not binary", + })?; + for key in array.keys().values() { + append_unique_primary_key( + values.value(*key as usize), + &mut builder, + last_primary_key, + has_last, + ); + } + } else if let Some(array) = pk_column.as_any().downcast_ref::() { + for value in array.iter().flatten() { + append_unique_primary_key(value, &mut builder, last_primary_key, has_last); + } + } else { + return UnexpectedSnafu { + reason: format!( + "primary-key column is neither binary nor dictionary, got {:?}", + pk_column.data_type() + ), + } + .fail(); + } + + let array = builder.finish(); + if array.is_empty() { + return Ok(None); + } + let batch = RecordBatch::try_new(primary_key_schema(), vec![Arc::new(array)]) + .context(NewRecordBatchSnafu)?; + Ok(Some(batch)) +} + +fn append_unique_primary_key( + value: &[u8], + builder: &mut BinaryBuilder, + last_primary_key: &mut Vec, + has_last: &mut bool, +) { + if !*has_last || last_primary_key != value { + builder.append_value(value); + last_primary_key.clear(); + last_primary_key.extend_from_slice(value); + *has_last = true; + } +} + +fn merge_primary_key_streams( + sources: Vec, + memory_pool: Arc, + metrics_set: &ExecutionPlanMetricsSet, + partition: usize, + consumer_name: &'static str, +) -> Result { + if sources.is_empty() { + return Ok(Box::pin(futures::stream::empty())); + } + if sources.len() == 1 { + return Ok(sources.into_iter().next().unwrap()); + } + + let schema = primary_key_schema(); + let df_sources = sources + .into_iter() + .map(|source| { + let stream = source.map_err(|error| DataFusionError::External(Box::new(error))); + Box::pin(RecordBatchStreamAdapter::new(schema.clone(), stream)) as _ + }) + .collect(); + let ordering = LexOrdering::new([PhysicalSortExpr { + expr: Arc::new(Column::new(PRIMARY_KEY_COLUMN_NAME, 0)), + options: SortOptions { + descending: false, + nulls_first: false, + }, + }]) + // Safe to unwrap because `LexOrdering::new` returns `None` only for empty + // input, and this array always contains one sort expression. + .unwrap(); + let reservation = MemoryConsumer::new(consumer_name).register(&memory_pool); + let mut merged = StreamingMergeBuilder::new() + .with_streams(df_sources) + .with_schema(schema) + .with_expressions(&ordering) + .with_metrics(BaselineMetrics::new(metrics_set, partition)) + .with_batch_size(DEFAULT_READ_BATCH_SIZE) + .with_reservation(reservation) + .build() + .context(MergeCandidateSeriesSnafu)?; + + Ok(Box::pin(try_stream! { + while let Some(batch) = merged.next().await { + yield batch.context(MergeCandidateSeriesSnafu)?; + } + })) +} + +fn decode_metric_series( + mut input: BoxedRecordBatchStream, + metadata: store_api::metadata::RegionMetadataRef, +) -> Result { + let codec = SparsePrimaryKeyCodec::new(&metadata); + Ok(Box::pin(try_stream! { + let mut last_series = None; + let mut output = Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE); + while let Some(batch) = input.try_next().await? { + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: "merged candidate primary key is not binary", + })?; + for primary_key in array.iter().flatten() { + let (table_id, tsid) = codec + .decode_ids(primary_key) + .context(crate::error::DecodeSnafu)?; + let series = MetricSeriesId { table_id, tsid }; + if last_series == Some(series) { + continue; + } + last_series = Some(series); + output.push(series); + if output.len() == CANDIDATE_SERIES_BATCH_SIZE { + yield std::mem::replace( + &mut output, + Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE), + ); + } + } + } + if !output.is_empty() { + yield output; + } + })) +} + +#[cfg(test)] +mod tests { + use datafusion::execution::memory_pool::UnboundedMemoryPool; + use datafusion_expr::{col, lit}; + use datatypes::arrow::array::{ArrayRef, DictionaryArray, UInt32Array}; + use datatypes::arrow::datatypes::UInt32Type; + use futures::TryStreamExt; + use store_api::codec::PrimaryKeyEncoding; + use table::predicate::Predicate; + + use super::*; + use crate::test_util::sst_util::sst_region_metadata_with_encoding; + + fn binary_batch(values: &[&[u8]]) -> RecordBatch { + RecordBatch::try_new( + primary_key_schema(), + vec![Arc::new(BinaryArray::from_iter_values( + values.iter().copied(), + ))], + ) + .unwrap() + } + + fn dictionary_batch(values: &[&[u8]], keys: &[u32]) -> RecordBatch { + let dict_values: ArrayRef = Arc::new(BinaryArray::from_iter_values(values.iter().copied())); + let dict = + DictionaryArray::::try_new(UInt32Array::from(keys.to_vec()), dict_values) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new_dictionary( + PRIMARY_KEY_COLUMN_NAME, + DataType::UInt32, + DataType::Binary, + false, + )])); + RecordBatch::try_new(schema, vec![Arc::new(dict)]).unwrap() + } + + #[tokio::test] + async fn candidate_stream_normalizes_and_deduplicates_primary_keys() { + let input = Box::pin(futures::stream::iter(vec![ + Ok(dictionary_batch(&[b"a", b"b"], &[0, 0, 1])), + Ok(binary_batch(&[b"b", b"c", b"c"])), + ])); + let batches = candidate_primary_key_stream(input, None) + .try_collect::>() + .await + .unwrap(); + let actual = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect::>(); + assert_eq!(actual, vec![b"a".as_slice(), b"b", b"c"]); + } + + #[tokio::test] + async fn candidate_stream_filters_primary_keys_before_merge() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let codec = SparsePrimaryKeyCodec::new(&metadata); + let mut table_1 = Vec::new(); + let mut table_2 = Vec::new(); + codec.encode_internal(1, 10, &mut table_1).unwrap(); + codec.encode_internal(2, 20, &mut table_2).unwrap(); + + let predicate = Predicate::new(vec![ + col(store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME).eq(lit(1_u32)), + ]); + let filter = build_primary_key_filter(&metadata, Some(&predicate)); + let input = Box::pin(futures::stream::iter(vec![Ok(dictionary_batch( + &[table_1.as_slice(), table_2.as_slice()], + &[0, 1], + ))])); + let batches = candidate_primary_key_stream(input, filter) + .try_collect::>() + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + let array = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(array.len(), 1); + assert_eq!(array.value(0), table_1); + } + + #[tokio::test] + async fn decode_metric_series_yields_groups_of_500() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let codec = SparsePrimaryKeyCodec::new(&metadata); + let primary_keys = (0..501_u64) + .map(|tsid| { + let mut primary_key = Vec::new(); + codec.encode_internal(1, tsid, &mut primary_key).unwrap(); + primary_key + }) + .collect::>(); + let batch = binary_batch(&primary_keys.iter().map(Vec::as_slice).collect::>()); + let source = Box::pin(futures::stream::iter(vec![Ok(batch)])); + let metrics = ExecutionPlanMetricsSet::new(); + let pool = Arc::new(UnboundedMemoryPool::default()); + let merged = + merge_primary_key_streams(vec![source], pool, &metrics, 0, "candidate-test").unwrap(); + let groups = decode_metric_series(merged, metadata) + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(groups.iter().map(Vec::len).collect::>(), [500, 1]); + assert_eq!( + groups[0][0], + MetricSeriesId { + table_id: 1, + tsid: 0 + } + ); + assert_eq!( + groups[1][0], + MetricSeriesId { + table_id: 1, + tsid: 500 + } + ); + } + + #[tokio::test] + async fn merge_primary_keys_globally_sorts_and_deduplicates_series() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let codec = SparsePrimaryKeyCodec::new(&metadata); + let encode = |tsid| { + let mut primary_key = Vec::new(); + codec.encode_internal(1, tsid, &mut primary_key).unwrap(); + primary_key + }; + let keys_1 = [encode(1), encode(3)]; + let mut alternate_key_for_series_1 = encode(1); + alternate_key_for_series_1.push(0); + let keys_2 = [alternate_key_for_series_1, encode(2)]; + let sources = vec![ + Box::pin(futures::stream::iter(vec![Ok(binary_batch( + &keys_1.iter().map(Vec::as_slice).collect::>(), + ))])) as BoxedRecordBatchStream, + Box::pin(futures::stream::iter(vec![Ok(binary_batch( + &keys_2.iter().map(Vec::as_slice).collect::>(), + ))])), + ]; + + let merged = merge_primary_key_streams( + sources, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + "candidate-merge-test", + ) + .unwrap(); + let groups = decode_metric_series(merged, metadata) + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!( + groups, + vec![vec![ + MetricSeriesId { + table_id: 1, + tsid: 1, + }, + MetricSeriesId { + table_id: 1, + tsid: 2, + }, + MetricSeriesId { + table_id: 1, + tsid: 3, + }, + ]] + ); + } +} diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index 52c74ef70a..64c94938e7 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -27,6 +27,7 @@ use datatypes::arrow::array::{Array as _, ArrayRef, BooleanArray}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::record_batch::RecordBatch; use datatypes::schema::Schema; +use futures::StreamExt; use mito_codec::row_converter::PrimaryKeyCodec; use parquet::arrow::arrow_reader::RowSelection; use parquet::file::metadata::ParquetMetaData; @@ -49,6 +50,8 @@ use crate::sst::file::FileHandle; use crate::sst::parquet::flat_format::{ DecodedPrimaryKeys, FlatReadFormat, decode_primary_keys, time_index_column_index, }; +use crate::sst::parquet::json_align::ProjectedRecordBatchStream; +use crate::sst::parquet::prefilter::CachedPrimaryKeyFilter; use crate::sst::parquet::reader::{ FlatRowGroupReader, MaybeFilter, RowGroupBuildContext, RowGroupReaderBuilder, SimpleFilterContext, @@ -95,6 +98,11 @@ pub struct FileRange { } impl FileRange { + /// Builds the encoded-primary-key filter selected for this file. + pub(crate) fn primary_key_filter(&self) -> Option { + self.context.reader_builder.primary_key_filter() + } + /// Creates a new [FileRange]. pub(crate) fn new( context: FileRangeContextRef, @@ -224,6 +232,44 @@ impl FileRange { Ok(Some(flat_prune_reader)) } + /// Creates a reader that returns only the encoded primary-key column. + /// + /// The returned primary keys are compatible with the expected region metadata. + pub(crate) async fn primary_key_reader( + &self, + fetch_metrics: Option<&ParquetFetchMetrics>, + ) -> Result> { + if !self.in_dynamic_filter_range() { + return Ok(None); + } + + let stream = self + .context + .reader_builder + .build_primary_key(self.context.build_context( + self.row_group_idx, + self.row_selection.clone(), + fetch_metrics, + )) + .await?; + if self.context.compat_batch().is_none() { + return Ok(Some(stream)); + } + + let context = self.context.clone(); + let stream = stream + .map(move |batch| { + let batch = batch?; + let compat = context.compat_batch().context(UnexpectedSnafu { + reason: "Primary-key compatibility helper is missing", + })?; + let primary_key = compat.compat_primary_key(batch.column(0))?; + RecordBatch::try_new(batch.schema(), vec![primary_key]).context(NewRecordBatchSnafu) + }) + .boxed(); + Ok(Some(stream)) + } + /// Returns the helper to compat batches. pub(crate) fn compat_batch(&self) -> Option<&FlatCompatBatch> { self.context.compat_batch() diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index f8fb0dc668..514add7073 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -29,7 +29,7 @@ use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::datatypes::SchemaRef; use datatypes::arrow::record_batch::RecordBatch; use futures::StreamExt; -use mito_codec::row_converter::{PrimaryKeyCodec, PrimaryKeyFilter}; +use mito_codec::row_converter::{PrimaryKeyCodec, PrimaryKeyFilter, build_primary_key_codec}; use parquet::arrow::ProjectionMask; use parquet::arrow::arrow_reader::RowSelection; use parquet::schema::types::SchemaDescriptor; @@ -262,6 +262,34 @@ pub(crate) struct BulkFilterPlan { pub(crate) pk_filters: Option>>, } +/// Builds an encoded-primary-key filter from the supported tag predicates. +/// +/// Predicates on fields, timestamps, or unsupported expression shapes are intentionally +/// omitted. Callers use this as a pruning filter and must preserve the full predicate for +/// authoritative filtering later in the scan. +pub(crate) fn build_primary_key_filter( + metadata: &RegionMetadataRef, + predicate: Option<&Predicate>, +) -> Option { + let filters = predicate + .into_iter() + .flat_map(|predicate| predicate.exprs()) + .filter_map(|expr| SimpleFilterContext::new_opt(metadata, None, expr)) + .filter_map(|filter_ctx| { + (filter_ctx.semantic_type() == SemanticType::Tag) + .then(|| filter_ctx.filter().as_filter().cloned()) + .flatten() + }) + .collect::>(); + if filters.is_empty() { + return None; + } + + let codec = build_primary_key_codec(metadata.as_ref()); + let filter = codec.primary_key_filter(metadata, Arc::new(filters)); + Some(CachedPrimaryKeyFilter::new(filter)) +} + /// How the parquet reader should apply each predicate. /// /// The reader runs in two phases. Predicates routed into `prefilter_builder` @@ -584,12 +612,9 @@ impl PrefilterContextBuilder { /// Builds a [PrefilterContext] for a specific row group. pub(crate) fn build(&self) -> PrefilterContext { - let pk_filter = self.pk_filters.as_ref().map(|pk_filters| { - let pk_filter = self - .codec - .primary_key_filter(&self.metadata, Arc::clone(pk_filters)); - Box::new(CachedPrimaryKeyFilter::new(pk_filter)) as Box - }); + let pk_filter = self + .build_primary_key_filter() + .map(|filter| Box::new(filter) as Box); PrefilterContext { pk_filter, filters: self.filters.clone(), @@ -599,6 +624,16 @@ impl PrefilterContextBuilder { arrow_schema: self.arrow_schema.clone(), } } + + /// Builds a fresh encoded-primary-key filter selected by this plan. + pub(crate) fn build_primary_key_filter(&self) -> Option { + self.pk_filters.as_ref().map(|pk_filters| { + let filter = self + .codec + .primary_key_filter(&self.metadata, Arc::clone(pk_filters)); + CachedPrimaryKeyFilter::new(filter) + }) + } } const PREFILTER_COLUMN_RATIO_THRESHOLD: f64 = 0.5; @@ -1530,24 +1565,36 @@ mod tests { vec!["field_0"] ); - let field_0 = metadata.column_by_name("field_0").unwrap().column_id; - let ts = metadata.time_index_column().column_id; + let metric_metadata: RegionMetadataRef = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let field_0 = metric_metadata.column_by_name("field_0").unwrap().column_id; + let ts = metric_metadata.time_index_column().column_id; let projected_read_format = FlatReadFormat::new( - metadata.clone(), + metric_metadata.clone(), ReadColumns::from_deduped_column_ids([field_0, ts]), None, "test", true, ) .unwrap(); + let metric_codec = build_primary_key_codec(metric_metadata.as_ref()); let pk_prefilter_plan = build_reader_filter_plan( Some(&Predicate::new(vec![col("tag_0").eq(lit("a"))])), None, PreFilterMode::All, &projected_read_format, - &codec, + &metric_codec, ); assert!(pk_prefilter_plan.prefilter_builder.is_some()); + assert!( + pk_prefilter_plan + .prefilter_builder + .as_ref() + .unwrap() + .build_primary_key_filter() + .is_some() + ); assert!(pk_prefilter_plan.remaining_simple_filters.is_empty()); } diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index ea2c0e9d47..934d23c0e1 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -41,10 +41,11 @@ use parquet::arrow::{ProjectionMask, parquet_to_arrow_schema}; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; use parquet::file::properties::DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT; use partition::expr::PartitionExpr; -use snafu::ResultExt; +use snafu::{OptionExt, ResultExt}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; use store_api::region_request::PathType; +use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{ColumnId, FileId}; use table::predicate::Predicate; @@ -54,6 +55,7 @@ use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst use crate::error::ApplyVectorIndexSnafu; use crate::error::{ ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu, + UnexpectedSnafu, }; use crate::metrics::{ PRECISE_FILTER_ROWS_TOTAL, READ_ROW_GROUPS_TOTAL, READ_ROWS_IN_ROW_GROUP_TOTAL, @@ -83,7 +85,7 @@ use crate::sst::parquet::format::{INTERNAL_COLUMN_NUM, need_override_sequence}; use crate::sst::parquet::json_align::{NestedSchemaAligner, ProjectedRecordBatchStream}; use crate::sst::parquet::metadata::MetadataLoader; use crate::sst::parquet::prefilter::{ - PrefilterContextBuilder, build_reader_filter_plan, execute_prefilter, + CachedPrimaryKeyFilter, PrefilterContextBuilder, build_reader_filter_plan, execute_prefilter, }; use crate::sst::parquet::push_decoder::{ SstParquetRangeFetcher, build_sst_parquet_record_batch_stream, @@ -1836,6 +1838,13 @@ impl RowGroupReaderBuilder { self.prefilter_builder.is_some() } + /// Builds the encoded-primary-key filter selected by the reader filter plan. + pub(crate) fn primary_key_filter(&self) -> Option { + self.prefilter_builder + .as_ref() + .and_then(PrefilterContextBuilder::build_primary_key_filter) + } + /// Builds a parquet record batch stream to read the row group at `row_group_idx`. /// /// If prefiltering is applicable (based on `build_ctx`), this performs a two-phase read: @@ -1894,6 +1903,34 @@ impl RowGroupReaderBuilder { self.make_projected_stream(stream) } + /// Builds a stream that reads only the encoded primary-key column. + /// + /// It preserves the normal reader's binary-or-dictionary decision. This path deliberately + /// skips the normal prefilter pass: the caller reads `__primary_key` once and applies all + /// encoded-primary-key filters to the returned batches. + pub(crate) async fn build_primary_key( + &self, + build_ctx: RowGroupBuildContext<'_>, + ) -> Result { + let parquet_schema = self.parquet_meta.file_metadata().schema_descr(); + let primary_key_index = parquet_schema + .columns() + .iter() + .position(|column| column.name() == PRIMARY_KEY_COLUMN_NAME) + .context(UnexpectedSnafu { + reason: "SST does not contain __primary_key", + })?; + let projection = ProjectionMask::leaves(parquet_schema, [primary_key_index]); + + self.build_with_projection( + build_ctx.row_group_idx, + build_ctx.row_selection, + projection, + build_ctx.fetch_metrics, + ) + .await + } + fn make_projected_stream( &self, stream: ProjectedRecordBatchStream,