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