Skip to main content

mito2/sst/parquet/
reader.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Parquet reader.
16
17#[cfg(feature = "vector_index")]
18use std::collections::BTreeSet;
19use std::collections::HashSet;
20use std::sync::Arc;
21use std::time::{Duration, Instant};
22
23use api::v1::SemanticType;
24use common_recordbatch::filter::SimpleFilterEvaluator;
25use common_telemetry::{error, tracing, warn};
26use datafusion::physical_plan::PhysicalExpr;
27use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
28use datafusion_expr::utils::expr_to_columns;
29use datafusion_expr::{Expr, Volatility};
30use datatypes::arrow::array::ArrayRef;
31use datatypes::arrow::datatypes::{Field, Schema as ArrowSchema, SchemaRef};
32use datatypes::arrow::record_batch::RecordBatch;
33use datatypes::data_type::ConcreteDataType;
34use datatypes::extension::json::is_json2_extension_type;
35use datatypes::prelude::DataType;
36use futures::StreamExt;
37use mito_codec::row_converter::build_primary_key_codec;
38use object_store::ObjectStore;
39use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions, RowSelection};
40use parquet::arrow::{ProjectionMask, parquet_to_arrow_schema};
41use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
42use parquet::file::properties::DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT;
43use partition::expr::PartitionExpr;
44use snafu::{OptionExt, ResultExt};
45use store_api::codec::PrimaryKeyEncoding;
46use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef};
47use store_api::region_request::PathType;
48use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
49use store_api::storage::{ColumnId, FileId};
50use table::predicate::Predicate;
51
52use crate::cache::index::result_cache::PredicateKey;
53use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta};
54#[cfg(feature = "vector_index")]
55use crate::error::ApplyVectorIndexSnafu;
56use crate::error::{
57    ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu,
58    UnexpectedSnafu,
59};
60use crate::metrics::{
61    PRECISE_FILTER_ROWS_TOTAL, READ_ROW_GROUPS_TOTAL, READ_ROWS_IN_ROW_GROUP_TOTAL,
62    READ_ROWS_TOTAL, READ_STAGE_ELAPSED,
63};
64use crate::read::flat_projection::CompactionProjectionMapper;
65use crate::read::prune::FlatPruneReader;
66use crate::read::read_columns::ReadColumns;
67use crate::sst::file::FileHandle;
68use crate::sst::index::bloom_filter::applier::{
69    BloomFilterIndexApplierRef, BloomFilterIndexApplyMetrics,
70};
71use crate::sst::index::fulltext_index::applier::{
72    FulltextIndexApplierRef, FulltextIndexApplyMetrics,
73};
74use crate::sst::index::inverted_index::applier::{
75    InvertedIndexApplierRef, InvertedIndexApplyMetrics,
76};
77#[cfg(feature = "vector_index")]
78use crate::sst::index::vector_index::applier::VectorIndexApplierRef;
79use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
80use crate::sst::parquet::file_range::{
81    FileRangeContext, FileRangeContextRef, PartitionFilterContext, PreFilterMode, RangeBase,
82};
83use crate::sst::parquet::flat_format::{FlatReadFormat, primary_key_column_index};
84use crate::sst::parquet::format::{INTERNAL_COLUMN_NUM, need_override_sequence};
85use crate::sst::parquet::json_align::{NestedSchemaAligner, ProjectedRecordBatchStream};
86use crate::sst::parquet::metadata::MetadataLoader;
87use crate::sst::parquet::prefilter::{
88    PrefilterContextBuilder, build_reader_filter_plan, execute_prefilter,
89};
90use crate::sst::parquet::push_decoder::{
91    SstParquetRangeFetcher, build_sst_parquet_record_batch_stream,
92};
93use crate::sst::parquet::read_columns::{ProjectionMaskPlan, build_projection_plan};
94use crate::sst::parquet::row_group::ParquetFetchMetrics;
95use crate::sst::parquet::row_selection::RowGroupSelection;
96use crate::sst::parquet::stats::RowGroupPruningStats;
97use crate::sst::{override_pk_field_to_binary, tag_maybe_to_dictionary_field};
98
99const INDEX_TYPE_FULLTEXT: &str = "fulltext";
100
101/// Number of leading row groups sampled by [`should_read_pk_as_binary`].
102const MAX_ROW_GROUPS_TO_CHECK_PK: usize = 4;
103
104/// Returns `true` if the `__primary_key` chunk in any of the first
105/// [`MAX_ROW_GROUPS_TO_CHECK_PK`] row groups exceeds the dictionary page size
106/// limit, signalling the writer likely fell back to plain encoding.
107fn should_read_pk_as_binary(parquet_meta: &ParquetMetaData) -> bool {
108    should_read_pk_as_binary_with_limit(parquet_meta, DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT)
109}
110
111fn should_read_pk_as_binary_with_limit(
112    parquet_meta: &ParquetMetaData,
113    dict_page_size_limit: usize,
114) -> bool {
115    let num_columns = parquet_meta.file_metadata().schema_descr().num_columns();
116    if num_columns < INTERNAL_COLUMN_NUM {
117        return false;
118    }
119    let pk_idx = primary_key_column_index(num_columns);
120    parquet_meta
121        .row_groups()
122        .iter()
123        .take(MAX_ROW_GROUPS_TO_CHECK_PK)
124        .any(|rg| rg.column(pk_idx).uncompressed_size() as usize > dict_page_size_limit)
125}
126const INDEX_TYPE_INVERTED: &str = "inverted";
127const INDEX_TYPE_BLOOM: &str = "bloom filter";
128const INDEX_TYPE_VECTOR: &str = "vector";
129
130macro_rules! handle_index_error {
131    ($err:expr, $file_handle:expr, $index_type:expr) => {
132        if cfg!(any(test, feature = "test")) {
133            panic!(
134                "Failed to apply {} index, region_id: {}, file_id: {}, err: {:?}",
135                $index_type,
136                $file_handle.region_id(),
137                $file_handle.file_id(),
138                $err
139            );
140        } else {
141            warn!(
142                $err; "Failed to apply {} index, region_id: {}, file_id: {}",
143                $index_type,
144                $file_handle.region_id(),
145                $file_handle.file_id()
146            );
147        }
148    };
149}
150
151/// Parquet SST reader builder.
152pub struct ParquetReaderBuilder {
153    /// SST directory.
154    table_dir: String,
155    /// Path type for generating file paths.
156    path_type: PathType,
157    file_handle: FileHandle,
158    object_store: ObjectStore,
159    /// Predicate to push down.
160    predicate: Option<Predicate>,
161    /// The columns to read.
162    ///
163    /// `None` reads all columns. Due to schema change, the projection
164    /// can contain columns not in the parquet file.
165    read_cols: Option<ReadColumns>,
166    /// Strategy to cache SST data.
167    cache_strategy: CacheStrategy,
168    /// Index appliers.
169    inverted_index_appliers: [Option<InvertedIndexApplierRef>; 2],
170    bloom_filter_index_appliers: [Option<BloomFilterIndexApplierRef>; 2],
171    fulltext_index_appliers: [Option<FulltextIndexApplierRef>; 2],
172    /// Vector index applier for KNN search.
173    #[cfg(feature = "vector_index")]
174    vector_index_applier: Option<VectorIndexApplierRef>,
175    /// Over-fetched k for vector index scan.
176    #[cfg(feature = "vector_index")]
177    vector_index_k: Option<usize>,
178    /// Expected metadata of the region while reading the SST.
179    /// This is usually the latest metadata of the region. The reader use
180    /// it get the correct column id of a column by name.
181    expected_metadata: Option<RegionMetadataRef>,
182    /// Whether this reader is for compaction.
183    compaction: bool,
184    /// Mode to pre-filter columns.
185    pre_filter_mode: PreFilterMode,
186    /// Whether to run the reduced-column predicate prefilter pass.
187    enable_predicate_prefilter: bool,
188    /// Whether to decode primary key values eagerly when reading primary key format SSTs.
189    decode_primary_key_values: bool,
190    page_index_policy: PageIndexPolicy,
191    defer_optional_page_index: bool,
192    /// Scan-wide hint for rows in a decoded batch.
193    batch_size: usize,
194}
195
196impl ParquetReaderBuilder {
197    /// Returns a new [ParquetReaderBuilder] to read specific SST.
198    pub fn new(
199        table_dir: String,
200        path_type: PathType,
201        file_handle: FileHandle,
202        object_store: ObjectStore,
203    ) -> ParquetReaderBuilder {
204        ParquetReaderBuilder {
205            table_dir,
206            path_type,
207            file_handle,
208            object_store,
209            predicate: None,
210            read_cols: None,
211            cache_strategy: CacheStrategy::Disabled,
212            inverted_index_appliers: [None, None],
213            bloom_filter_index_appliers: [None, None],
214            fulltext_index_appliers: [None, None],
215            #[cfg(feature = "vector_index")]
216            vector_index_applier: None,
217            #[cfg(feature = "vector_index")]
218            vector_index_k: None,
219            expected_metadata: None,
220            compaction: false,
221            pre_filter_mode: PreFilterMode::All,
222            enable_predicate_prefilter: true,
223            decode_primary_key_values: false,
224            page_index_policy: Default::default(),
225            defer_optional_page_index: false,
226            batch_size: DEFAULT_READ_BATCH_SIZE,
227        }
228    }
229
230    /// Sets the scan-wide hint for rows in a decoded batch.
231    #[must_use]
232    pub(crate) fn batch_size(mut self, batch_size: usize) -> Self {
233        self.batch_size = batch_size.clamp(1, DEFAULT_READ_BATCH_SIZE);
234        self
235    }
236
237    /// Attaches the predicate to the builder.
238    #[must_use]
239    pub fn predicate(mut self, predicate: Option<Predicate>) -> ParquetReaderBuilder {
240        self.predicate = predicate;
241        self
242    }
243
244    /// Attaches the projection to the builder.
245    ///
246    /// The reader only applies the projection to fields.
247    #[must_use]
248    pub fn projection(mut self, read_cols: Option<ReadColumns>) -> ParquetReaderBuilder {
249        self.read_cols = read_cols;
250        self
251    }
252
253    /// Attaches the cache to the builder.
254    #[must_use]
255    pub fn cache(mut self, cache: CacheStrategy) -> ParquetReaderBuilder {
256        self.cache_strategy = cache;
257        self
258    }
259
260    /// Attaches the inverted index appliers to the builder.
261    #[must_use]
262    pub(crate) fn inverted_index_appliers(
263        mut self,
264        index_appliers: [Option<InvertedIndexApplierRef>; 2],
265    ) -> Self {
266        self.inverted_index_appliers = index_appliers;
267        self
268    }
269
270    /// Attaches the bloom filter index appliers to the builder.
271    #[must_use]
272    pub(crate) fn bloom_filter_index_appliers(
273        mut self,
274        index_appliers: [Option<BloomFilterIndexApplierRef>; 2],
275    ) -> Self {
276        self.bloom_filter_index_appliers = index_appliers;
277        self
278    }
279
280    /// Attaches the fulltext index appliers to the builder.
281    #[must_use]
282    pub(crate) fn fulltext_index_appliers(
283        mut self,
284        index_appliers: [Option<FulltextIndexApplierRef>; 2],
285    ) -> Self {
286        self.fulltext_index_appliers = index_appliers;
287        self
288    }
289
290    /// Attaches the vector index applier to the builder.
291    #[cfg(feature = "vector_index")]
292    #[must_use]
293    pub(crate) fn vector_index_applier(
294        mut self,
295        applier: Option<VectorIndexApplierRef>,
296        k: Option<usize>,
297    ) -> Self {
298        self.vector_index_applier = applier;
299        self.vector_index_k = k;
300        self
301    }
302
303    /// Attaches the expected metadata to the builder.
304    #[must_use]
305    pub fn expected_metadata(mut self, expected_metadata: Option<RegionMetadataRef>) -> Self {
306        self.expected_metadata = expected_metadata;
307        self
308    }
309
310    /// Sets the compaction flag.
311    #[must_use]
312    pub fn compaction(mut self, compaction: bool) -> Self {
313        self.compaction = compaction;
314        self
315    }
316
317    /// Sets the pre-filter mode.
318    #[must_use]
319    pub(crate) fn pre_filter_mode(mut self, pre_filter_mode: PreFilterMode) -> Self {
320        self.pre_filter_mode = pre_filter_mode;
321        self
322    }
323
324    /// Sets whether to run the reduced-column predicate prefilter pass.
325    #[must_use]
326    pub(crate) fn enable_predicate_prefilter(mut self, enable: bool) -> Self {
327        self.enable_predicate_prefilter = enable;
328        self
329    }
330
331    /// Decodes primary key values eagerly when reading primary key format SSTs.
332    #[must_use]
333    pub(crate) fn decode_primary_key_values(mut self, decode: bool) -> Self {
334        self.decode_primary_key_values = decode;
335        self
336    }
337
338    #[must_use]
339    pub fn page_index_policy(mut self, page_index_policy: PageIndexPolicy) -> Self {
340        self.page_index_policy = page_index_policy;
341        self
342    }
343
344    /// Defers loading optional page indexes until row-level selections can use them.
345    #[must_use]
346    pub(crate) fn deferred_optional_page_index(mut self) -> Self {
347        self.page_index_policy = PageIndexPolicy::Optional;
348        self.defer_optional_page_index = true;
349        self
350    }
351
352    /// Builds a [ParquetReader].
353    ///
354    /// This needs to perform IO operation.
355    #[tracing::instrument(
356        skip_all,
357        fields(
358            region_id = %self.file_handle.region_id(),
359            file_id = %self.file_handle.file_id()
360        )
361    )]
362    pub async fn build(&self) -> Result<Option<ParquetReader>> {
363        let mut metrics = ReaderMetrics::default();
364
365        let Some((context, selection)) = self.build_reader_input_inner(&mut metrics).await? else {
366            return Ok(None);
367        };
368        ParquetReader::new(Arc::new(context), selection)
369            .await
370            .map(Some)
371    }
372
373    /// Builds a [FileRangeContext] and collects row groups to read.
374    ///
375    /// This needs to perform IO operation.
376    #[tracing::instrument(
377        skip_all,
378        fields(
379            region_id = %self.file_handle.region_id(),
380            file_id = %self.file_handle.file_id()
381        )
382    )]
383    pub async fn build_reader_input(
384        &self,
385        metrics: &mut ReaderMetrics,
386    ) -> Result<Option<(FileRangeContext, RowGroupSelection)>> {
387        self.build_reader_input_inner(metrics).await
388    }
389
390    async fn build_reader_input_inner(
391        &self,
392        metrics: &mut ReaderMetrics,
393    ) -> Result<Option<(FileRangeContext, RowGroupSelection)>> {
394        let start = Instant::now();
395
396        let file_path = self.file_handle.file_path(&self.table_dir, self.path_type);
397        let file_size = self.file_handle.meta_ref().file_size;
398
399        // Loads parquet metadata of the file.
400        let initial_page_index_policy = if self.defer_optional_page_index
401            && self.page_index_policy == PageIndexPolicy::Optional
402        {
403            PageIndexPolicy::Skip
404        } else {
405            self.page_index_policy
406        };
407        let (sst_meta, mut cache_miss) = self
408            .read_parquet_metadata(
409                &file_path,
410                file_size,
411                &mut metrics.metadata_cache_metrics,
412                initial_page_index_policy,
413            )
414            .await?;
415        let mut parquet_meta = sst_meta.parquet_metadata();
416        let mut parquet_metadata_size = sst_meta.parquet_metadata_size();
417        let region_meta = sst_meta.region_metadata();
418        let region_partition_expr_str = self
419            .expected_metadata
420            .as_ref()
421            .and_then(|meta| meta.partition_expr.as_ref())
422            .map(|expr| expr.as_str());
423        let (_, is_same_region_partition) = Self::is_same_region_partition(
424            region_partition_expr_str,
425            self.file_handle.meta_ref().partition_expr.as_ref(),
426        )?;
427        // Skip auto convert when:
428        // - compaction is enabled
429        // - region partition expr is same with file partition expr (no need to auto convert)
430        let skip_auto_convert = self.compaction && is_same_region_partition;
431
432        // Build a compaction projection helper when:
433        // - compaction is enabled
434        // - region partition expr differs from file partition expr
435        // - flat format is enabled
436        // - primary key encoding is sparse
437        //
438        // This is applied after row-group filtering to align batches with flat output schema
439        // before compat handling.
440        let compaction_projection_mapper = if self.compaction
441            && !is_same_region_partition
442            && region_meta.primary_key_encoding == PrimaryKeyEncoding::Sparse
443        {
444            Some(CompactionProjectionMapper::try_new(&region_meta)?)
445        } else {
446            None
447        };
448
449        let read_cols = if let Some(read_cols) = &self.read_cols {
450            read_cols.clone()
451        } else {
452            let expected_meta = self.expected_metadata.as_ref().unwrap_or(&region_meta);
453            // Lists all column ids to read, we always use the expected metadata if possible.
454            ReadColumns::from_deduped_column_ids(
455                expected_meta
456                    .column_metadatas
457                    .iter()
458                    .map(|col| col.column_id),
459            )
460        };
461
462        let file_metadata = parquet_meta.file_metadata();
463        let parquet_schema_desc = file_metadata.schema_descr();
464        let file_schema =
465            parquet_to_arrow_schema(parquet_schema_desc, file_metadata.key_value_metadata())
466                .context(ParquetToArrowSchemaSnafu { file: &file_path })?;
467        let mut read_format = FlatReadFormat::new(
468            region_meta.clone(),
469            read_cols,
470            Some(Arc::new(file_schema)),
471            &file_path,
472            skip_auto_convert,
473        )?;
474        if need_override_sequence(&parquet_meta) {
475            read_format
476                .set_override_sequence(self.file_handle.meta_ref().sequence.map(|x| x.get()));
477        }
478
479        // Computes the projection mask.
480        let parquet_read_cols = read_format.parquet_read_columns();
481        let projection_plan = build_projection_plan(parquet_read_cols, parquet_schema_desc);
482        let has_nested_projection = parquet_read_cols.has_nested();
483        let selection = self
484            .row_groups_to_read(&read_format, &parquet_meta, &mut metrics.filter_metrics)
485            .await;
486
487        if selection.is_empty() {
488            metrics.build_cost += start.elapsed();
489            return Ok(None);
490        }
491
492        let prune_schema = self
493            .expected_metadata
494            .as_ref()
495            .map(|meta| meta.schema.clone())
496            .unwrap_or_else(|| region_meta.schema.clone());
497
498        let dyn_filters = if let Some(predicate) = &self.predicate {
499            predicate.dyn_filters().as_ref().clone()
500        } else {
501            vec![]
502        };
503
504        let codec = build_primary_key_codec(read_format.metadata());
505
506        let filter_plan = build_reader_filter_plan(
507            self.predicate.as_ref(),
508            self.expected_metadata.as_deref(),
509            self.pre_filter_mode,
510            self.enable_predicate_prefilter,
511            &read_format,
512            &codec,
513        );
514
515        if self.defer_optional_page_index
516            && self.page_index_policy == PageIndexPolicy::Optional
517            && (filter_plan.prefilter_builder.is_some()
518                || has_row_level_selection(&selection, &parquet_meta))
519        {
520            let (sst_meta, page_index_cache_miss) = self
521                .read_parquet_metadata(
522                    &file_path,
523                    file_size,
524                    &mut metrics.metadata_cache_metrics,
525                    PageIndexPolicy::Optional,
526                )
527                .await?;
528            parquet_meta = sst_meta.parquet_metadata();
529            parquet_metadata_size = sst_meta.parquet_metadata_size();
530            cache_miss |= page_index_cache_miss;
531        }
532
533        // Trigger background download if metadata had a cache miss and selection is not empty
534        if cache_miss && !selection.is_empty() {
535            use crate::cache::file_cache::{FileType, IndexKey};
536            let index_key = IndexKey::new(
537                self.file_handle.region_id(),
538                self.file_handle.file_id().file_id(),
539                FileType::Parquet,
540            );
541            self.cache_strategy.maybe_download_background(
542                index_key,
543                file_path.clone(),
544                self.object_store.clone(),
545                file_size,
546            );
547        }
548
549        // Create ArrowReaderMetadata for async stream building.
550        let mut arrow_reader_options = ArrowReaderOptions::new();
551        if !read_format
552            .arrow_schema()
553            .fields()
554            .iter()
555            .any(is_json2_extension_type)
556        {
557            // Read `__primary_key` as Binary when it's too large for dictionary
558            // encoding; convert_batch wraps it back to a DictionaryArray.
559            let schema_for_reader = if should_read_pk_as_binary(&parquet_meta) {
560                read_format.set_pk_as_binary()?;
561                override_pk_field_to_binary(read_format.arrow_schema())
562            } else {
563                read_format.arrow_schema().clone()
564            };
565            arrow_reader_options = arrow_reader_options.with_schema(schema_for_reader);
566        }
567        let arrow_metadata =
568            ArrowReaderMetadata::try_new(parquet_meta.clone(), arrow_reader_options)
569                .context(ReadDataPartSnafu)?;
570
571        let output_schema = read_format.output_arrow_schema()?;
572
573        let reader_builder = RowGroupReaderBuilder {
574            file_handle: self.file_handle.clone(),
575            file_path,
576            parquet_meta,
577            parquet_metadata_size,
578            arrow_metadata,
579            output_schema,
580            object_store: self.object_store.clone(),
581            projection: projection_plan,
582            has_nested_projection,
583            cache_strategy: self.cache_strategy.clone(),
584            prefilter_builder: filter_plan.prefilter_builder,
585            batch_size: self.batch_size,
586        };
587
588        let partition_filter = self.build_partition_filter(&read_format, &prune_schema)?;
589
590        let context = FileRangeContext::new(
591            reader_builder,
592            RangeBase {
593                filters: filter_plan.remaining_simple_filters,
594                dyn_filters,
595                read_format,
596                expected_metadata: self.expected_metadata.clone(),
597                prune_schema,
598                codec,
599                compat_batch: None,
600                compaction_projection_mapper,
601                pre_filter_mode: self.pre_filter_mode,
602                partition_filter,
603            },
604        );
605
606        metrics.build_cost += start.elapsed();
607
608        Ok(Some((context, selection)))
609    }
610
611    fn is_same_region_partition(
612        region_partition_expr_str: Option<&str>,
613        file_partition_expr: Option<&PartitionExpr>,
614    ) -> Result<(Option<PartitionExpr>, bool)> {
615        let region_partition_expr = match region_partition_expr_str {
616            Some(expr_str) => crate::region::parse_partition_expr(Some(expr_str))?,
617            None => None,
618        };
619
620        let is_same = region_partition_expr.as_ref() == file_partition_expr;
621        Ok((region_partition_expr, is_same))
622    }
623
624    /// Compare partition expressions from expected metadata and file metadata,
625    /// and build a partition filter if they differ.
626    fn build_partition_filter(
627        &self,
628        read_format: &FlatReadFormat,
629        prune_schema: &Arc<datatypes::schema::Schema>,
630    ) -> Result<Option<PartitionFilterContext>> {
631        let region_partition_expr_str = self
632            .expected_metadata
633            .as_ref()
634            .and_then(|meta| meta.partition_expr.as_ref());
635        let file_partition_expr_ref = self.file_handle.meta_ref().partition_expr.as_ref();
636
637        let (region_partition_expr, is_same_region_partition) = Self::is_same_region_partition(
638            region_partition_expr_str.map(|s| s.as_str()),
639            file_partition_expr_ref,
640        )?;
641
642        if is_same_region_partition {
643            return Ok(None);
644        }
645
646        let Some(region_partition_expr) = region_partition_expr else {
647            return Ok(None);
648        };
649
650        // Collect columns referenced by the partition expression.
651        let mut referenced_columns = HashSet::new();
652        region_partition_expr.collect_column_names(&mut referenced_columns);
653
654        // Build a partition_schema containing only referenced columns.
655        let partition_schema = Arc::new(datatypes::schema::Schema::new(
656            prune_schema
657                .column_schemas()
658                .iter()
659                .filter(|col| referenced_columns.contains(&col.name))
660                .map(|col| {
661                    if let Some(column_meta) = read_format.metadata().column_by_name(&col.name)
662                        && column_meta.semantic_type == SemanticType::Tag
663                        && col.data_type.is_string()
664                    {
665                        let field = Arc::new(Field::new(
666                            &col.name,
667                            col.data_type.as_arrow_type(),
668                            col.is_nullable(),
669                        ));
670                        let dict_field = tag_maybe_to_dictionary_field(&col.data_type, &field);
671                        let mut column = col.clone();
672                        column.data_type =
673                            ConcreteDataType::from_arrow_type(dict_field.data_type());
674                        return column;
675                    }
676
677                    col.clone()
678                })
679                .collect::<Vec<_>>(),
680        ));
681
682        let region_partition_physical_expr = region_partition_expr
683            .try_as_physical_expr(partition_schema.arrow_schema())
684            .context(SerializePartitionExprSnafu)?;
685
686        Ok(Some(PartitionFilterContext {
687            region_partition_physical_expr,
688            partition_schema,
689        }))
690    }
691
692    /// Reads parquet metadata of specific file.
693    /// Returns (fused metadata, cache_miss_flag).
694    pub(crate) async fn read_parquet_metadata(
695        &self,
696        file_path: &str,
697        file_size: u64,
698        cache_metrics: &mut MetadataCacheMetrics,
699        page_index_policy: PageIndexPolicy,
700    ) -> Result<(Arc<CachedSstMeta>, bool)> {
701        let start = Instant::now();
702        let _t = READ_STAGE_ELAPSED
703            .with_label_values(&["read_parquet_metadata"])
704            .start_timer();
705
706        let file_id = self.file_handle.file_id();
707        // Tries to get from cache with metrics tracking.
708        if let Some(metadata) = self
709            .cache_strategy
710            .get_sst_meta_data(file_id, cache_metrics, page_index_policy)
711            .await
712        {
713            cache_metrics.metadata_load_cost += start.elapsed();
714            return Ok((metadata, false));
715        }
716
717        // Cache miss, load metadata directly.
718        let mut metadata_loader =
719            MetadataLoader::new(self.object_store.clone(), file_path, file_size);
720        metadata_loader.with_page_index_policy(page_index_policy);
721        let metadata = metadata_loader.load(cache_metrics).await?;
722
723        let decoded = if self.cache_strategy.sst_meta_cache_enabled() {
724            let metadata = prepare_sst_meta(file_path, metadata, None, page_index_policy).await?;
725            let decoded = metadata.decoded();
726            match metadata {
727                SstMetaPreparation::Prepared(metadata) => {
728                    self.cache_strategy
729                        .put_prepared_sst_meta(file_id, metadata, true);
730                }
731                SstMetaPreparation::DecodedOnly { encoding_error, .. } => {
732                    warn!(
733                        encoding_error;
734                        "Failed to encode SST metadata for cache, using decoded metadata for {}",
735                        file_path
736                    );
737                }
738            }
739            decoded
740        } else {
741            Arc::new(CachedSstMeta::try_new_with_page_index_policy(
742                file_path,
743                metadata,
744                None,
745                page_index_policy,
746            )?)
747        };
748
749        cache_metrics.metadata_load_cost += start.elapsed();
750        Ok((decoded, true))
751    }
752
753    /// Computes row groups to read, along with their respective row selections.
754    #[tracing::instrument(
755        skip_all,
756        fields(
757            region_id = %self.file_handle.region_id(),
758            file_id = %self.file_handle.file_id()
759        )
760    )]
761    async fn row_groups_to_read(
762        &self,
763        read_format: &FlatReadFormat,
764        parquet_meta: &ParquetMetaData,
765        metrics: &mut ReaderFilterMetrics,
766    ) -> RowGroupSelection {
767        let num_row_groups = parquet_meta.num_row_groups();
768        let num_rows = parquet_meta.file_metadata().num_rows();
769        if num_row_groups == 0 || num_rows == 0 {
770            return RowGroupSelection::default();
771        }
772
773        // Let's assume that the number of rows in the first row group
774        // can represent the `row_group_size` of the Parquet file.
775        let row_group_size = parquet_meta.row_group(0).num_rows() as usize;
776        if row_group_size == 0 {
777            return RowGroupSelection::default();
778        }
779
780        metrics.rg_total += num_row_groups;
781        metrics.rows_total += num_rows as usize;
782
783        // Compute skip_fields once for all pruning operations
784        let skip_fields = self.pre_filter_mode.skip_fields();
785
786        let mut output = self.row_groups_by_minmax(
787            read_format,
788            parquet_meta,
789            row_group_size,
790            num_rows as usize,
791            metrics,
792            skip_fields,
793        );
794        if output.is_empty() {
795            return output;
796        }
797
798        let fulltext_filtered = self
799            .prune_row_groups_by_fulltext_index(
800                row_group_size,
801                num_row_groups,
802                &mut output,
803                metrics,
804                skip_fields,
805            )
806            .await;
807        if output.is_empty() {
808            return output;
809        }
810
811        self.prune_row_groups_by_inverted_index(
812            read_format.metadata(),
813            row_group_size,
814            parquet_meta,
815            &mut output,
816            metrics,
817            skip_fields,
818        )
819        .await;
820        if output.is_empty() {
821            return output;
822        }
823
824        self.prune_row_groups_by_bloom_filter(
825            read_format.metadata(),
826            row_group_size,
827            parquet_meta,
828            &mut output,
829            metrics,
830            skip_fields,
831        )
832        .await;
833        if output.is_empty() {
834            return output;
835        }
836
837        if !fulltext_filtered {
838            self.prune_row_groups_by_fulltext_bloom(
839                row_group_size,
840                parquet_meta,
841                &mut output,
842                metrics,
843                skip_fields,
844            )
845            .await;
846        }
847        #[cfg(feature = "vector_index")]
848        {
849            self.prune_row_groups_by_vector_index(
850                row_group_size,
851                num_row_groups,
852                &mut output,
853                metrics,
854            )
855            .await;
856            if output.is_empty() {
857                return output;
858            }
859        }
860        output
861    }
862
863    /// Prunes row groups by fulltext index. Returns `true` if the row groups are pruned.
864    async fn prune_row_groups_by_fulltext_index(
865        &self,
866        row_group_size: usize,
867        num_row_groups: usize,
868        output: &mut RowGroupSelection,
869        metrics: &mut ReaderFilterMetrics,
870        skip_fields: bool,
871    ) -> bool {
872        if !self.file_handle.meta_ref().fulltext_index_available() {
873            return false;
874        }
875
876        let mut pruned = false;
877        // If skip_fields is true, only apply the first applier (for tags).
878        let appliers = if skip_fields {
879            &self.fulltext_index_appliers[..1]
880        } else {
881            &self.fulltext_index_appliers[..]
882        };
883        for index_applier in appliers.iter().flatten() {
884            let predicate_key = index_applier.predicate_key();
885            // Fast path: return early if the result is in the cache.
886            let cached = self
887                .cache_strategy
888                .index_result_cache()
889                .and_then(|cache| cache.get(predicate_key, self.file_handle.file_id().file_id()));
890            if let Some(result) = cached.as_ref()
891                && all_required_row_groups_searched(output, result)
892            {
893                apply_selection_and_update_metrics(output, result, metrics, INDEX_TYPE_FULLTEXT);
894                metrics.fulltext_index_cache_hit += 1;
895                pruned = true;
896                continue;
897            }
898
899            // Slow path: apply the index from the file.
900            metrics.fulltext_index_cache_miss += 1;
901            let file_size_hint = self.file_handle.meta_ref().index_file_size();
902            let apply_res = index_applier
903                .apply_fine(
904                    self.file_handle.index_id(),
905                    Some(file_size_hint),
906                    metrics.fulltext_index_apply_metrics.as_mut(),
907                )
908                .await;
909            let selection = match apply_res {
910                Ok(Some(res)) => {
911                    RowGroupSelection::from_row_ids(res, row_group_size, num_row_groups)
912                }
913                Ok(None) => continue,
914                Err(err) => {
915                    handle_index_error!(err, self.file_handle, INDEX_TYPE_FULLTEXT);
916                    continue;
917                }
918            };
919
920            self.apply_index_result_and_update_cache(
921                predicate_key,
922                self.file_handle.file_id().file_id(),
923                selection,
924                output,
925                metrics,
926                INDEX_TYPE_FULLTEXT,
927            );
928            pruned = true;
929        }
930        pruned
931    }
932
933    /// Applies index to prune row groups.
934    ///
935    /// TODO(zhongzc): Devise a mechanism to enforce the non-use of indices
936    /// as an escape route in case of index issues, and it can be used to test
937    /// the correctness of the index.
938    async fn prune_row_groups_by_inverted_index(
939        &self,
940        sst_metadata: &RegionMetadataRef,
941        row_group_size: usize,
942        parquet_meta: &ParquetMetaData,
943        output: &mut RowGroupSelection,
944        metrics: &mut ReaderFilterMetrics,
945        skip_fields: bool,
946    ) -> bool {
947        if !self.file_handle.meta_ref().inverted_index_available() {
948            return false;
949        }
950
951        let num_row_groups = parquet_meta.num_row_groups();
952        let total_row_count = parquet_meta.file_metadata().num_rows() as usize;
953        let mut pruned = false;
954        // If skip_fields is true, only apply the first applier (for tags).
955        let appliers = if skip_fields {
956            &self.inverted_index_appliers[..1]
957        } else {
958            &self.inverted_index_appliers[..]
959        };
960        for index_applier in appliers.iter().flatten() {
961            let Ok(Some(plan)) = index_applier
962                .plan_for_sst(sst_metadata)
963                .inspect_err(|e| warn!(e; "failed to build compatible plan for sst"))
964            else {
965                continue;
966            };
967
968            // Fast path: return early if the result is in the cache.
969            let cached = self.cache_strategy.index_result_cache().and_then(|cache| {
970                let file_id = self.file_handle.file_id().file_id();
971                cache.get(&plan.predicate_key, file_id)
972            });
973
974            if let Some(result) = cached.as_ref()
975                && all_required_row_groups_searched(output, result)
976            {
977                apply_selection_and_update_metrics(output, result, metrics, INDEX_TYPE_INVERTED);
978                metrics.inverted_index_cache_hit += 1;
979                pruned = true;
980                continue;
981            }
982
983            // Slow path: apply the index from the file.
984            metrics.inverted_index_cache_miss += 1;
985            let file_size_hint = self.file_handle.meta_ref().index_file_size();
986            let apply_res = index_applier
987                .apply(
988                    self.file_handle.index_id(),
989                    Some(file_size_hint),
990                    &plan.index_applier,
991                    metrics.inverted_index_apply_metrics.as_mut(),
992                )
993                .await;
994
995            let selection = match apply_res {
996                Ok(apply_output) => {
997                    let index_row_count = apply_output.total_row_count;
998                    let Some(selection) = RowGroupSelection::from_inverted_index_apply_output(
999                        row_group_size,
1000                        num_row_groups,
1001                        total_row_count,
1002                        apply_output,
1003                    ) else {
1004                        warn!(
1005                            "Ignore inverted index with mismatched row count, file_id: {:?}, index_row_count: {}, parquet_row_count: {}",
1006                            self.file_handle.file_id(),
1007                            index_row_count,
1008                            total_row_count,
1009                        );
1010                        continue;
1011                    };
1012                    selection
1013                }
1014                Err(err) => {
1015                    handle_index_error!(err, self.file_handle, INDEX_TYPE_INVERTED);
1016                    continue;
1017                }
1018            };
1019
1020            self.apply_index_result_and_update_cache(
1021                &plan.predicate_key,
1022                self.file_handle.file_id().file_id(),
1023                selection,
1024                output,
1025                metrics,
1026                INDEX_TYPE_INVERTED,
1027            );
1028            pruned = true;
1029        }
1030        pruned
1031    }
1032
1033    async fn prune_row_groups_by_bloom_filter(
1034        &self,
1035        sst_metadata: &RegionMetadataRef,
1036        row_group_size: usize,
1037        parquet_meta: &ParquetMetaData,
1038        output: &mut RowGroupSelection,
1039        metrics: &mut ReaderFilterMetrics,
1040        skip_fields: bool,
1041    ) -> bool {
1042        if !self.file_handle.meta_ref().bloom_filter_index_available() {
1043            return false;
1044        }
1045
1046        let mut pruned = false;
1047        // If skip_fields is true, only apply the first applier (for tags).
1048        let appliers = if skip_fields {
1049            &self.bloom_filter_index_appliers[..1]
1050        } else {
1051            &self.bloom_filter_index_appliers[..]
1052        };
1053        for index_applier in appliers.iter().flatten() {
1054            let Some(compatible_predicates) =
1055                index_applier.compatible_predicate_for_sst(sst_metadata)
1056            else {
1057                continue;
1058            };
1059            let predicate_key = PredicateKey::new_bloom(compatible_predicates.clone());
1060            // Fast path: return early if the result is in the cache.
1061            let cached = self.cache_strategy.index_result_cache().and_then(|cache| {
1062                let file_id = self.file_handle.file_id().file_id();
1063                cache.get(&predicate_key, file_id)
1064            });
1065            if let Some(result) = cached.as_ref()
1066                && all_required_row_groups_searched(output, result)
1067            {
1068                apply_selection_and_update_metrics(output, result, metrics, INDEX_TYPE_BLOOM);
1069                metrics.bloom_filter_cache_hit += 1;
1070                pruned = true;
1071                continue;
1072            }
1073
1074            // Slow path: apply the index from the file.
1075            metrics.bloom_filter_cache_miss += 1;
1076            let file_size_hint = self.file_handle.meta_ref().index_file_size();
1077            let rgs = parquet_meta.row_groups().iter().enumerate().map(|(i, rg)| {
1078                (
1079                    rg.num_rows() as usize,
1080                    // Optimize: only search the row group that required by `output` and not stored in `cached`.
1081                    output.contains_non_empty_row_group(i)
1082                        && cached
1083                            .as_ref()
1084                            .map(|c| !c.contains_row_group(i))
1085                            .unwrap_or(true),
1086                )
1087            });
1088            let apply_res = index_applier
1089                .apply(
1090                    self.file_handle.index_id(),
1091                    Some(file_size_hint),
1092                    &compatible_predicates,
1093                    rgs,
1094                    metrics.bloom_filter_apply_metrics.as_mut(),
1095                )
1096                .await;
1097            let mut selection = match apply_res {
1098                Ok(apply_output) => {
1099                    RowGroupSelection::from_row_ranges(apply_output, row_group_size)
1100                }
1101                Err(err) => {
1102                    handle_index_error!(err, self.file_handle, INDEX_TYPE_BLOOM);
1103                    continue;
1104                }
1105            };
1106
1107            // New searched row groups are added to `selection`, concat them with `cached`.
1108            if let Some(cached) = cached.as_ref() {
1109                selection.concat(cached);
1110            }
1111
1112            self.apply_index_result_and_update_cache(
1113                &predicate_key,
1114                self.file_handle.file_id().file_id(),
1115                selection,
1116                output,
1117                metrics,
1118                INDEX_TYPE_BLOOM,
1119            );
1120            pruned = true;
1121        }
1122        pruned
1123    }
1124
1125    /// Prunes row groups by vector index results.
1126    #[cfg(feature = "vector_index")]
1127    async fn prune_row_groups_by_vector_index(
1128        &self,
1129        row_group_size: usize,
1130        num_row_groups: usize,
1131        output: &mut RowGroupSelection,
1132        metrics: &mut ReaderFilterMetrics,
1133    ) {
1134        let Some(applier) = &self.vector_index_applier else {
1135            return;
1136        };
1137        let Some(k) = self.vector_index_k else {
1138            return;
1139        };
1140        if !self.file_handle.meta_ref().vector_index_available() {
1141            return;
1142        }
1143
1144        let file_size_hint = self.file_handle.meta_ref().index_file_size();
1145        let apply_res = applier
1146            .apply_with_k(self.file_handle.index_id(), Some(file_size_hint), k)
1147            .await;
1148        let row_ids = match apply_res {
1149            Ok(res) => res.row_offsets,
1150            Err(err) => {
1151                handle_index_error!(err, self.file_handle, INDEX_TYPE_VECTOR);
1152                return;
1153            }
1154        };
1155
1156        let selection = match vector_selection_from_offsets(row_ids, row_group_size, num_row_groups)
1157        {
1158            Ok(selection) => selection,
1159            Err(err) => {
1160                handle_index_error!(err, self.file_handle, INDEX_TYPE_VECTOR);
1161                return;
1162            }
1163        };
1164        metrics.rows_vector_selected += selection.row_count();
1165        apply_selection_and_update_metrics(output, &selection, metrics, INDEX_TYPE_VECTOR);
1166    }
1167
1168    async fn prune_row_groups_by_fulltext_bloom(
1169        &self,
1170        row_group_size: usize,
1171        parquet_meta: &ParquetMetaData,
1172        output: &mut RowGroupSelection,
1173        metrics: &mut ReaderFilterMetrics,
1174        skip_fields: bool,
1175    ) -> bool {
1176        if !self.file_handle.meta_ref().fulltext_index_available() {
1177            return false;
1178        }
1179
1180        let mut pruned = false;
1181        // If skip_fields is true, only apply the first applier (for tags).
1182        let appliers = if skip_fields {
1183            &self.fulltext_index_appliers[..1]
1184        } else {
1185            &self.fulltext_index_appliers[..]
1186        };
1187        for index_applier in appliers.iter().flatten() {
1188            let predicate_key = index_applier.predicate_key();
1189            // Fast path: return early if the result is in the cache.
1190            let cached = self
1191                .cache_strategy
1192                .index_result_cache()
1193                .and_then(|cache| cache.get(predicate_key, self.file_handle.file_id().file_id()));
1194            if let Some(result) = cached.as_ref()
1195                && all_required_row_groups_searched(output, result)
1196            {
1197                apply_selection_and_update_metrics(output, result, metrics, INDEX_TYPE_FULLTEXT);
1198                metrics.fulltext_index_cache_hit += 1;
1199                pruned = true;
1200                continue;
1201            }
1202
1203            // Slow path: apply the index from the file.
1204            metrics.fulltext_index_cache_miss += 1;
1205            let file_size_hint = self.file_handle.meta_ref().index_file_size();
1206            let rgs = parquet_meta.row_groups().iter().enumerate().map(|(i, rg)| {
1207                (
1208                    rg.num_rows() as usize,
1209                    // Optimize: only search the row group that required by `output` and not stored in `cached`.
1210                    output.contains_non_empty_row_group(i)
1211                        && cached
1212                            .as_ref()
1213                            .map(|c| !c.contains_row_group(i))
1214                            .unwrap_or(true),
1215                )
1216            });
1217            let apply_res = index_applier
1218                .apply_coarse(
1219                    self.file_handle.index_id(),
1220                    Some(file_size_hint),
1221                    rgs,
1222                    metrics.fulltext_index_apply_metrics.as_mut(),
1223                )
1224                .await;
1225            let mut selection = match apply_res {
1226                Ok(Some(apply_output)) => {
1227                    RowGroupSelection::from_row_ranges(apply_output, row_group_size)
1228                }
1229                Ok(None) => continue,
1230                Err(err) => {
1231                    handle_index_error!(err, self.file_handle, INDEX_TYPE_FULLTEXT);
1232                    continue;
1233                }
1234            };
1235
1236            // New searched row groups are added to `selection`, concat them with `cached`.
1237            if let Some(cached) = cached.as_ref() {
1238                selection.concat(cached);
1239            }
1240
1241            self.apply_index_result_and_update_cache(
1242                predicate_key,
1243                self.file_handle.file_id().file_id(),
1244                selection,
1245                output,
1246                metrics,
1247                INDEX_TYPE_FULLTEXT,
1248            );
1249            pruned = true;
1250        }
1251        pruned
1252    }
1253
1254    /// Computes row groups selection after min-max pruning.
1255    fn row_groups_by_minmax(
1256        &self,
1257        read_format: &FlatReadFormat,
1258        parquet_meta: &ParquetMetaData,
1259        row_group_size: usize,
1260        total_row_count: usize,
1261        metrics: &mut ReaderFilterMetrics,
1262        skip_fields: bool,
1263    ) -> RowGroupSelection {
1264        let Some(predicate) = &self.predicate else {
1265            return RowGroupSelection::new(row_group_size, total_row_count);
1266        };
1267
1268        let file_id = self.file_handle.file_id().file_id();
1269        let index_result_cache = self.cache_strategy.index_result_cache();
1270        let cached_minmax_key =
1271            if index_result_cache.is_some() && predicate.dyn_filters().is_empty() {
1272                // Cache min-max pruning results keyed by predicate expressions. This avoids repeatedly
1273                // building row-group pruning stats for identical predicates across queries.
1274                let mut exprs = predicate
1275                    .exprs()
1276                    .iter()
1277                    .map(|expr| format!("{expr:?}"))
1278                    .collect::<Vec<_>>();
1279                exprs.sort();
1280                let schema_version = self
1281                    .expected_metadata
1282                    .as_ref()
1283                    .map(|meta| meta.schema_version)
1284                    .unwrap_or_else(|| read_format.metadata().schema_version);
1285                Some(PredicateKey::new_minmax(
1286                    Arc::new(exprs),
1287                    schema_version,
1288                    skip_fields,
1289                ))
1290            } else {
1291                None
1292            };
1293
1294        if let Some(index_result_cache) = index_result_cache
1295            && let Some(predicate_key) = cached_minmax_key.as_ref()
1296        {
1297            if let Some(result) = index_result_cache.get(predicate_key, file_id) {
1298                metrics.minmax_cache_hit += 1;
1299                let num_row_groups = parquet_meta.num_row_groups();
1300                metrics.rg_minmax_filtered +=
1301                    num_row_groups.saturating_sub(result.row_group_count());
1302                return (*result).clone();
1303            }
1304
1305            metrics.minmax_cache_miss += 1;
1306        }
1307
1308        let region_meta = read_format.metadata();
1309        let row_groups = parquet_meta.row_groups();
1310        let stats = RowGroupPruningStats::new(
1311            row_groups,
1312            read_format,
1313            self.expected_metadata.clone(),
1314            skip_fields,
1315        );
1316        let prune_schema = self
1317            .expected_metadata
1318            .as_ref()
1319            .map(|meta| meta.schema.arrow_schema())
1320            .unwrap_or_else(|| region_meta.schema.arrow_schema());
1321
1322        // Here we use the schema of the SST to build the physical expression. If the column
1323        // in the SST doesn't have the same column id as the column in the expected metadata,
1324        // we will get a None statistics for that column.
1325        let mask = predicate.prune_with_stats(&stats, prune_schema);
1326        let output = RowGroupSelection::from_full_row_group_ids(
1327            mask.iter()
1328                .enumerate()
1329                .filter_map(|(row_group, keep)| keep.then_some(row_group)),
1330            row_group_size,
1331            total_row_count,
1332        );
1333
1334        metrics.rg_minmax_filtered += parquet_meta
1335            .num_row_groups()
1336            .saturating_sub(output.row_group_count());
1337
1338        if let Some(index_result_cache) = index_result_cache
1339            && let Some(predicate_key) = cached_minmax_key
1340        {
1341            index_result_cache.put(predicate_key, file_id, Arc::new(output.clone()));
1342        }
1343
1344        output
1345    }
1346
1347    fn apply_index_result_and_update_cache(
1348        &self,
1349        predicate_key: &PredicateKey,
1350        file_id: FileId,
1351        result: RowGroupSelection,
1352        output: &mut RowGroupSelection,
1353        metrics: &mut ReaderFilterMetrics,
1354        index_type: &str,
1355    ) {
1356        apply_selection_and_update_metrics(output, &result, metrics, index_type);
1357
1358        if let Some(index_result_cache) = &self.cache_strategy.index_result_cache() {
1359            index_result_cache.put(predicate_key.clone(), file_id, Arc::new(result));
1360        }
1361    }
1362}
1363
1364fn has_row_level_selection(selection: &RowGroupSelection, parquet_meta: &ParquetMetaData) -> bool {
1365    selection.iter().any(|(row_group_idx, row_selection)| {
1366        let Some(row_group) = parquet_meta.row_groups().get(*row_group_idx) else {
1367            return false;
1368        };
1369
1370        row_selection.row_count() != row_group.num_rows() as usize
1371            || row_selection.iter().any(|selector| selector.skip)
1372    })
1373}
1374
1375fn apply_selection_and_update_metrics(
1376    output: &mut RowGroupSelection,
1377    result: &RowGroupSelection,
1378    metrics: &mut ReaderFilterMetrics,
1379    index_type: &str,
1380) {
1381    let intersection = output.intersect(result);
1382
1383    let row_group_count = output.row_group_count() - intersection.row_group_count();
1384    let row_count = output.row_count() - intersection.row_count();
1385
1386    metrics.update_index_metrics(index_type, row_group_count, row_count);
1387
1388    *output = intersection;
1389}
1390
1391#[cfg(feature = "vector_index")]
1392fn vector_selection_from_offsets(
1393    row_offsets: Vec<u64>,
1394    row_group_size: usize,
1395    num_row_groups: usize,
1396) -> Result<RowGroupSelection> {
1397    let mut row_ids = BTreeSet::new();
1398    for offset in row_offsets {
1399        let row_id = u32::try_from(offset).map_err(|_| {
1400            ApplyVectorIndexSnafu {
1401                reason: format!("Row offset {} exceeds u32::MAX", offset),
1402            }
1403            .build()
1404        })?;
1405        row_ids.insert(row_id);
1406    }
1407    Ok(RowGroupSelection::from_row_ids(
1408        row_ids,
1409        row_group_size,
1410        num_row_groups,
1411    ))
1412}
1413
1414fn all_required_row_groups_searched(
1415    required_row_groups: &RowGroupSelection,
1416    cached_row_groups: &RowGroupSelection,
1417) -> bool {
1418    required_row_groups.iter().all(|(rg_id, _)| {
1419        // Row group with no rows is not required to search.
1420        !required_row_groups.contains_non_empty_row_group(*rg_id)
1421            // The row group is already searched.
1422            || cached_row_groups.contains_row_group(*rg_id)
1423    })
1424}
1425
1426/// Metrics of filtering rows groups and rows.
1427#[derive(Debug, Default, Clone)]
1428pub(crate) struct ReaderFilterMetrics {
1429    /// Number of row groups before filtering.
1430    pub(crate) rg_total: usize,
1431    /// Number of row groups filtered by fulltext index.
1432    pub(crate) rg_fulltext_filtered: usize,
1433    /// Number of row groups filtered by inverted index.
1434    pub(crate) rg_inverted_filtered: usize,
1435    /// Number of row groups filtered by min-max index.
1436    pub(crate) rg_minmax_filtered: usize,
1437    /// Number of row groups filtered by bloom filter index.
1438    pub(crate) rg_bloom_filtered: usize,
1439    /// Number of row groups filtered by vector index.
1440    pub(crate) rg_vector_filtered: usize,
1441
1442    /// Number of rows in row group before filtering.
1443    pub(crate) rows_total: usize,
1444    /// Number of rows in row group filtered by fulltext index.
1445    pub(crate) rows_fulltext_filtered: usize,
1446    /// Number of rows in row group filtered by inverted index.
1447    pub(crate) rows_inverted_filtered: usize,
1448    /// Number of rows in row group filtered by bloom filter index.
1449    pub(crate) rows_bloom_filtered: usize,
1450    /// Number of rows filtered by vector index.
1451    pub(crate) rows_vector_filtered: usize,
1452    /// Number of rows selected by vector index.
1453    pub(crate) rows_vector_selected: usize,
1454    /// Number of rows filtered by precise filter.
1455    pub(crate) rows_precise_filtered: usize,
1456
1457    /// Number of index result cache hits for fulltext index.
1458    pub(crate) fulltext_index_cache_hit: usize,
1459    /// Number of index result cache misses for fulltext index.
1460    pub(crate) fulltext_index_cache_miss: usize,
1461    /// Number of index result cache hits for inverted index.
1462    pub(crate) inverted_index_cache_hit: usize,
1463    /// Number of index result cache misses for inverted index.
1464    pub(crate) inverted_index_cache_miss: usize,
1465    /// Number of index result cache hits for bloom filter index.
1466    pub(crate) bloom_filter_cache_hit: usize,
1467    /// Number of index result cache misses for bloom filter index.
1468    pub(crate) bloom_filter_cache_miss: usize,
1469    /// Number of index result cache hits for minmax pruning.
1470    pub(crate) minmax_cache_hit: usize,
1471    /// Number of index result cache misses for minmax pruning.
1472    pub(crate) minmax_cache_miss: usize,
1473
1474    /// Optional metrics for inverted index applier.
1475    pub(crate) inverted_index_apply_metrics: Option<InvertedIndexApplyMetrics>,
1476    /// Optional metrics for bloom filter index applier.
1477    pub(crate) bloom_filter_apply_metrics: Option<BloomFilterIndexApplyMetrics>,
1478    /// Optional metrics for fulltext index applier.
1479    pub(crate) fulltext_index_apply_metrics: Option<FulltextIndexApplyMetrics>,
1480
1481    /// Number of pruner builder cache hits.
1482    pub(crate) pruner_cache_hit: usize,
1483    /// Number of pruner builder cache misses.
1484    pub(crate) pruner_cache_miss: usize,
1485    /// Duration spent waiting for pruner to build file ranges.
1486    pub(crate) pruner_prune_cost: Duration,
1487    /// Number of files filtered by manifest time-range pruning.
1488    pub(crate) files_time_range_pruned: usize,
1489}
1490
1491impl ReaderFilterMetrics {
1492    /// Adds `other` metrics to this metrics.
1493    pub(crate) fn merge_from(&mut self, other: &ReaderFilterMetrics) {
1494        self.rg_total += other.rg_total;
1495        self.rg_fulltext_filtered += other.rg_fulltext_filtered;
1496        self.rg_inverted_filtered += other.rg_inverted_filtered;
1497        self.rg_minmax_filtered += other.rg_minmax_filtered;
1498        self.rg_bloom_filtered += other.rg_bloom_filtered;
1499        self.rg_vector_filtered += other.rg_vector_filtered;
1500
1501        self.rows_total += other.rows_total;
1502        self.rows_fulltext_filtered += other.rows_fulltext_filtered;
1503        self.rows_inverted_filtered += other.rows_inverted_filtered;
1504        self.rows_bloom_filtered += other.rows_bloom_filtered;
1505        self.rows_vector_filtered += other.rows_vector_filtered;
1506        self.rows_vector_selected += other.rows_vector_selected;
1507        self.rows_precise_filtered += other.rows_precise_filtered;
1508
1509        self.fulltext_index_cache_hit += other.fulltext_index_cache_hit;
1510        self.fulltext_index_cache_miss += other.fulltext_index_cache_miss;
1511        self.inverted_index_cache_hit += other.inverted_index_cache_hit;
1512        self.inverted_index_cache_miss += other.inverted_index_cache_miss;
1513        self.bloom_filter_cache_hit += other.bloom_filter_cache_hit;
1514        self.bloom_filter_cache_miss += other.bloom_filter_cache_miss;
1515        self.minmax_cache_hit += other.minmax_cache_hit;
1516        self.minmax_cache_miss += other.minmax_cache_miss;
1517
1518        self.pruner_cache_hit += other.pruner_cache_hit;
1519        self.pruner_cache_miss += other.pruner_cache_miss;
1520        self.pruner_prune_cost += other.pruner_prune_cost;
1521        self.files_time_range_pruned += other.files_time_range_pruned;
1522
1523        // Merge optional applier metrics
1524        if let Some(other_metrics) = &other.inverted_index_apply_metrics {
1525            self.inverted_index_apply_metrics
1526                .get_or_insert_with(Default::default)
1527                .merge_from(other_metrics);
1528        }
1529        if let Some(other_metrics) = &other.bloom_filter_apply_metrics {
1530            self.bloom_filter_apply_metrics
1531                .get_or_insert_with(Default::default)
1532                .merge_from(other_metrics);
1533        }
1534        if let Some(other_metrics) = &other.fulltext_index_apply_metrics {
1535            self.fulltext_index_apply_metrics
1536                .get_or_insert_with(Default::default)
1537                .merge_from(other_metrics);
1538        }
1539    }
1540
1541    /// Reports metrics.
1542    pub(crate) fn observe(&self) {
1543        READ_ROW_GROUPS_TOTAL
1544            .with_label_values(&["before_filtering"])
1545            .inc_by(self.rg_total as u64);
1546        READ_ROW_GROUPS_TOTAL
1547            .with_label_values(&["fulltext_index_filtered"])
1548            .inc_by(self.rg_fulltext_filtered as u64);
1549        READ_ROW_GROUPS_TOTAL
1550            .with_label_values(&["inverted_index_filtered"])
1551            .inc_by(self.rg_inverted_filtered as u64);
1552        READ_ROW_GROUPS_TOTAL
1553            .with_label_values(&["minmax_index_filtered"])
1554            .inc_by(self.rg_minmax_filtered as u64);
1555        READ_ROW_GROUPS_TOTAL
1556            .with_label_values(&["bloom_filter_index_filtered"])
1557            .inc_by(self.rg_bloom_filtered as u64);
1558        READ_ROW_GROUPS_TOTAL
1559            .with_label_values(&["vector_index_filtered"])
1560            .inc_by(self.rg_vector_filtered as u64);
1561
1562        PRECISE_FILTER_ROWS_TOTAL
1563            .with_label_values(&["parquet"])
1564            .inc_by(self.rows_precise_filtered as u64);
1565        READ_ROWS_IN_ROW_GROUP_TOTAL
1566            .with_label_values(&["before_filtering"])
1567            .inc_by(self.rows_total as u64);
1568        READ_ROWS_IN_ROW_GROUP_TOTAL
1569            .with_label_values(&["fulltext_index_filtered"])
1570            .inc_by(self.rows_fulltext_filtered as u64);
1571        READ_ROWS_IN_ROW_GROUP_TOTAL
1572            .with_label_values(&["inverted_index_filtered"])
1573            .inc_by(self.rows_inverted_filtered as u64);
1574        READ_ROWS_IN_ROW_GROUP_TOTAL
1575            .with_label_values(&["bloom_filter_index_filtered"])
1576            .inc_by(self.rows_bloom_filtered as u64);
1577        READ_ROWS_IN_ROW_GROUP_TOTAL
1578            .with_label_values(&["vector_index_filtered"])
1579            .inc_by(self.rows_vector_filtered as u64);
1580    }
1581
1582    fn update_index_metrics(&mut self, index_type: &str, row_group_count: usize, row_count: usize) {
1583        match index_type {
1584            INDEX_TYPE_FULLTEXT => {
1585                self.rg_fulltext_filtered += row_group_count;
1586                self.rows_fulltext_filtered += row_count;
1587            }
1588            INDEX_TYPE_INVERTED => {
1589                self.rg_inverted_filtered += row_group_count;
1590                self.rows_inverted_filtered += row_count;
1591            }
1592            INDEX_TYPE_BLOOM => {
1593                self.rg_bloom_filtered += row_group_count;
1594                self.rows_bloom_filtered += row_count;
1595            }
1596            INDEX_TYPE_VECTOR => {
1597                self.rg_vector_filtered += row_group_count;
1598                self.rows_vector_filtered += row_count;
1599            }
1600            _ => {}
1601        }
1602    }
1603}
1604
1605#[cfg(all(test, feature = "vector_index"))]
1606mod vector_index_tests {
1607    use super::*;
1608
1609    #[test]
1610    fn test_vector_selection_from_offsets() {
1611        let row_group_size = 4;
1612        let num_row_groups = 3;
1613        let selection =
1614            vector_selection_from_offsets(vec![0, 1, 5, 9], row_group_size, num_row_groups)
1615                .unwrap();
1616
1617        assert_eq!(selection.row_group_count(), 3);
1618        assert_eq!(selection.row_count(), 4);
1619        assert!(selection.contains_non_empty_row_group(0));
1620        assert!(selection.contains_non_empty_row_group(1));
1621        assert!(selection.contains_non_empty_row_group(2));
1622    }
1623
1624    #[test]
1625    fn test_vector_selection_from_offsets_out_of_range() {
1626        let row_group_size = 4;
1627        let num_row_groups = 2;
1628        let selection = vector_selection_from_offsets(
1629            vec![0, 7, u64::from(u32::MAX) + 1],
1630            row_group_size,
1631            num_row_groups,
1632        );
1633        assert!(selection.is_err());
1634    }
1635
1636    #[test]
1637    fn test_vector_selection_updates_metrics() {
1638        let row_group_size = 4;
1639        let total_rows = 8;
1640        let mut output = RowGroupSelection::new(row_group_size, total_rows);
1641        let selection = vector_selection_from_offsets(vec![1], row_group_size, 2).unwrap();
1642        let mut metrics = ReaderFilterMetrics::default();
1643
1644        apply_selection_and_update_metrics(
1645            &mut output,
1646            &selection,
1647            &mut metrics,
1648            INDEX_TYPE_VECTOR,
1649        );
1650
1651        assert_eq!(metrics.rg_vector_filtered, 1);
1652        assert_eq!(metrics.rows_vector_filtered, 7);
1653        assert_eq!(output.row_count(), 1);
1654    }
1655}
1656
1657/// Metrics for parquet metadata cache operations.
1658#[derive(Default, Clone, Copy)]
1659pub struct MetadataCacheMetrics {
1660    /// Number of memory cache hits for parquet metadata.
1661    pub mem_cache_hit: usize,
1662    /// Number of file cache hits for parquet metadata.
1663    pub file_cache_hit: usize,
1664    /// Number of cache misses for parquet metadata.
1665    pub cache_miss: usize,
1666    /// Duration to load parquet metadata.
1667    pub metadata_load_cost: Duration,
1668    /// Number of read operations performed.
1669    pub num_reads: usize,
1670    /// Total bytes read from storage.
1671    pub bytes_read: u64,
1672}
1673
1674impl std::fmt::Debug for MetadataCacheMetrics {
1675    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1676        let Self {
1677            mem_cache_hit,
1678            file_cache_hit,
1679            cache_miss,
1680            metadata_load_cost,
1681            num_reads,
1682            bytes_read,
1683        } = self;
1684
1685        if self.is_empty() {
1686            return write!(f, "{{}}");
1687        }
1688        write!(f, "{{")?;
1689
1690        write!(f, "\"metadata_load_cost\":\"{:?}\"", metadata_load_cost)?;
1691
1692        if *mem_cache_hit > 0 {
1693            write!(f, ", \"mem_cache_hit\":{}", mem_cache_hit)?;
1694        }
1695        if *file_cache_hit > 0 {
1696            write!(f, ", \"file_cache_hit\":{}", file_cache_hit)?;
1697        }
1698        if *cache_miss > 0 {
1699            write!(f, ", \"cache_miss\":{}", cache_miss)?;
1700        }
1701        if *num_reads > 0 {
1702            write!(f, ", \"num_reads\":{}", num_reads)?;
1703        }
1704        if *bytes_read > 0 {
1705            write!(f, ", \"bytes_read\":{}", bytes_read)?;
1706        }
1707
1708        write!(f, "}}")
1709    }
1710}
1711
1712impl MetadataCacheMetrics {
1713    /// Returns true if the metrics are empty (contain no meaningful data).
1714    pub(crate) fn is_empty(&self) -> bool {
1715        self.metadata_load_cost.is_zero()
1716    }
1717
1718    /// Adds `other` metrics to this metrics.
1719    pub(crate) fn merge_from(&mut self, other: &MetadataCacheMetrics) {
1720        self.mem_cache_hit += other.mem_cache_hit;
1721        self.file_cache_hit += other.file_cache_hit;
1722        self.cache_miss += other.cache_miss;
1723        self.metadata_load_cost += other.metadata_load_cost;
1724        self.num_reads += other.num_reads;
1725        self.bytes_read += other.bytes_read;
1726    }
1727}
1728
1729/// Parquet reader metrics.
1730#[derive(Debug, Default, Clone)]
1731pub struct ReaderMetrics {
1732    /// Filtered row groups and rows metrics.
1733    pub(crate) filter_metrics: ReaderFilterMetrics,
1734    /// Duration to build the parquet reader.
1735    pub(crate) build_cost: Duration,
1736    /// Duration to scan the reader.
1737    pub(crate) scan_cost: Duration,
1738    /// Number of record batches read.
1739    pub(crate) num_record_batches: usize,
1740    /// Number of batches decoded.
1741    pub(crate) num_batches: usize,
1742    /// Number of rows read.
1743    pub(crate) num_rows: usize,
1744    /// Metrics for parquet metadata cache.
1745    pub(crate) metadata_cache_metrics: MetadataCacheMetrics,
1746    /// Optional metrics for page/row group fetch operations.
1747    pub(crate) fetch_metrics: Option<Arc<ParquetFetchMetrics>>,
1748    /// Memory size of metadata loaded for building file ranges.
1749    pub(crate) metadata_mem_size: isize,
1750    /// Number of file range builders created.
1751    pub(crate) num_range_builders: isize,
1752}
1753
1754impl ReaderMetrics {
1755    /// Adds `other` metrics to this metrics.
1756    pub(crate) fn merge_from(&mut self, other: &ReaderMetrics) {
1757        self.filter_metrics.merge_from(&other.filter_metrics);
1758        self.build_cost += other.build_cost;
1759        self.scan_cost += other.scan_cost;
1760        self.num_record_batches += other.num_record_batches;
1761        self.num_batches += other.num_batches;
1762        self.num_rows += other.num_rows;
1763        self.metadata_cache_metrics
1764            .merge_from(&other.metadata_cache_metrics);
1765        if let Some(other_fetch) = &other.fetch_metrics {
1766            if let Some(self_fetch) = &self.fetch_metrics {
1767                self_fetch.merge_from(other_fetch);
1768            } else {
1769                self.fetch_metrics = Some(other_fetch.clone());
1770            }
1771        }
1772        self.metadata_mem_size += other.metadata_mem_size;
1773        self.num_range_builders += other.num_range_builders;
1774    }
1775
1776    /// Reports total rows.
1777    pub(crate) fn observe_rows(&self, read_type: &str) {
1778        READ_ROWS_TOTAL
1779            .with_label_values(&[read_type])
1780            .inc_by(self.num_rows as u64);
1781    }
1782}
1783
1784/// Builder to build a parquet record batch stream for a row group.
1785pub(crate) struct RowGroupReaderBuilder {
1786    /// SST file to read.
1787    ///
1788    /// Holds the file handle to avoid the file purge it.
1789    file_handle: FileHandle,
1790    /// Path of the file.
1791    file_path: String,
1792    /// Metadata of the parquet file.
1793    parquet_meta: Arc<ParquetMetaData>,
1794    /// Immutable metadata size, computed once when the footer is decoded.
1795    parquet_metadata_size: usize,
1796    /// Arrow reader metadata for building async stream.
1797    arrow_metadata: ArrowReaderMetadata,
1798    /// Projected output schema aligned with `projection.projected_root_presence`.
1799    output_schema: SchemaRef,
1800    /// Object store as an Operator.
1801    object_store: ObjectStore,
1802    /// Projection mask.
1803    projection: ProjectionMaskPlan,
1804    /// Whether projected read columns include nested paths.
1805    has_nested_projection: bool,
1806    /// Cache.
1807    cache_strategy: CacheStrategy,
1808    /// Pre-built prefilter state. `None` if prefiltering is not applicable.
1809    prefilter_builder: Option<PrefilterContextBuilder>,
1810    /// Hint for rows in a decoded batch.
1811    batch_size: usize,
1812}
1813
1814/// Context passed to [RowGroupReaderBuilder::build()] carrying all information
1815/// needed for prefiltering decisions.
1816pub(crate) struct RowGroupBuildContext<'a> {
1817    /// Index of the row group to read.
1818    pub(crate) row_group_idx: usize,
1819    /// Row selection for the row group. `None` means all rows.
1820    pub(crate) row_selection: Option<RowSelection>,
1821    /// Metrics for tracking fetch operations.
1822    pub(crate) fetch_metrics: Option<&'a ParquetFetchMetrics>,
1823}
1824
1825impl RowGroupReaderBuilder {
1826    /// Path of the file to read.
1827    pub(crate) fn file_path(&self) -> &str {
1828        &self.file_path
1829    }
1830
1831    /// Handle of the file to read.
1832    pub(crate) fn file_handle(&self) -> &FileHandle {
1833        &self.file_handle
1834    }
1835
1836    pub(crate) fn parquet_metadata(&self) -> &Arc<ParquetMetaData> {
1837        &self.parquet_meta
1838    }
1839
1840    pub(crate) fn parquet_metadata_size(&self) -> usize {
1841        self.parquet_metadata_size
1842    }
1843
1844    pub(crate) fn cache_strategy(&self) -> &CacheStrategy {
1845        &self.cache_strategy
1846    }
1847
1848    pub(crate) fn has_predicate_prefilter(&self) -> bool {
1849        self.prefilter_builder.is_some()
1850    }
1851
1852    /// Builds a parquet record batch stream to read the row group at `row_group_idx`.
1853    ///
1854    /// If prefiltering is applicable (based on `build_ctx`), this performs a two-phase read:
1855    /// 1. Reads only the prefilter columns (e.g. PK column), applies filters to get a refined row selection
1856    /// 2. Reads the full projection with the refined row selection
1857    ///
1858    /// The prefilter pass is *best-effort pruning*, not the precise filter for the query.
1859    /// Predicates that cannot be lowered to prefilter columns (column not projected,
1860    /// expression not supported, etc.) are silently skipped. Correctness rests on the
1861    /// DataFusion `FilterExec` above this reader, which always re-applies the original
1862    /// predicate. With predicate prefiltering enabled, tag and timestamp predicates that
1863    /// flow through [`SimpleFilterEvaluator`] are enforced precisely in this pass. See
1864    /// [`build_reader_filter_plan`] for the bucketing rules and disabled mode.
1865    ///
1866    /// When the prefilter result selects no rows, the second read still issues but
1867    /// parquet-rs short-circuits before any column-chunk IO: the row-group state machine
1868    /// jumps to `Finished` once it sees `num_rows_selected() == 0`, so no fast path is
1869    /// added here.
1870    pub(crate) async fn build(
1871        &self,
1872        build_ctx: RowGroupBuildContext<'_>,
1873    ) -> Result<ProjectedRecordBatchStream> {
1874        let prefilter_ctx = self.prefilter_builder.as_ref().map(|b| b.build());
1875
1876        let Some(mut prefilter_ctx) = prefilter_ctx else {
1877            // No prefilter applicable, build stream with full projection.
1878            let stream = self
1879                .build_with_projection(
1880                    build_ctx.row_group_idx,
1881                    build_ctx.row_selection,
1882                    self.projection.mask.clone(),
1883                    build_ctx.fetch_metrics,
1884                )
1885                .await?;
1886            return self.make_projected_stream(stream);
1887        };
1888
1889        let prefilter_start = Instant::now();
1890        let prefilter_result = execute_prefilter(&mut prefilter_ctx, self, &build_ctx).await?;
1891        if let Some(metrics) = build_ctx.fetch_metrics {
1892            let mut data = metrics.data.lock().unwrap();
1893            data.prefilter_cost += prefilter_start.elapsed();
1894            data.prefilter_filtered_rows += prefilter_result.filtered_rows;
1895        }
1896
1897        let refined_selection = Some(prefilter_result.refined_selection);
1898
1899        let stream = self
1900            .build_with_projection(
1901                build_ctx.row_group_idx,
1902                refined_selection,
1903                self.projection.mask.clone(),
1904                build_ctx.fetch_metrics,
1905            )
1906            .await?;
1907        self.make_projected_stream(stream)
1908    }
1909
1910    /// Builds the normal projection without running the generic predicate prefilter.
1911    ///
1912    /// The series reader uses this after computing its own primary-key-only row
1913    /// selection.
1914    pub(crate) async fn build_without_prefilter(
1915        &self,
1916        build_ctx: RowGroupBuildContext<'_>,
1917    ) -> Result<ProjectedRecordBatchStream> {
1918        let stream = self
1919            .build_with_projection(
1920                build_ctx.row_group_idx,
1921                build_ctx.row_selection,
1922                self.projection.mask.clone(),
1923                build_ctx.fetch_metrics,
1924            )
1925            .await?;
1926        self.make_projected_stream(stream)
1927    }
1928
1929    /// Builds a stream that reads only the encoded primary-key column.
1930    ///
1931    /// It preserves the normal reader's binary-or-dictionary decision. This path deliberately
1932    /// skips the normal prefilter pass: the caller reads `__primary_key` once and applies all
1933    /// encoded-primary-key filters to the returned batches.
1934    pub(crate) async fn build_primary_key(
1935        &self,
1936        build_ctx: RowGroupBuildContext<'_>,
1937    ) -> Result<ProjectedRecordBatchStream> {
1938        let parquet_schema = self.parquet_meta.file_metadata().schema_descr();
1939        let primary_key_index = parquet_schema
1940            .columns()
1941            .iter()
1942            .position(|column| column.name() == PRIMARY_KEY_COLUMN_NAME)
1943            .context(UnexpectedSnafu {
1944                reason: "SST does not contain __primary_key",
1945            })?;
1946        let projection = ProjectionMask::leaves(parquet_schema, [primary_key_index]);
1947
1948        self.build_with_projection(
1949            build_ctx.row_group_idx,
1950            build_ctx.row_selection,
1951            projection,
1952            build_ctx.fetch_metrics,
1953        )
1954        .await
1955    }
1956
1957    fn make_projected_stream(
1958        &self,
1959        stream: ProjectedRecordBatchStream,
1960    ) -> Result<ProjectedRecordBatchStream> {
1961        if !self.has_nested_projection {
1962            return Ok(stream);
1963        }
1964
1965        Ok(NestedSchemaAligner::new(
1966            stream,
1967            self.projection.projected_root_presence.clone(),
1968            self.output_schema.clone(),
1969        )?
1970        .boxed())
1971    }
1972
1973    /// Builds a parquet record batch stream with a custom projection mask.
1974    pub(crate) async fn build_with_projection(
1975        &self,
1976        row_group_idx: usize,
1977        row_selection: Option<RowSelection>,
1978        projection: ProjectionMask,
1979        fetch_metrics: Option<&ParquetFetchMetrics>,
1980    ) -> Result<ProjectedRecordBatchStream> {
1981        let range_fetcher = SstParquetRangeFetcher::new(
1982            self.file_handle.file_id(),
1983            self.file_path.clone(),
1984            self.object_store.clone(),
1985            self.cache_strategy.clone(),
1986            row_group_idx,
1987            fetch_metrics.cloned(),
1988        );
1989
1990        build_sst_parquet_record_batch_stream(
1991            self.arrow_metadata.clone(),
1992            row_group_idx,
1993            row_selection,
1994            projection,
1995            range_fetcher,
1996            self.file_path.clone(),
1997            self.batch_size,
1998        )
1999    }
2000}
2001
2002#[derive(Clone)]
2003/// The filter to evaluate or the prune result of the default value.
2004pub(crate) enum MaybeFilter {
2005    /// The filter to evaluate.
2006    Filter(SimpleFilterEvaluator),
2007    /// The filter matches the default value.
2008    Matched,
2009    /// The filter is pruned.
2010    Pruned,
2011}
2012
2013impl MaybeFilter {
2014    /// Returns the inner filter when it is available.
2015    pub(crate) fn as_filter(&self) -> Option<&SimpleFilterEvaluator> {
2016        match self {
2017            MaybeFilter::Filter(filter) => Some(filter),
2018            MaybeFilter::Matched | MaybeFilter::Pruned => None,
2019        }
2020    }
2021}
2022
2023#[derive(Clone)]
2024/// Context to evaluate the column filter for a parquet file.
2025pub(crate) struct SimpleFilterContext {
2026    /// Filter to evaluate.
2027    filter: MaybeFilter,
2028    /// Debug string of the original logical expression.
2029    expr_str: String,
2030    /// Id of the column to evaluate.
2031    column_id: ColumnId,
2032    /// Semantic type of the column.
2033    semantic_type: SemanticType,
2034}
2035
2036impl SimpleFilterContext {
2037    /// Creates a context for the `expr`.
2038    ///
2039    /// Returns None if the column to filter doesn't exist in the SST metadata or the
2040    /// expected metadata.
2041    pub(crate) fn new_opt(
2042        sst_meta: &RegionMetadataRef,
2043        expected_meta: Option<&RegionMetadata>,
2044        expr: &Expr,
2045    ) -> Option<Self> {
2046        let filter = SimpleFilterEvaluator::try_new(expr)?;
2047        let expr_str = format!("{expr:?}");
2048        let (column_metadata, maybe_filter) = match expected_meta {
2049            Some(meta) => {
2050                // Gets the column metadata from the expected metadata.
2051                let column = meta.column_by_name(filter.column_name())?;
2052                // Checks if the column is present in the SST metadata. We still uses the
2053                // column from the expected metadata.
2054                match sst_meta.column_by_id(column.column_id) {
2055                    Some(sst_column) => {
2056                        debug_assert_eq!(column.semantic_type, sst_column.semantic_type);
2057                        // Schema evolution can make field columns with the same id have
2058                        // different concrete data types across SSTs. In that case,
2059                        // evaluating this simple filter against current SST column may
2060                        // raise an invalid cross-type comparison error (e.g. Float64 == Utf8).
2061                        let maybe_filter = if sst_column.column_schema.data_type
2062                            == column.column_schema.data_type
2063                        {
2064                            MaybeFilter::Filter(filter)
2065                        } else {
2066                            // Altering tag or timestamp column types is not allowed,
2067                            // so only field columns can reach this branch.
2068                            debug_assert_eq!(column.semantic_type, SemanticType::Field);
2069                            return None;
2070                        };
2071                        (column, maybe_filter)
2072                    }
2073                    None => {
2074                        // If the column is not present in the SST metadata, we evaluate the filter
2075                        // against the default value of the column.
2076                        // If we can't evaluate the filter, we return None.
2077                        if pruned_by_default(&filter, column)? {
2078                            (column, MaybeFilter::Pruned)
2079                        } else {
2080                            (column, MaybeFilter::Matched)
2081                        }
2082                    }
2083                }
2084            }
2085            None => {
2086                let column = sst_meta.column_by_name(filter.column_name())?;
2087                (column, MaybeFilter::Filter(filter))
2088            }
2089        };
2090
2091        Some(Self {
2092            filter: maybe_filter,
2093            expr_str,
2094            column_id: column_metadata.column_id,
2095            semantic_type: column_metadata.semantic_type,
2096        })
2097    }
2098
2099    /// Returns the filter to evaluate.
2100    pub(crate) fn filter(&self) -> &MaybeFilter {
2101        &self.filter
2102    }
2103
2104    /// Returns the original logical expression string.
2105    pub(crate) fn expr_str(&self) -> &str {
2106        &self.expr_str
2107    }
2108
2109    /// Returns the column id.
2110    pub(crate) fn column_id(&self) -> ColumnId {
2111        self.column_id
2112    }
2113
2114    /// Returns the semantic type of the column.
2115    pub(crate) fn semantic_type(&self) -> SemanticType {
2116        self.semantic_type
2117    }
2118}
2119
2120/// Context to evaluate a physical expression for a parquet file.
2121#[derive(Clone)]
2122pub(crate) struct PhysicalFilterContext {
2123    /// Filter to evaluate.
2124    filter: Arc<dyn PhysicalExpr>,
2125    /// Debug string of the original logical expression.
2126    expr_str: String,
2127    /// Id of the column to evaluate.
2128    column_id: ColumnId,
2129    /// Name of the column to evaluate.
2130    column_name: String,
2131    /// Semantic type of the column.
2132    semantic_type: SemanticType,
2133    /// Schema containing only the referenced column.
2134    schema: SchemaRef,
2135    /// Whether the original logical expression is immutable across queries.
2136    immutable: bool,
2137}
2138
2139impl PhysicalFilterContext {
2140    /// Creates a context for the `expr`.
2141    ///
2142    /// Returns None if the expression doesn't reference exactly one column or the
2143    /// column to filter doesn't exist in the SST metadata or the expected metadata.
2144    pub(crate) fn new_opt(
2145        sst_meta: &RegionMetadataRef,
2146        expected_meta: Option<&RegionMetadata>,
2147        read_format: &FlatReadFormat,
2148        expr: &Expr,
2149    ) -> Option<Self> {
2150        if !Self::is_prefilter_candidate(expr) {
2151            return None;
2152        }
2153        let expr_str = format!("{expr:?}");
2154        let column_name = Self::single_column_name(expr)?;
2155        let column_metadata = match expected_meta {
2156            Some(meta) => {
2157                let column = meta.column_by_name(&column_name)?;
2158                let sst_column = sst_meta.column_by_id(column.column_id)?;
2159                // Physical expr requires the column name to match the SST column name.
2160                if sst_column.column_schema.name != column_name {
2161                    return None;
2162                }
2163                column
2164            }
2165            None => sst_meta.column_by_name(&column_name)?,
2166        };
2167
2168        // The column must be present in the projected arrow schema for the
2169        // prefilter to be able to read it.
2170        let (_, field) = read_format.arrow_schema().column_with_name(&column_name)?;
2171        let field = field.clone();
2172        let schema = Arc::new(ArrowSchema::new(vec![field]));
2173        let physical_expr = Predicate::to_physical_expr(expr, &schema)
2174            .inspect_err(|e| {
2175                error!(e; "Unable to build physical filter for {expr}, schema: {schema:?}");
2176            })
2177            .ok()?;
2178        let immutable = expr_is_immutable(expr);
2179
2180        Some(Self {
2181            filter: physical_expr,
2182            expr_str,
2183            column_id: column_metadata.column_id,
2184            column_name,
2185            semantic_type: column_metadata.semantic_type,
2186            schema,
2187            immutable,
2188        })
2189    }
2190
2191    /// Returns true if the expression is a variant we want to evaluate as a
2192    /// physical prefilter. Binary exprs are intentionally excluded because
2193    /// [`SimpleFilterEvaluator`] already handles them.
2194    // TODO(yingwen): extend more expressions if necessary. For example, allow some cheap scalar functions (e.g. `lower`, `length`, date truncations)
2195    fn is_prefilter_candidate(expr: &Expr) -> bool {
2196        if !matches!(
2197            expr,
2198            Expr::InList(_) | Expr::IsNull(_) | Expr::IsNotNull(_) | Expr::Between(_)
2199        ) {
2200            return false;
2201        }
2202
2203        // If any functions are found in the expr, it will be not considered as worthy enough to
2204        // be evaluated in the prefilter. At last, prefilter reads the Parquet files one more time.
2205        !expr
2206            .exists(|e| Ok(matches!(e, Expr::ScalarFunction(_))))
2207            .unwrap_or(false)
2208    }
2209
2210    fn single_column_name(expr: &Expr) -> Option<String> {
2211        let mut columns = HashSet::new();
2212        if expr_to_columns(expr, &mut columns).is_err() {
2213            return None;
2214        }
2215        if columns.len() != 1 {
2216            return None;
2217        }
2218        columns.iter().next().map(|column| column.name.clone())
2219    }
2220
2221    /// Returns the filter to evaluate.
2222    pub(crate) fn filter(&self) -> &Arc<dyn PhysicalExpr> {
2223        &self.filter
2224    }
2225
2226    /// Returns the original logical expression string.
2227    pub(crate) fn expr_str(&self) -> &str {
2228        &self.expr_str
2229    }
2230
2231    /// Returns the column id.
2232    pub(crate) fn column_id(&self) -> ColumnId {
2233        self.column_id
2234    }
2235
2236    /// Returns the column name.
2237    pub(crate) fn column_name(&self) -> &str {
2238        &self.column_name
2239    }
2240
2241    /// Returns the semantic type of the column.
2242    pub(crate) fn semantic_type(&self) -> SemanticType {
2243        self.semantic_type
2244    }
2245
2246    /// Returns the schema containing only the referenced column.
2247    pub(crate) fn schema(&self) -> &SchemaRef {
2248        &self.schema
2249    }
2250
2251    /// Returns true if the original logical expression is immutable across queries.
2252    pub(crate) fn is_immutable(&self) -> bool {
2253        self.immutable
2254    }
2255}
2256
2257fn expr_is_immutable(expr: &Expr) -> bool {
2258    let mut is_immutable = true;
2259    let _ = expr.apply(|expr| match expr {
2260        Expr::ScalarFunction(function)
2261            if function.func.signature().volatility != Volatility::Immutable =>
2262        {
2263            is_immutable = false;
2264            Ok(TreeNodeRecursion::Stop)
2265        }
2266        Expr::ScalarVariable(_, _) => {
2267            is_immutable = false;
2268            Ok(TreeNodeRecursion::Stop)
2269        }
2270        _ => Ok(TreeNodeRecursion::Continue),
2271    });
2272    is_immutable
2273}
2274
2275/// Prune a column by its default value.
2276/// Returns false if we can't create the default value or evaluate the filter.
2277fn pruned_by_default(filter: &SimpleFilterEvaluator, column: &ColumnMetadata) -> Option<bool> {
2278    let value = column.column_schema.create_default().ok().flatten()?;
2279    let scalar_value = value
2280        .try_to_scalar_value(&column.column_schema.data_type)
2281        .ok()?;
2282    let matches = filter.evaluate_scalar(&scalar_value).ok()?;
2283    Some(!matches)
2284}
2285
2286/// Parquet batch reader to read our SST format.
2287pub struct ParquetReader {
2288    /// File range context.
2289    context: FileRangeContextRef,
2290    /// Row group selection to read.
2291    selection: RowGroupSelection,
2292    /// Reader of current row group.
2293    reader: Option<FlatPruneReader>,
2294    /// Metrics for tracking row group fetch operations.
2295    fetch_metrics: ParquetFetchMetrics,
2296}
2297
2298impl ParquetReader {
2299    #[tracing::instrument(
2300        skip_all,
2301        fields(
2302            region_id = %self.context.reader_builder().file_handle.region_id(),
2303            file_id = %self.context.reader_builder().file_handle.file_id()
2304        )
2305    )]
2306    pub async fn next_record_batch(&mut self) -> Result<Option<RecordBatch>> {
2307        loop {
2308            if let Some(reader) = &mut self.reader {
2309                if let Some(batch) = reader.next_batch().await? {
2310                    return Ok(Some(batch));
2311                }
2312                self.reader = None;
2313                continue;
2314            }
2315
2316            let Some((row_group_idx, row_selection)) = self.selection.pop_first() else {
2317                return Ok(None);
2318            };
2319
2320            let skip_fields = self.context.pre_filter_mode().skip_fields();
2321            let parquet_reader = self
2322                .context
2323                .reader_builder()
2324                .build(self.context.build_context(
2325                    row_group_idx,
2326                    Some(row_selection),
2327                    Some(&self.fetch_metrics),
2328                ))
2329                .await?;
2330            self.reader = Some(FlatPruneReader::new_with_row_group_reader(
2331                self.context.clone(),
2332                FlatRowGroupReader::new(self.context.clone(), parquet_reader),
2333                skip_fields,
2334            ));
2335        }
2336    }
2337    /// Creates a new reader.
2338    #[tracing::instrument(
2339        skip_all,
2340        fields(
2341            region_id = %context.reader_builder().file_handle.region_id(),
2342            file_id = %context.reader_builder().file_handle.file_id()
2343        )
2344    )]
2345    pub(crate) async fn new(
2346        context: FileRangeContextRef,
2347        mut selection: RowGroupSelection,
2348    ) -> Result<Self> {
2349        let fetch_metrics = ParquetFetchMetrics::default();
2350        let reader = if let Some((row_group_idx, row_selection)) = selection.pop_first() {
2351            let skip_fields = context.pre_filter_mode().skip_fields();
2352            let parquet_reader = context
2353                .reader_builder()
2354                .build(context.build_context(
2355                    row_group_idx,
2356                    Some(row_selection),
2357                    Some(&fetch_metrics),
2358                ))
2359                .await?;
2360            Some(FlatPruneReader::new_with_row_group_reader(
2361                context.clone(),
2362                FlatRowGroupReader::new(context.clone(), parquet_reader),
2363                skip_fields,
2364            ))
2365        } else {
2366            None
2367        };
2368
2369        Ok(ParquetReader {
2370            context,
2371            selection,
2372            reader,
2373            fetch_metrics,
2374        })
2375    }
2376
2377    /// Returns the metadata of the SST.
2378    pub fn metadata(&self) -> &RegionMetadataRef {
2379        self.context.read_format().metadata()
2380    }
2381
2382    pub fn parquet_metadata(&self) -> Arc<ParquetMetaData> {
2383        self.context.reader_builder().parquet_meta.clone()
2384    }
2385}
2386
2387/// Reader to read a row group of a parquet file in flat format, returning RecordBatch.
2388pub(crate) struct FlatRowGroupReader {
2389    /// Context for file ranges.
2390    context: FileRangeContextRef,
2391    /// Inner parquet record batch stream.
2392    stream: ProjectedRecordBatchStream,
2393    /// Cached sequence array to override sequences.
2394    override_sequence: Option<ArrayRef>,
2395}
2396
2397impl FlatRowGroupReader {
2398    /// Creates a new flat reader from file range.
2399    pub(crate) fn new(context: FileRangeContextRef, stream: ProjectedRecordBatchStream) -> Self {
2400        // The batch length from the reader should be less than or equal to DEFAULT_READ_BATCH_SIZE.
2401        let override_sequence = context
2402            .read_format()
2403            .new_override_sequence_array(DEFAULT_READ_BATCH_SIZE);
2404
2405        Self {
2406            context,
2407            stream,
2408            override_sequence,
2409        }
2410    }
2411
2412    /// Returns the next RecordBatch.
2413    pub(crate) async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
2414        match self.stream.next().await {
2415            Some(batch_result) => {
2416                let record_batch = batch_result?;
2417
2418                let record_batch = self
2419                    .context
2420                    .read_format()
2421                    .convert_batch(record_batch, self.override_sequence.as_ref())?;
2422                Ok(Some(record_batch))
2423            }
2424            None => Ok(None),
2425        }
2426    }
2427}
2428
2429#[cfg(test)]
2430mod tests {
2431    use std::any::Any;
2432    use std::fmt::{Debug, Formatter};
2433    use std::sync::{Arc, LazyLock};
2434
2435    use common_error::ext::WhateverResult;
2436    use common_function::scalars::json::json_get::JsonGetWithType;
2437    use common_function::scalars::udf::create_udf;
2438    use common_recordbatch::ext::RecordBatchExt;
2439    use datafusion::arrow::datatypes::DataType;
2440    use datafusion_common::ScalarValue;
2441    use datafusion_expr::expr::ScalarFunction;
2442    use datafusion_expr::{
2443        ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
2444        col, lit,
2445    };
2446    use datatypes::arrow::array::{ArrayRef, Int64Array, StringArray, StructArray};
2447    use datatypes::arrow::datatypes::{Fields, Schema};
2448    use datatypes::arrow::record_batch::RecordBatch;
2449    use datatypes::extension::json::Json2ExtensionType;
2450    use datatypes::prelude::ConcreteDataType;
2451    use datatypes::schema::ColumnSchema;
2452    use object_store::services::Memory;
2453    use parquet::arrow::ArrowWriter;
2454    use parquet::file::properties::WriterProperties;
2455    use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
2456    use store_api::region_request::PathType;
2457    use store_api::storage::RegionId;
2458    use table::predicate::Predicate;
2459
2460    use super::*;
2461    use crate::sst::parquet::metadata::MetadataLoader;
2462    use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns};
2463    use crate::test_util::sst_util::{sst_file_handle, sst_region_metadata};
2464
2465    #[test]
2466    fn test_skip_prefilter_for_json_get() -> WhateverResult<()> {
2467        fn json_get_expr(base: Expr, path: &str) -> Expr {
2468            let json_get = Arc::new(create_udf(Arc::new(JsonGetWithType::default())));
2469            Expr::ScalarFunction(ScalarFunction::new_udf(json_get, vec![base, lit(path)]))
2470        }
2471
2472        let metadata = Arc::new(sst_region_metadata());
2473        let format = FlatReadFormat::new(
2474            metadata.clone(),
2475            ReadColumns::from_deduped_column_ids(
2476                metadata.column_metadatas.iter().map(|c| c.column_id),
2477            ),
2478            None,
2479            "test",
2480            true,
2481        )?;
2482        let new_filter =
2483            |expr: Expr| PhysicalFilterContext::new_opt(&metadata, None, &format, &expr);
2484
2485        let json_get = || json_get_expr(col("field_0"), "a.b");
2486
2487        let regular_expr = col("field_0").is_null();
2488        assert!(new_filter(regular_expr).is_some());
2489
2490        let is_null = json_get().is_null();
2491        assert!(new_filter(is_null).is_none());
2492
2493        let is_not_null = json_get().is_not_null();
2494        assert!(new_filter(is_not_null).is_none());
2495
2496        let in_list = json_get().in_list(vec![lit("value")], false);
2497        assert!(new_filter(in_list).is_none());
2498
2499        let in_list_nested = col("field_0").in_list(vec![json_get()], false);
2500        assert!(new_filter(in_list_nested).is_none());
2501
2502        let between = json_get().between(lit(1_u64), lit(10_u64));
2503        assert!(new_filter(between).is_none());
2504
2505        let between_nested = col("field_0").between(json_get(), lit(10_u64));
2506        assert!(new_filter(between_nested).is_none());
2507
2508        Ok(())
2509    }
2510
2511    #[tokio::test]
2512    async fn test_nested_projection_reads_partial_json2_physical_fields() -> WhateverResult<()> {
2513        // Write a full JSON2-like Arrow struct:
2514        // j: { a: { x: int, y: string }, b: string }.
2515        // The test later requests only j.a.x and verifies that the physical Parquet projection
2516        // does not materialize j.a.y or j.b.
2517
2518        let xy_fields = Fields::from(vec![
2519            Arc::new(Field::new("x", DataType::Int64, true)),
2520            Arc::new(Field::new("y", DataType::Utf8, true)),
2521        ]);
2522        let a_field = Arc::new(Field::new("a", DataType::Struct(xy_fields.clone()), true));
2523        let b_field = Arc::new(Field::new("b", DataType::Utf8, true));
2524        let json_fields = Fields::from(vec![a_field, b_field]);
2525        let json_field = Field::new("j", DataType::Struct(json_fields.clone()), true)
2526            .with_extension_type(Json2ExtensionType::default());
2527        let schema = Arc::new(Schema::new(vec![json_field]));
2528
2529        let a_array = Arc::new(StructArray::new(
2530            xy_fields,
2531            vec![
2532                Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef,
2533                Arc::new(StringArray::from_iter_values(["x1", "x2", "x3"])) as ArrayRef,
2534            ],
2535            None,
2536        )) as ArrayRef;
2537        let b_array = Arc::new(StringArray::from_iter_values(["b1", "b2", "b3"])) as ArrayRef;
2538        let j_array =
2539            Arc::new(StructArray::new(json_fields, vec![a_array, b_array], None)) as ArrayRef;
2540        let columns = vec![j_array];
2541
2542        let batch = RecordBatch::try_new(schema, columns).map_err(|e| e.to_string())?;
2543
2544        // Persist the complete nested schema to an in-memory Parquet file so the projection is
2545        // exercised through parquet-rs rather than a mock.
2546
2547        let object_store = ObjectStore::new(Memory::default())
2548            .map_err(|e| e.to_string())?
2549            .finish();
2550        let file_handle = sst_file_handle(0, 1);
2551        let file_path = file_handle.file_path("test_table", PathType::Bare);
2552
2553        let mut parquet_bytes = Vec::new();
2554        ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), None)
2555            .and_then(|mut w| {
2556                w.write(&batch)?;
2557                Ok(w)
2558            })
2559            .and_then(|w| w.close())
2560            .map_err(|e| e.to_string())?;
2561        let file_size = parquet_bytes.len() as u64;
2562        object_store
2563            .write(&file_path, parquet_bytes)
2564            .await
2565            .map_err(|e| e.to_string())?;
2566
2567        let mut cache_metrics = MetadataCacheMetrics::default();
2568        let loader = MetadataLoader::new(object_store.clone(), &file_path, file_size);
2569        let parquet_meta = loader.load(&mut cache_metrics).await?;
2570        let parquet_schema = parquet_meta.file_metadata().schema_descr();
2571        assert_eq!(3, parquet_schema.num_columns());
2572
2573        // Ask Parquet to read only the deepest requested JSON2 path. This should select the single
2574        // leaf j.a.x and avoid both sibling leaves j.a.y and j.b.
2575
2576        let projection =
2577            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
2578                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
2579            )]);
2580        let projection_plan = build_projection_plan(&projection, parquet_schema);
2581        assert_eq!(vec![true], projection_plan.projected_root_presence);
2582        assert_eq!(
2583            projection_plan.mask,
2584            ProjectionMask::leaves(parquet_schema, vec![0])
2585        );
2586
2587        // Read through the low-level stream directly.
2588
2589        let arrow_metadata =
2590            ArrowReaderMetadata::try_new(Arc::new(parquet_meta), ArrowReaderOptions::new())
2591                .map_err(|e| e.to_string())?;
2592        let fetcher = SstParquetRangeFetcher::new(
2593            file_handle.file_id(),
2594            file_path.clone(),
2595            object_store,
2596            CacheStrategy::Disabled,
2597            0,
2598            None,
2599        );
2600        let mut stream = build_sst_parquet_record_batch_stream(
2601            arrow_metadata,
2602            0,
2603            None,
2604            projection_plan.mask,
2605            fetcher,
2606            file_path,
2607            1024,
2608        )?;
2609
2610        let Some(batch) = stream.next().await.transpose()? else {
2611            unreachable!()
2612        };
2613        let expected = r#"
2614+-------------+
2615| j           |
2616+-------------+
2617| {a: {x: 1}} |
2618| {a: {x: 2}} |
2619| {a: {x: 3}} |
2620+-------------+
2621"#;
2622        assert_eq!(batch.pretty_print(), expected.trim());
2623        Ok(())
2624    }
2625
2626    #[tokio::test(flavor = "current_thread")]
2627    async fn test_minmax_predicate_key_not_built_when_index_result_cache_disabled() {
2628        #[derive(Eq, PartialEq, Hash)]
2629        struct PanicDebugUdf;
2630
2631        impl Debug for PanicDebugUdf {
2632            fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
2633                panic!("minmax predicate key should not format exprs when cache is disabled");
2634            }
2635        }
2636
2637        impl ScalarUDFImpl for PanicDebugUdf {
2638            fn as_any(&self) -> &dyn Any {
2639                self
2640            }
2641
2642            fn name(&self) -> &str {
2643                "panic_debug_udf"
2644            }
2645
2646            fn signature(&self) -> &Signature {
2647                static SIGNATURE: LazyLock<Signature> =
2648                    LazyLock::new(|| Signature::variadic_any(Volatility::Immutable));
2649                &SIGNATURE
2650            }
2651
2652            fn return_type(&self, _arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
2653                Ok(DataType::Int64)
2654            }
2655
2656            fn invoke_with_args(
2657                &self,
2658                _args: ScalarFunctionArgs,
2659            ) -> datafusion_common::Result<ColumnarValue> {
2660                Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(1))))
2661            }
2662        }
2663
2664        let object_store = ObjectStore::new(Memory::default()).unwrap().finish();
2665        let file_handle = sst_file_handle(0, 1);
2666        let table_dir = "test_table".to_string();
2667        let path_type = PathType::Bare;
2668        let file_path = file_handle.file_path(&table_dir, path_type);
2669
2670        let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
2671        let batch = RecordBatch::try_from_iter([("col", col)]).unwrap();
2672        let mut parquet_bytes = Vec::new();
2673        let mut writer = ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), None).unwrap();
2674        writer.write(&batch).unwrap();
2675        writer.close().unwrap();
2676        let file_size = parquet_bytes.len() as u64;
2677        object_store.write(&file_path, parquet_bytes).await.unwrap();
2678
2679        let region_metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
2680        let read_format = FlatReadFormat::new(
2681            region_metadata.clone(),
2682            ReadColumns::from_deduped_column_ids(
2683                region_metadata
2684                    .column_metadatas
2685                    .iter()
2686                    .map(|column| column.column_id),
2687            ),
2688            None,
2689            &file_path,
2690            false,
2691        )
2692        .unwrap();
2693
2694        let mut cache_metrics = MetadataCacheMetrics::default();
2695        let loader = MetadataLoader::new(object_store.clone(), &file_path, file_size);
2696        let parquet_meta = loader.load(&mut cache_metrics).await.unwrap();
2697
2698        let udf = Arc::new(ScalarUDF::new_from_impl(PanicDebugUdf));
2699        let predicate = Predicate::new(vec![Expr::ScalarFunction(ScalarFunction::new_udf(
2700            udf,
2701            vec![],
2702        ))]);
2703        let builder = ParquetReaderBuilder::new(table_dir, path_type, file_handle, object_store)
2704            .predicate(Some(predicate))
2705            .cache(CacheStrategy::Disabled);
2706
2707        let row_group_size = parquet_meta.row_group(0).num_rows() as usize;
2708        let total_row_count = parquet_meta.file_metadata().num_rows() as usize;
2709        let mut metrics = ReaderFilterMetrics::default();
2710        let selection = builder.row_groups_by_minmax(
2711            &read_format,
2712            &parquet_meta,
2713            row_group_size,
2714            total_row_count,
2715            &mut metrics,
2716            false,
2717        );
2718
2719        assert!(!selection.is_empty());
2720    }
2721
2722    #[test]
2723    fn test_expr_is_immutable_checks_scalar_function_volatility() {
2724        #[derive(Debug, PartialEq, Eq, Hash)]
2725        struct TestVolatilityUdf {
2726            name: String,
2727            signature: Signature,
2728        }
2729
2730        impl TestVolatilityUdf {
2731            fn new(name: &str, volatility: Volatility) -> Self {
2732                Self {
2733                    name: name.to_string(),
2734                    signature: Signature::variadic_any(volatility),
2735                }
2736            }
2737        }
2738
2739        impl ScalarUDFImpl for TestVolatilityUdf {
2740            fn as_any(&self) -> &dyn Any {
2741                self
2742            }
2743
2744            fn name(&self) -> &str {
2745                &self.name
2746            }
2747
2748            fn signature(&self) -> &Signature {
2749                &self.signature
2750            }
2751
2752            fn return_type(&self, _arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
2753                Ok(DataType::Int64)
2754            }
2755
2756            fn invoke_with_args(
2757                &self,
2758                _args: ScalarFunctionArgs,
2759            ) -> datafusion_common::Result<ColumnarValue> {
2760                Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(1))))
2761            }
2762        }
2763
2764        let expr = |name: &str, volatility| {
2765            Expr::ScalarFunction(ScalarFunction::new_udf(
2766                Arc::new(ScalarUDF::new_from_impl(TestVolatilityUdf::new(
2767                    name, volatility,
2768                ))),
2769                vec![],
2770            ))
2771        };
2772
2773        assert!(expr_is_immutable(&expr(
2774            "immutable_udf",
2775            Volatility::Immutable
2776        )));
2777        assert!(!expr_is_immutable(&expr("stable_udf", Volatility::Stable)));
2778        assert!(!expr_is_immutable(&expr(
2779            "volatile_udf",
2780            Volatility::Volatile
2781        )));
2782
2783        let scalar_variable = Expr::ScalarVariable(
2784            Arc::new(Field::new("@@version", DataType::Utf8, false)),
2785            vec!["@@version".to_string()],
2786        );
2787        assert!(!expr_is_immutable(&scalar_variable));
2788    }
2789
2790    #[tokio::test(flavor = "current_thread")]
2791    async fn test_has_row_level_selection() {
2792        let object_store = ObjectStore::new(Memory::default()).unwrap().finish();
2793        let file_path = "row_level_selection.parquet";
2794
2795        let col = Arc::new(Int64Array::from_iter_values([1, 2, 3, 4, 5])) as ArrayRef;
2796        let batch = RecordBatch::try_from_iter([("col", col)]).unwrap();
2797        let props = WriterProperties::builder()
2798            .set_max_row_group_row_count(Some(3))
2799            .build();
2800        let mut parquet_bytes = Vec::new();
2801        let mut writer =
2802            ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap();
2803        writer.write(&batch).unwrap();
2804        writer.close().unwrap();
2805        let file_size = parquet_bytes.len() as u64;
2806        object_store.write(file_path, parquet_bytes).await.unwrap();
2807
2808        let mut cache_metrics = MetadataCacheMetrics::default();
2809        let loader = MetadataLoader::new(object_store, file_path, file_size);
2810        let parquet_meta = loader.load(&mut cache_metrics).await.unwrap();
2811        assert_eq!(2, parquet_meta.num_row_groups());
2812
2813        let full_row_groups = RowGroupSelection::from_full_row_group_ids([0, 1], 3, 5);
2814        assert!(!has_row_level_selection(&full_row_groups, &parquet_meta));
2815
2816        let prefix_selection = RowGroupSelection::from_row_ranges(vec![(0, vec![0..1, 1..2])], 3);
2817        assert!(has_row_level_selection(&prefix_selection, &parquet_meta));
2818
2819        let interior_selection = RowGroupSelection::from_row_ranges(vec![(0, vec![1..2, 2..3])], 3);
2820        assert!(has_row_level_selection(&interior_selection, &parquet_meta));
2821    }
2822
2823    fn expected_metadata_with_reused_tag_name(
2824        old_metadata: &RegionMetadata,
2825    ) -> Arc<RegionMetadata> {
2826        let mut builder = RegionMetadataBuilder::new(old_metadata.region_id);
2827        builder
2828            .push_column_metadata(ColumnMetadata {
2829                column_schema: ColumnSchema::new(
2830                    "tag_0".to_string(),
2831                    ConcreteDataType::string_datatype(),
2832                    true,
2833                ),
2834                semantic_type: SemanticType::Tag,
2835                column_id: 10,
2836            })
2837            .push_column_metadata(ColumnMetadata {
2838                column_schema: ColumnSchema::new(
2839                    "tag_1".to_string(),
2840                    ConcreteDataType::string_datatype(),
2841                    true,
2842                ),
2843                semantic_type: SemanticType::Tag,
2844                column_id: 1,
2845            })
2846            .push_column_metadata(ColumnMetadata {
2847                column_schema: ColumnSchema::new(
2848                    "field_0".to_string(),
2849                    ConcreteDataType::uint64_datatype(),
2850                    true,
2851                ),
2852                semantic_type: SemanticType::Field,
2853                column_id: 2,
2854            })
2855            .push_column_metadata(ColumnMetadata {
2856                column_schema: ColumnSchema::new(
2857                    "ts".to_string(),
2858                    ConcreteDataType::timestamp_millisecond_datatype(),
2859                    false,
2860                ),
2861                semantic_type: SemanticType::Timestamp,
2862                column_id: 3,
2863            })
2864            .primary_key(vec![10, 1]);
2865
2866        Arc::new(builder.build().unwrap())
2867    }
2868
2869    #[test]
2870    fn test_simple_filter_context_uses_default_value_for_mismatched_expected_metadata() {
2871        let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
2872        let expected_metadata = expected_metadata_with_reused_tag_name(metadata.as_ref());
2873        let ctx = SimpleFilterContext::new_opt(
2874            &metadata,
2875            Some(expected_metadata.as_ref()),
2876            &col("tag_0").eq(lit("a")),
2877        )
2878        .unwrap();
2879        assert!(matches!(
2880            ctx.filter(),
2881            MaybeFilter::Matched | MaybeFilter::Pruned
2882        ));
2883    }
2884
2885    #[test]
2886    fn test_simple_filter_context_drops_mismatched_field_filter() {
2887        let (sst_metadata, latest_metadata) = mock_metadata();
2888        let ctx = SimpleFilterContext::new_opt(
2889            &sst_metadata,
2890            Some(latest_metadata.as_ref()),
2891            &col("field_0").eq(lit(1_i64)),
2892        );
2893
2894        assert!(ctx.is_none());
2895    }
2896
2897    fn mock_metadata() -> (RegionMetadataRef, RegionMetadataRef) {
2898        let region_id = RegionId::new(1, 1);
2899        let make_tag_0 = || ColumnMetadata {
2900            column_schema: ColumnSchema::new(
2901                "tag_0".to_string(),
2902                ConcreteDataType::string_datatype(),
2903                true,
2904            ),
2905            semantic_type: SemanticType::Tag,
2906            column_id: 0,
2907        };
2908        let make_ts = || ColumnMetadata {
2909            column_schema: ColumnSchema::new(
2910                "ts".to_string(),
2911                ConcreteDataType::timestamp_millisecond_datatype(),
2912                false,
2913            ),
2914            semantic_type: SemanticType::Timestamp,
2915            column_id: 2,
2916        };
2917        let make_field_0 = |data_type| ColumnMetadata {
2918            column_schema: ColumnSchema::new("field_0".to_string(), data_type, true),
2919            semantic_type: SemanticType::Field,
2920            column_id: 1,
2921        };
2922
2923        let mut sst_builder = RegionMetadataBuilder::new(region_id);
2924        sst_builder
2925            .push_column_metadata(make_tag_0())
2926            .push_column_metadata(make_field_0(ConcreteDataType::uint64_datatype()))
2927            .push_column_metadata(make_ts())
2928            .primary_key(vec![0]);
2929        let sst_metadata = Arc::new(sst_builder.build().unwrap());
2930
2931        let mut expected_builder = RegionMetadataBuilder::new(region_id);
2932        expected_builder
2933            .push_column_metadata(make_tag_0())
2934            .push_column_metadata(make_field_0(ConcreteDataType::int64_datatype()))
2935            .push_column_metadata(make_ts())
2936            .primary_key(vec![0]);
2937
2938        let expected_metadata = Arc::new(expected_builder.build().unwrap());
2939
2940        (sst_metadata, expected_metadata)
2941    }
2942
2943    #[test]
2944    fn test_physical_filter_context_skips_renamed_column() {
2945        let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
2946        let expected_metadata = expected_metadata_with_reused_tag_name(metadata.as_ref());
2947        let read_format = FlatReadFormat::new(
2948            metadata.clone(),
2949            ReadColumns::from_deduped_column_ids(
2950                metadata.column_metadatas.iter().map(|c| c.column_id),
2951            ),
2952            None,
2953            "test",
2954            true,
2955        )
2956        .unwrap();
2957
2958        let ctx = PhysicalFilterContext::new_opt(
2959            &metadata,
2960            Some(expected_metadata.as_ref()),
2961            &read_format,
2962            &col("tag_0").in_list(vec![lit("a"), lit("b")], false),
2963        );
2964
2965        assert!(ctx.is_none());
2966    }
2967
2968    #[test]
2969    fn test_physical_filter_context_only_accepts_prefilter_candidates() {
2970        let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
2971        let read_format = FlatReadFormat::new(
2972            metadata.clone(),
2973            ReadColumns::from_deduped_column_ids(
2974                metadata.column_metadatas.iter().map(|c| c.column_id),
2975            ),
2976            None,
2977            "test",
2978            true,
2979        )
2980        .unwrap();
2981
2982        // InList is on the allowlist — should build a context.
2983        let in_list = col("tag_0").in_list(vec![lit("a"), lit("b")], false);
2984        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &in_list).is_some());
2985
2986        // NOT IN uses the same variant with `negated: true` — also accepted.
2987        let not_in = col("tag_0").in_list(vec![lit("a"), lit("b")], true);
2988        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &not_in).is_some());
2989
2990        // IS NULL / IS NOT NULL are accepted.
2991        let is_null = col("tag_0").is_null();
2992        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &is_null).is_some());
2993        let is_not_null = col("tag_0").is_not_null();
2994        assert!(
2995            PhysicalFilterContext::new_opt(&metadata, None, &read_format, &is_not_null).is_some()
2996        );
2997
2998        // BETWEEN is accepted.
2999        let between = col("field_0").between(lit(1_u64), lit(10_u64));
3000        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &between).is_some());
3001
3002        // Binary expr is handled by SimpleFilterEvaluator — rejected here.
3003        let binary = col("tag_0").eq(lit("a"));
3004        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &binary).is_none());
3005    }
3006
3007    fn write_test_parquet_with_pk_column(values: &[&[u8]]) -> bytes::Bytes {
3008        use datatypes::arrow::array::{
3009            BinaryArray, TimestampMillisecondArray, UInt8Array, UInt64Array,
3010        };
3011        use datatypes::arrow::datatypes::{Field as ArrowField, Schema as ArrowSchema, TimeUnit};
3012        use store_api::storage::consts::{
3013            OP_TYPE_COLUMN_NAME, PRIMARY_KEY_COLUMN_NAME, SEQUENCE_COLUMN_NAME,
3014        };
3015
3016        let n = values.len();
3017        let schema = Arc::new(ArrowSchema::new(vec![
3018            ArrowField::new(
3019                "ts",
3020                DataType::Timestamp(TimeUnit::Millisecond, None),
3021                false,
3022            ),
3023            ArrowField::new(PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false),
3024            ArrowField::new(SEQUENCE_COLUMN_NAME, DataType::UInt64, false),
3025            ArrowField::new(OP_TYPE_COLUMN_NAME, DataType::UInt8, false),
3026        ]));
3027        let ts: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![0_i64; n]));
3028        let pk: ArrayRef = Arc::new(BinaryArray::from_iter_values(values.iter().copied()));
3029        let seq: ArrayRef = Arc::new(UInt64Array::from(vec![0_u64; n]));
3030        let op: ArrayRef = Arc::new(UInt8Array::from(vec![0_u8; n]));
3031        let batch = RecordBatch::try_new(schema.clone(), vec![ts, pk, seq, op]).unwrap();
3032
3033        let mut bytes = Vec::new();
3034        let mut writer = ArrowWriter::try_new(&mut bytes, schema, None).unwrap();
3035        writer.write(&batch).unwrap();
3036        writer.close().unwrap();
3037        bytes::Bytes::from(bytes)
3038    }
3039
3040    fn load_parquet_meta(bytes: bytes::Bytes) -> Arc<ParquetMetaData> {
3041        let builder =
3042            parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
3043        builder.metadata().clone()
3044    }
3045
3046    #[test]
3047    fn test_should_read_pk_as_binary_small_chunk_returns_false() {
3048        let bytes = write_test_parquet_with_pk_column(&[b"a", b"b", b"c"]);
3049        let meta = load_parquet_meta(bytes);
3050
3051        assert!(!should_read_pk_as_binary_with_limit(&meta, 1024));
3052    }
3053
3054    #[test]
3055    fn test_should_read_pk_as_binary_large_chunk_returns_true() {
3056        let owned: Vec<Vec<u8>> = (0..512u32)
3057            .map(|i| {
3058                let mut v = vec![0u8; 16];
3059                v[..4].copy_from_slice(&i.to_le_bytes());
3060                v
3061            })
3062            .collect();
3063        let refs: Vec<&[u8]> = owned.iter().map(|v| v.as_slice()).collect();
3064        let bytes = write_test_parquet_with_pk_column(&refs);
3065        let meta = load_parquet_meta(bytes);
3066
3067        assert!(should_read_pk_as_binary_with_limit(&meta, 1024));
3068    }
3069}