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::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta};
57use crate::error::{
58    ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu,
59    UnexpectedSnafu,
60};
61use crate::metrics::{
62    PRECISE_FILTER_ROWS_TOTAL, READ_ROW_GROUPS_TOTAL, READ_ROWS_IN_ROW_GROUP_TOTAL,
63    READ_ROWS_TOTAL, READ_STAGE_ELAPSED,
64};
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            .map(|(column_id, layout)| {
659                let column = region_meta
660                    .column_by_id(*column_id)
661                    .context(UnexpectedSnafu {
662                        reason: format!("JSON2 target column by id {column_id} does not exist"),
663                    })?;
664                Ok((column.column_schema.name.clone(), layout.clone()))
665            })
666            .collect::<Result<HashMap<_, _>>>()?;
667
668        let reader_builder = RowGroupReaderBuilder {
669            memory_tracker: self.memory_tracker.clone(),
670            file_handle: self.file_handle.clone(),
671            file_path,
672            parquet_meta,
673            parquet_metadata_size,
674            arrow_metadata,
675            output_schema,
676            json2_rewrite_targets,
677            object_store: self.object_store.clone(),
678            projection: projection_plan,
679            has_nested_projection,
680            cache_strategy: self.cache_strategy.clone(),
681            prefilter_builder: filter_plan.prefilter_builder,
682            batch_size: self.batch_size,
683        };
684
685        let partition_filter = self.build_partition_filter(&read_format, &prune_schema)?;
686
687        let range_index_store = self.series_index.as_ref().and_then(|context| {
688            let region_id = self
689                .expected_metadata
690                .as_ref()
691                .unwrap_or(&region_meta)
692                .region_id;
693            (self.file_handle.region_id() == region_id
694                && context
695                    .version
696                    .range_indexes
697                    .contains_key(&self.file_handle.file_id().file_id()))
698            .then(|| context.store.clone())
699        });
700        let context = FileRangeContext::new(
701            reader_builder,
702            RangeBase {
703                filters: filter_plan.remaining_simple_filters,
704                dyn_filters,
705                read_format,
706                expected_metadata: self.expected_metadata.clone(),
707                prune_schema,
708                codec,
709                compat_batch: None,
710                compaction_projection_mapper,
711                pre_filter_mode: self.pre_filter_mode,
712                partition_filter,
713            },
714            range_index_store,
715        );
716
717        metrics.build_cost += start.elapsed();
718
719        Ok(Some((context, selection)))
720    }
721
722    fn is_same_region_partition(
723        region_partition_expr_str: Option<&str>,
724        file_partition_expr: Option<&PartitionExpr>,
725    ) -> Result<(Option<PartitionExpr>, bool)> {
726        let region_partition_expr = match region_partition_expr_str {
727            Some(expr_str) => crate::region::parse_partition_expr(Some(expr_str))?,
728            None => None,
729        };
730
731        let is_same = region_partition_expr.as_ref() == file_partition_expr;
732        Ok((region_partition_expr, is_same))
733    }
734
735    /// Compare partition expressions from expected metadata and file metadata,
736    /// and build a partition filter if they differ.
737    fn build_partition_filter(
738        &self,
739        read_format: &FlatReadFormat,
740        prune_schema: &Arc<datatypes::schema::Schema>,
741    ) -> Result<Option<PartitionFilterContext>> {
742        let region_partition_expr_str = self
743            .expected_metadata
744            .as_ref()
745            .and_then(|meta| meta.partition_expr.as_ref());
746        let file_partition_expr_ref = self.file_handle.meta_ref().partition_expr.as_ref();
747
748        let (region_partition_expr, is_same_region_partition) = Self::is_same_region_partition(
749            region_partition_expr_str.map(|s| s.as_str()),
750            file_partition_expr_ref,
751        )?;
752
753        if is_same_region_partition {
754            return Ok(None);
755        }
756
757        let Some(region_partition_expr) = region_partition_expr else {
758            return Ok(None);
759        };
760
761        // Collect columns referenced by the partition expression.
762        let mut referenced_columns = HashSet::new();
763        region_partition_expr.collect_column_names(&mut referenced_columns);
764
765        // Build a partition_schema containing only referenced columns.
766        let partition_schema = Arc::new(datatypes::schema::Schema::new(
767            prune_schema
768                .column_schemas()
769                .iter()
770                .filter(|col| referenced_columns.contains(&col.name))
771                .map(|col| {
772                    if let Some(column_meta) = read_format.metadata().column_by_name(&col.name)
773                        && column_meta.semantic_type == SemanticType::Tag
774                        && col.data_type.is_string()
775                    {
776                        let field = Arc::new(Field::new(
777                            &col.name,
778                            col.data_type.as_arrow_type(),
779                            col.is_nullable(),
780                        ));
781                        let dict_field = tag_maybe_to_dictionary_field(&col.data_type, &field);
782                        let mut column = col.clone();
783                        column.data_type =
784                            ConcreteDataType::from_arrow_type(dict_field.data_type());
785                        return column;
786                    }
787
788                    col.clone()
789                })
790                .collect::<Vec<_>>(),
791        ));
792
793        let region_partition_physical_expr = region_partition_expr
794            .try_as_physical_expr(partition_schema.arrow_schema())
795            .context(SerializePartitionExprSnafu)?;
796
797        Ok(Some(PartitionFilterContext {
798            region_partition_physical_expr,
799            partition_schema,
800        }))
801    }
802
803    /// Reads parquet metadata of specific file.
804    /// Returns (fused metadata, cache_miss_flag).
805    pub(crate) async fn read_parquet_metadata(
806        &self,
807        file_path: &str,
808        file_size: u64,
809        cache_metrics: &mut MetadataCacheMetrics,
810        page_index_policy: PageIndexPolicy,
811    ) -> Result<(Arc<CachedSstMeta>, bool)> {
812        let start = Instant::now();
813        let _t = READ_STAGE_ELAPSED
814            .with_label_values(&["read_parquet_metadata"])
815            .start_timer();
816
817        let file_id = self.file_handle.file_id();
818        // Tries to get from cache with metrics tracking.
819        if let Some(metadata) = self
820            .cache_strategy
821            .get_sst_meta_data(file_id, cache_metrics, page_index_policy)
822            .await
823        {
824            cache_metrics.metadata_load_cost += start.elapsed();
825            return Ok((metadata, false));
826        }
827
828        // Cache miss, load metadata directly.
829        let mut metadata_loader =
830            MetadataLoader::new(self.object_store.clone(), file_path, file_size);
831        metadata_loader.with_page_index_policy(page_index_policy);
832        let metadata = metadata_loader.load(cache_metrics).await?;
833
834        let decoded = if self.cache_strategy.sst_meta_cache_enabled() {
835            let metadata = prepare_sst_meta(
836                file_path,
837                metadata,
838                None,
839                page_index_policy,
840                &self.cache_strategy.sst_meta_runtime(),
841            )
842            .await?;
843            let decoded = metadata.decoded();
844            match metadata {
845                SstMetaPreparation::Prepared(metadata) => {
846                    self.cache_strategy
847                        .put_prepared_sst_meta(file_id, metadata, true);
848                }
849                SstMetaPreparation::DecodedOnly { encoding_error, .. } => {
850                    warn!(
851                        encoding_error;
852                        "Failed to encode SST metadata for cache, using decoded metadata for {}",
853                        file_path
854                    );
855                }
856            }
857            decoded
858        } else {
859            Arc::new(CachedSstMeta::try_new_with_page_index_policy(
860                file_path,
861                metadata,
862                None,
863                page_index_policy,
864            )?)
865        };
866
867        cache_metrics.metadata_load_cost += start.elapsed();
868        Ok((decoded, true))
869    }
870
871    /// Computes row groups to read, along with their respective row selections.
872    #[tracing::instrument(
873        skip_all,
874        fields(
875            region_id = %self.file_handle.region_id(),
876            file_id = %self.file_handle.file_id()
877        )
878    )]
879    async fn row_groups_to_read(
880        &self,
881        read_format: &FlatReadFormat,
882        parquet_meta: &ParquetMetaData,
883        metrics: &mut ReaderFilterMetrics,
884    ) -> RowGroupSelection {
885        let num_row_groups = parquet_meta.num_row_groups();
886        let num_rows = parquet_meta.file_metadata().num_rows();
887        if num_row_groups == 0 || num_rows == 0 {
888            return RowGroupSelection::default();
889        }
890
891        // Let's assume that the number of rows in the first row group
892        // can represent the `row_group_size` of the Parquet file.
893        let row_group_size = parquet_meta.row_group(0).num_rows() as usize;
894        if row_group_size == 0 {
895            return RowGroupSelection::default();
896        }
897
898        metrics.rg_total += num_row_groups;
899        metrics.rows_total += num_rows as usize;
900
901        // Compute skip_fields once for all pruning operations
902        let skip_fields = self.pre_filter_mode.skip_fields();
903
904        let mut output = self.row_groups_by_minmax(
905            read_format,
906            parquet_meta,
907            row_group_size,
908            num_rows as usize,
909            metrics,
910            skip_fields,
911        );
912        if output.is_empty() {
913            return output;
914        }
915
916        let fulltext_filtered = self
917            .prune_row_groups_by_fulltext_index(
918                row_group_size,
919                num_row_groups,
920                &mut output,
921                metrics,
922                skip_fields,
923            )
924            .await;
925        if output.is_empty() {
926            return output;
927        }
928
929        self.prune_row_groups_by_inverted_index(
930            read_format.metadata(),
931            row_group_size,
932            parquet_meta,
933            &mut output,
934            metrics,
935            skip_fields,
936        )
937        .await;
938        if output.is_empty() {
939            return output;
940        }
941
942        self.prune_row_groups_by_bloom_filter(
943            read_format.metadata(),
944            row_group_size,
945            parquet_meta,
946            &mut output,
947            metrics,
948            skip_fields,
949        )
950        .await;
951        if output.is_empty() {
952            return output;
953        }
954
955        if !fulltext_filtered {
956            self.prune_row_groups_by_fulltext_bloom(
957                row_group_size,
958                parquet_meta,
959                &mut output,
960                metrics,
961                skip_fields,
962            )
963            .await;
964        }
965        output
966    }
967
968    /// Prunes row groups by fulltext index. Returns `true` if the row groups are pruned.
969    async fn prune_row_groups_by_fulltext_index(
970        &self,
971        row_group_size: usize,
972        num_row_groups: usize,
973        output: &mut RowGroupSelection,
974        metrics: &mut ReaderFilterMetrics,
975        skip_fields: bool,
976    ) -> bool {
977        if !self.file_handle.meta_ref().fulltext_index_available() {
978            return false;
979        }
980
981        let mut pruned = false;
982        // If skip_fields is true, only apply the first applier (for tags).
983        let appliers = if skip_fields {
984            &self.fulltext_index_appliers[..1]
985        } else {
986            &self.fulltext_index_appliers[..]
987        };
988        for index_applier in appliers.iter().flatten() {
989            let predicate_key = index_applier.predicate_key();
990            // Fast path: return early if the result is in the cache.
991            let cached = self
992                .cache_strategy
993                .index_result_cache()
994                .and_then(|cache| cache.get(predicate_key, self.file_handle.file_id().file_id()));
995            if let Some(result) = cached.as_ref()
996                && all_required_row_groups_searched(output, result)
997            {
998                apply_selection_and_update_metrics(output, result, metrics, INDEX_TYPE_FULLTEXT);
999                metrics.fulltext_index_cache_hit += 1;
1000                pruned = true;
1001                continue;
1002            }
1003
1004            // Slow path: apply the index from the file.
1005            metrics.fulltext_index_cache_miss += 1;
1006            let file_size_hint = self.file_handle.meta_ref().index_file_size();
1007            let apply_res = index_applier
1008                .apply_fine(
1009                    self.file_handle.index_id(),
1010                    Some(file_size_hint),
1011                    metrics.fulltext_index_apply_metrics.as_mut(),
1012                )
1013                .await;
1014            let selection = match apply_res {
1015                Ok(Some(res)) => {
1016                    RowGroupSelection::from_row_ids(res, row_group_size, num_row_groups)
1017                }
1018                Ok(None) => continue,
1019                Err(err) => {
1020                    handle_index_error!(err, self.file_handle, INDEX_TYPE_FULLTEXT);
1021                    continue;
1022                }
1023            };
1024
1025            self.apply_index_result_and_update_cache(
1026                predicate_key,
1027                self.file_handle.file_id().file_id(),
1028                selection,
1029                output,
1030                metrics,
1031                INDEX_TYPE_FULLTEXT,
1032            );
1033            pruned = true;
1034        }
1035        pruned
1036    }
1037
1038    /// Applies index to prune row groups.
1039    ///
1040    /// TODO(zhongzc): Devise a mechanism to enforce the non-use of indices
1041    /// as an escape route in case of index issues, and it can be used to test
1042    /// the correctness of the index.
1043    async fn prune_row_groups_by_inverted_index(
1044        &self,
1045        sst_metadata: &RegionMetadataRef,
1046        row_group_size: usize,
1047        parquet_meta: &ParquetMetaData,
1048        output: &mut RowGroupSelection,
1049        metrics: &mut ReaderFilterMetrics,
1050        skip_fields: bool,
1051    ) -> bool {
1052        if !self.file_handle.meta_ref().inverted_index_available() {
1053            return false;
1054        }
1055
1056        let num_row_groups = parquet_meta.num_row_groups();
1057        let total_row_count = parquet_meta.file_metadata().num_rows() as usize;
1058        let mut pruned = false;
1059        // If skip_fields is true, only apply the first applier (for tags).
1060        let appliers = if skip_fields {
1061            &self.inverted_index_appliers[..1]
1062        } else {
1063            &self.inverted_index_appliers[..]
1064        };
1065        for index_applier in appliers.iter().flatten() {
1066            let Ok(Some(plan)) = index_applier
1067                .plan_for_sst(sst_metadata)
1068                .inspect_err(|e| warn!(e; "failed to build compatible plan for sst"))
1069            else {
1070                continue;
1071            };
1072
1073            // Fast path: return early if the result is in the cache.
1074            let cached = self.cache_strategy.index_result_cache().and_then(|cache| {
1075                let file_id = self.file_handle.file_id().file_id();
1076                cache.get(&plan.predicate_key, file_id)
1077            });
1078
1079            if let Some(result) = cached.as_ref()
1080                && all_required_row_groups_searched(output, result)
1081            {
1082                apply_selection_and_update_metrics(output, result, metrics, INDEX_TYPE_INVERTED);
1083                metrics.inverted_index_cache_hit += 1;
1084                pruned = true;
1085                continue;
1086            }
1087
1088            // Slow path: apply the index from the file.
1089            metrics.inverted_index_cache_miss += 1;
1090            let file_size_hint = self.file_handle.meta_ref().index_file_size();
1091            let apply_res = index_applier
1092                .apply(
1093                    self.file_handle.index_id(),
1094                    Some(file_size_hint),
1095                    &plan.index_applier,
1096                    metrics.inverted_index_apply_metrics.as_mut(),
1097                )
1098                .await;
1099
1100            let selection = match apply_res {
1101                Ok(apply_output) => {
1102                    let index_row_count = apply_output.total_row_count;
1103                    let Some(selection) = RowGroupSelection::from_inverted_index_apply_output(
1104                        row_group_size,
1105                        num_row_groups,
1106                        total_row_count,
1107                        apply_output,
1108                    ) else {
1109                        warn!(
1110                            "Ignore inverted index with mismatched row count, file_id: {:?}, index_row_count: {}, parquet_row_count: {}",
1111                            self.file_handle.file_id(),
1112                            index_row_count,
1113                            total_row_count,
1114                        );
1115                        continue;
1116                    };
1117                    selection
1118                }
1119                Err(err) => {
1120                    handle_index_error!(err, self.file_handle, INDEX_TYPE_INVERTED);
1121                    continue;
1122                }
1123            };
1124
1125            self.apply_index_result_and_update_cache(
1126                &plan.predicate_key,
1127                self.file_handle.file_id().file_id(),
1128                selection,
1129                output,
1130                metrics,
1131                INDEX_TYPE_INVERTED,
1132            );
1133            pruned = true;
1134        }
1135        pruned
1136    }
1137
1138    async fn prune_row_groups_by_bloom_filter(
1139        &self,
1140        sst_metadata: &RegionMetadataRef,
1141        row_group_size: usize,
1142        parquet_meta: &ParquetMetaData,
1143        output: &mut RowGroupSelection,
1144        metrics: &mut ReaderFilterMetrics,
1145        skip_fields: bool,
1146    ) -> bool {
1147        if !self.file_handle.meta_ref().bloom_filter_index_available() {
1148            return false;
1149        }
1150
1151        let mut pruned = false;
1152        // If skip_fields is true, only apply the first applier (for tags).
1153        let appliers = if skip_fields {
1154            &self.bloom_filter_index_appliers[..1]
1155        } else {
1156            &self.bloom_filter_index_appliers[..]
1157        };
1158        for index_applier in appliers.iter().flatten() {
1159            let Some(compatible_predicates) =
1160                index_applier.compatible_predicate_for_sst(sst_metadata)
1161            else {
1162                continue;
1163            };
1164            let predicate_key = PredicateKey::new_bloom(compatible_predicates.clone());
1165            // Fast path: return early if the result is in the cache.
1166            let cached = self.cache_strategy.index_result_cache().and_then(|cache| {
1167                let file_id = self.file_handle.file_id().file_id();
1168                cache.get(&predicate_key, file_id)
1169            });
1170            if let Some(result) = cached.as_ref()
1171                && all_required_row_groups_searched(output, result)
1172            {
1173                apply_selection_and_update_metrics(output, result, metrics, INDEX_TYPE_BLOOM);
1174                metrics.bloom_filter_cache_hit += 1;
1175                pruned = true;
1176                continue;
1177            }
1178
1179            // Slow path: apply the index from the file.
1180            metrics.bloom_filter_cache_miss += 1;
1181            let file_size_hint = self.file_handle.meta_ref().index_file_size();
1182            let rgs = parquet_meta.row_groups().iter().enumerate().map(|(i, rg)| {
1183                (
1184                    rg.num_rows() as usize,
1185                    // Optimize: only search the row group that required by `output` and not stored in `cached`.
1186                    output.contains_non_empty_row_group(i)
1187                        && cached
1188                            .as_ref()
1189                            .map(|c| !c.contains_row_group(i))
1190                            .unwrap_or(true),
1191                )
1192            });
1193            let apply_res = index_applier
1194                .apply(
1195                    self.file_handle.index_id(),
1196                    Some(file_size_hint),
1197                    &compatible_predicates,
1198                    rgs,
1199                    metrics.bloom_filter_apply_metrics.as_mut(),
1200                )
1201                .await;
1202            let mut selection = match apply_res {
1203                Ok(apply_output) => {
1204                    RowGroupSelection::from_row_ranges(apply_output, row_group_size)
1205                }
1206                Err(err) => {
1207                    handle_index_error!(err, self.file_handle, INDEX_TYPE_BLOOM);
1208                    continue;
1209                }
1210            };
1211
1212            // New searched row groups are added to `selection`, concat them with `cached`.
1213            if let Some(cached) = cached.as_ref() {
1214                selection.concat(cached);
1215            }
1216
1217            self.apply_index_result_and_update_cache(
1218                &predicate_key,
1219                self.file_handle.file_id().file_id(),
1220                selection,
1221                output,
1222                metrics,
1223                INDEX_TYPE_BLOOM,
1224            );
1225            pruned = true;
1226        }
1227        pruned
1228    }
1229
1230    async fn prune_row_groups_by_fulltext_bloom(
1231        &self,
1232        row_group_size: usize,
1233        parquet_meta: &ParquetMetaData,
1234        output: &mut RowGroupSelection,
1235        metrics: &mut ReaderFilterMetrics,
1236        skip_fields: bool,
1237    ) -> bool {
1238        if !self.file_handle.meta_ref().fulltext_index_available() {
1239            return false;
1240        }
1241
1242        let mut pruned = false;
1243        // If skip_fields is true, only apply the first applier (for tags).
1244        let appliers = if skip_fields {
1245            &self.fulltext_index_appliers[..1]
1246        } else {
1247            &self.fulltext_index_appliers[..]
1248        };
1249        for index_applier in appliers.iter().flatten() {
1250            let predicate_key = index_applier.predicate_key();
1251            // Fast path: return early if the result is in the cache.
1252            let cached = self
1253                .cache_strategy
1254                .index_result_cache()
1255                .and_then(|cache| cache.get(predicate_key, self.file_handle.file_id().file_id()));
1256            if let Some(result) = cached.as_ref()
1257                && all_required_row_groups_searched(output, result)
1258            {
1259                apply_selection_and_update_metrics(output, result, metrics, INDEX_TYPE_FULLTEXT);
1260                metrics.fulltext_index_cache_hit += 1;
1261                pruned = true;
1262                continue;
1263            }
1264
1265            // Slow path: apply the index from the file.
1266            metrics.fulltext_index_cache_miss += 1;
1267            let file_size_hint = self.file_handle.meta_ref().index_file_size();
1268            let rgs = parquet_meta.row_groups().iter().enumerate().map(|(i, rg)| {
1269                (
1270                    rg.num_rows() as usize,
1271                    // Optimize: only search the row group that required by `output` and not stored in `cached`.
1272                    output.contains_non_empty_row_group(i)
1273                        && cached
1274                            .as_ref()
1275                            .map(|c| !c.contains_row_group(i))
1276                            .unwrap_or(true),
1277                )
1278            });
1279            let apply_res = index_applier
1280                .apply_coarse(
1281                    self.file_handle.index_id(),
1282                    Some(file_size_hint),
1283                    rgs,
1284                    metrics.fulltext_index_apply_metrics.as_mut(),
1285                )
1286                .await;
1287            let mut selection = match apply_res {
1288                Ok(Some(apply_output)) => {
1289                    RowGroupSelection::from_row_ranges(apply_output, row_group_size)
1290                }
1291                Ok(None) => continue,
1292                Err(err) => {
1293                    handle_index_error!(err, self.file_handle, INDEX_TYPE_FULLTEXT);
1294                    continue;
1295                }
1296            };
1297
1298            // New searched row groups are added to `selection`, concat them with `cached`.
1299            if let Some(cached) = cached.as_ref() {
1300                selection.concat(cached);
1301            }
1302
1303            self.apply_index_result_and_update_cache(
1304                predicate_key,
1305                self.file_handle.file_id().file_id(),
1306                selection,
1307                output,
1308                metrics,
1309                INDEX_TYPE_FULLTEXT,
1310            );
1311            pruned = true;
1312        }
1313        pruned
1314    }
1315
1316    /// Computes row groups selection after min-max pruning.
1317    fn row_groups_by_minmax(
1318        &self,
1319        read_format: &FlatReadFormat,
1320        parquet_meta: &ParquetMetaData,
1321        row_group_size: usize,
1322        total_row_count: usize,
1323        metrics: &mut ReaderFilterMetrics,
1324        skip_fields: bool,
1325    ) -> RowGroupSelection {
1326        let Some(predicate) = &self.predicate else {
1327            return RowGroupSelection::new(row_group_size, total_row_count);
1328        };
1329
1330        let file_id = self.file_handle.file_id().file_id();
1331        let index_result_cache = self.cache_strategy.index_result_cache();
1332        let cached_minmax_key =
1333            if index_result_cache.is_some() && predicate.dyn_filters().is_empty() {
1334                // Cache min-max pruning results keyed by predicate expressions. This avoids repeatedly
1335                // building row-group pruning stats for identical predicates across queries.
1336                let mut exprs = predicate
1337                    .exprs()
1338                    .iter()
1339                    .map(|expr| format!("{expr:?}"))
1340                    .collect::<Vec<_>>();
1341                exprs.sort();
1342                let schema_version = self
1343                    .expected_metadata
1344                    .as_ref()
1345                    .map(|meta| meta.schema_version)
1346                    .unwrap_or_else(|| read_format.metadata().schema_version);
1347                Some(PredicateKey::new_minmax(
1348                    Arc::new(exprs),
1349                    schema_version,
1350                    skip_fields,
1351                ))
1352            } else {
1353                None
1354            };
1355
1356        if let Some(index_result_cache) = index_result_cache
1357            && let Some(predicate_key) = cached_minmax_key.as_ref()
1358        {
1359            if let Some(result) = index_result_cache.get(predicate_key, file_id) {
1360                metrics.minmax_cache_hit += 1;
1361                let num_row_groups = parquet_meta.num_row_groups();
1362                metrics.rg_minmax_filtered +=
1363                    num_row_groups.saturating_sub(result.row_group_count());
1364                return (*result).clone();
1365            }
1366
1367            metrics.minmax_cache_miss += 1;
1368        }
1369
1370        let region_meta = read_format.metadata();
1371        let row_groups = parquet_meta.row_groups();
1372        let stats = RowGroupPruningStats::new(
1373            row_groups,
1374            read_format,
1375            self.expected_metadata.clone(),
1376            skip_fields,
1377        );
1378        let prune_schema = self
1379            .expected_metadata
1380            .as_ref()
1381            .map(|meta| meta.schema.arrow_schema())
1382            .unwrap_or_else(|| region_meta.schema.arrow_schema());
1383
1384        // Here we use the schema of the SST to build the physical expression. If the column
1385        // in the SST doesn't have the same column id as the column in the expected metadata,
1386        // we will get a None statistics for that column.
1387        let mask = predicate.prune_with_stats(&stats, prune_schema);
1388        let output = RowGroupSelection::from_full_row_group_ids(
1389            mask.iter()
1390                .enumerate()
1391                .filter_map(|(row_group, keep)| keep.then_some(row_group)),
1392            row_group_size,
1393            total_row_count,
1394        );
1395
1396        metrics.rg_minmax_filtered += parquet_meta
1397            .num_row_groups()
1398            .saturating_sub(output.row_group_count());
1399
1400        if let Some(index_result_cache) = index_result_cache
1401            && let Some(predicate_key) = cached_minmax_key
1402        {
1403            index_result_cache.put(predicate_key, file_id, Arc::new(output.clone()));
1404        }
1405
1406        output
1407    }
1408
1409    fn apply_index_result_and_update_cache(
1410        &self,
1411        predicate_key: &PredicateKey,
1412        file_id: FileId,
1413        result: RowGroupSelection,
1414        output: &mut RowGroupSelection,
1415        metrics: &mut ReaderFilterMetrics,
1416        index_type: &str,
1417    ) {
1418        apply_selection_and_update_metrics(output, &result, metrics, index_type);
1419
1420        if let Some(index_result_cache) = &self.cache_strategy.index_result_cache() {
1421            index_result_cache.put(predicate_key.clone(), file_id, Arc::new(result));
1422        }
1423    }
1424}
1425
1426fn has_row_level_selection(selection: &RowGroupSelection, parquet_meta: &ParquetMetaData) -> bool {
1427    selection.iter().any(|(row_group_idx, row_selection)| {
1428        let Some(row_group) = parquet_meta.row_groups().get(*row_group_idx) else {
1429            return false;
1430        };
1431
1432        row_selection.row_count() != row_group.num_rows() as usize
1433            || row_selection.iter().any(|selector| selector.skip)
1434    })
1435}
1436
1437fn apply_selection_and_update_metrics(
1438    output: &mut RowGroupSelection,
1439    result: &RowGroupSelection,
1440    metrics: &mut ReaderFilterMetrics,
1441    index_type: &str,
1442) {
1443    let intersection = output.intersect(result);
1444
1445    let row_group_count = output.row_group_count() - intersection.row_group_count();
1446    let row_count = output.row_count() - intersection.row_count();
1447
1448    metrics.update_index_metrics(index_type, row_group_count, row_count);
1449
1450    *output = intersection;
1451}
1452
1453fn all_required_row_groups_searched(
1454    required_row_groups: &RowGroupSelection,
1455    cached_row_groups: &RowGroupSelection,
1456) -> bool {
1457    required_row_groups.iter().all(|(rg_id, _)| {
1458        // Row group with no rows is not required to search.
1459        !required_row_groups.contains_non_empty_row_group(*rg_id)
1460            // The row group is already searched.
1461            || cached_row_groups.contains_row_group(*rg_id)
1462    })
1463}
1464
1465/// Metrics of filtering rows groups and rows.
1466#[derive(Debug, Default, Clone)]
1467pub(crate) struct ReaderFilterMetrics {
1468    /// Number of row groups before filtering.
1469    pub(crate) rg_total: usize,
1470    /// Number of row groups filtered by fulltext index.
1471    pub(crate) rg_fulltext_filtered: usize,
1472    /// Number of row groups filtered by inverted index.
1473    pub(crate) rg_inverted_filtered: usize,
1474    /// Number of row groups filtered by min-max index.
1475    pub(crate) rg_minmax_filtered: usize,
1476    /// Number of row groups filtered by bloom filter index.
1477    pub(crate) rg_bloom_filtered: usize,
1478
1479    /// Number of rows in row group before filtering.
1480    pub(crate) rows_total: usize,
1481    /// Number of rows in row group filtered by fulltext index.
1482    pub(crate) rows_fulltext_filtered: usize,
1483    /// Number of rows in row group filtered by inverted index.
1484    pub(crate) rows_inverted_filtered: usize,
1485    /// Number of rows in row group filtered by bloom filter index.
1486    pub(crate) rows_bloom_filtered: usize,
1487    /// Number of rows filtered by precise filter.
1488    pub(crate) rows_precise_filtered: usize,
1489
1490    /// Number of index result cache hits for fulltext index.
1491    pub(crate) fulltext_index_cache_hit: usize,
1492    /// Number of index result cache misses for fulltext index.
1493    pub(crate) fulltext_index_cache_miss: usize,
1494    /// Number of index result cache hits for inverted index.
1495    pub(crate) inverted_index_cache_hit: usize,
1496    /// Number of index result cache misses for inverted index.
1497    pub(crate) inverted_index_cache_miss: usize,
1498    /// Number of index result cache hits for bloom filter index.
1499    pub(crate) bloom_filter_cache_hit: usize,
1500    /// Number of index result cache misses for bloom filter index.
1501    pub(crate) bloom_filter_cache_miss: usize,
1502    /// Number of index result cache hits for minmax pruning.
1503    pub(crate) minmax_cache_hit: usize,
1504    /// Number of index result cache misses for minmax pruning.
1505    pub(crate) minmax_cache_miss: usize,
1506
1507    /// Optional metrics for inverted index applier.
1508    pub(crate) inverted_index_apply_metrics: Option<InvertedIndexApplyMetrics>,
1509    /// Optional metrics for bloom filter index applier.
1510    pub(crate) bloom_filter_apply_metrics: Option<BloomFilterIndexApplyMetrics>,
1511    /// Optional metrics for fulltext index applier.
1512    pub(crate) fulltext_index_apply_metrics: Option<FulltextIndexApplyMetrics>,
1513
1514    /// Number of pruner builder cache hits.
1515    pub(crate) pruner_cache_hit: usize,
1516    /// Number of pruner builder cache misses.
1517    pub(crate) pruner_cache_miss: usize,
1518    /// Duration spent waiting for pruner to build file ranges.
1519    pub(crate) pruner_prune_cost: Duration,
1520    /// Number of files filtered by manifest time-range pruning.
1521    pub(crate) files_time_range_pruned: usize,
1522}
1523
1524impl ReaderFilterMetrics {
1525    /// Adds `other` metrics to this metrics.
1526    pub(crate) fn merge_from(&mut self, other: &ReaderFilterMetrics) {
1527        self.rg_total += other.rg_total;
1528        self.rg_fulltext_filtered += other.rg_fulltext_filtered;
1529        self.rg_inverted_filtered += other.rg_inverted_filtered;
1530        self.rg_minmax_filtered += other.rg_minmax_filtered;
1531        self.rg_bloom_filtered += other.rg_bloom_filtered;
1532
1533        self.rows_total += other.rows_total;
1534        self.rows_fulltext_filtered += other.rows_fulltext_filtered;
1535        self.rows_inverted_filtered += other.rows_inverted_filtered;
1536        self.rows_bloom_filtered += other.rows_bloom_filtered;
1537        self.rows_precise_filtered += other.rows_precise_filtered;
1538
1539        self.fulltext_index_cache_hit += other.fulltext_index_cache_hit;
1540        self.fulltext_index_cache_miss += other.fulltext_index_cache_miss;
1541        self.inverted_index_cache_hit += other.inverted_index_cache_hit;
1542        self.inverted_index_cache_miss += other.inverted_index_cache_miss;
1543        self.bloom_filter_cache_hit += other.bloom_filter_cache_hit;
1544        self.bloom_filter_cache_miss += other.bloom_filter_cache_miss;
1545        self.minmax_cache_hit += other.minmax_cache_hit;
1546        self.minmax_cache_miss += other.minmax_cache_miss;
1547
1548        self.pruner_cache_hit += other.pruner_cache_hit;
1549        self.pruner_cache_miss += other.pruner_cache_miss;
1550        self.pruner_prune_cost += other.pruner_prune_cost;
1551        self.files_time_range_pruned += other.files_time_range_pruned;
1552
1553        // Merge optional applier metrics
1554        if let Some(other_metrics) = &other.inverted_index_apply_metrics {
1555            self.inverted_index_apply_metrics
1556                .get_or_insert_with(Default::default)
1557                .merge_from(other_metrics);
1558        }
1559        if let Some(other_metrics) = &other.bloom_filter_apply_metrics {
1560            self.bloom_filter_apply_metrics
1561                .get_or_insert_with(Default::default)
1562                .merge_from(other_metrics);
1563        }
1564        if let Some(other_metrics) = &other.fulltext_index_apply_metrics {
1565            self.fulltext_index_apply_metrics
1566                .get_or_insert_with(Default::default)
1567                .merge_from(other_metrics);
1568        }
1569    }
1570
1571    /// Reports metrics.
1572    pub(crate) fn observe(&self) {
1573        READ_ROW_GROUPS_TOTAL
1574            .with_label_values(&["before_filtering"])
1575            .inc_by(self.rg_total as u64);
1576        READ_ROW_GROUPS_TOTAL
1577            .with_label_values(&["fulltext_index_filtered"])
1578            .inc_by(self.rg_fulltext_filtered as u64);
1579        READ_ROW_GROUPS_TOTAL
1580            .with_label_values(&["inverted_index_filtered"])
1581            .inc_by(self.rg_inverted_filtered as u64);
1582        READ_ROW_GROUPS_TOTAL
1583            .with_label_values(&["minmax_index_filtered"])
1584            .inc_by(self.rg_minmax_filtered as u64);
1585        READ_ROW_GROUPS_TOTAL
1586            .with_label_values(&["bloom_filter_index_filtered"])
1587            .inc_by(self.rg_bloom_filtered as u64);
1588
1589        PRECISE_FILTER_ROWS_TOTAL
1590            .with_label_values(&["parquet"])
1591            .inc_by(self.rows_precise_filtered as u64);
1592        READ_ROWS_IN_ROW_GROUP_TOTAL
1593            .with_label_values(&["before_filtering"])
1594            .inc_by(self.rows_total as u64);
1595        READ_ROWS_IN_ROW_GROUP_TOTAL
1596            .with_label_values(&["fulltext_index_filtered"])
1597            .inc_by(self.rows_fulltext_filtered as u64);
1598        READ_ROWS_IN_ROW_GROUP_TOTAL
1599            .with_label_values(&["inverted_index_filtered"])
1600            .inc_by(self.rows_inverted_filtered as u64);
1601        READ_ROWS_IN_ROW_GROUP_TOTAL
1602            .with_label_values(&["bloom_filter_index_filtered"])
1603            .inc_by(self.rows_bloom_filtered as u64);
1604    }
1605
1606    fn update_index_metrics(&mut self, index_type: &str, row_group_count: usize, row_count: usize) {
1607        match index_type {
1608            INDEX_TYPE_FULLTEXT => {
1609                self.rg_fulltext_filtered += row_group_count;
1610                self.rows_fulltext_filtered += row_count;
1611            }
1612            INDEX_TYPE_INVERTED => {
1613                self.rg_inverted_filtered += row_group_count;
1614                self.rows_inverted_filtered += row_count;
1615            }
1616            INDEX_TYPE_BLOOM => {
1617                self.rg_bloom_filtered += row_group_count;
1618                self.rows_bloom_filtered += row_count;
1619            }
1620            _ => {}
1621        }
1622    }
1623}
1624
1625/// Metrics for parquet metadata cache operations.
1626#[derive(Default, Clone, Copy)]
1627pub struct MetadataCacheMetrics {
1628    /// Number of memory cache hits for parquet metadata.
1629    pub mem_cache_hit: usize,
1630    /// Number of file cache hits for parquet metadata.
1631    pub file_cache_hit: usize,
1632    /// Number of cache misses for parquet metadata.
1633    pub cache_miss: usize,
1634    /// Duration to load parquet metadata.
1635    pub metadata_load_cost: Duration,
1636    /// Number of read operations performed.
1637    pub num_reads: usize,
1638    /// Total bytes read from storage.
1639    pub bytes_read: u64,
1640}
1641
1642impl std::fmt::Debug for MetadataCacheMetrics {
1643    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1644        let Self {
1645            mem_cache_hit,
1646            file_cache_hit,
1647            cache_miss,
1648            metadata_load_cost,
1649            num_reads,
1650            bytes_read,
1651        } = self;
1652
1653        if self.is_empty() {
1654            return write!(f, "{{}}");
1655        }
1656        write!(f, "{{")?;
1657
1658        write!(f, "\"metadata_load_cost\":\"{:?}\"", metadata_load_cost)?;
1659
1660        if *mem_cache_hit > 0 {
1661            write!(f, ", \"mem_cache_hit\":{}", mem_cache_hit)?;
1662        }
1663        if *file_cache_hit > 0 {
1664            write!(f, ", \"file_cache_hit\":{}", file_cache_hit)?;
1665        }
1666        if *cache_miss > 0 {
1667            write!(f, ", \"cache_miss\":{}", cache_miss)?;
1668        }
1669        if *num_reads > 0 {
1670            write!(f, ", \"num_reads\":{}", num_reads)?;
1671        }
1672        if *bytes_read > 0 {
1673            write!(f, ", \"bytes_read\":{}", bytes_read)?;
1674        }
1675
1676        write!(f, "}}")
1677    }
1678}
1679
1680impl MetadataCacheMetrics {
1681    /// Returns true if the metrics are empty (contain no meaningful data).
1682    pub(crate) fn is_empty(&self) -> bool {
1683        self.metadata_load_cost.is_zero()
1684    }
1685
1686    /// Adds `other` metrics to this metrics.
1687    pub(crate) fn merge_from(&mut self, other: &MetadataCacheMetrics) {
1688        self.mem_cache_hit += other.mem_cache_hit;
1689        self.file_cache_hit += other.file_cache_hit;
1690        self.cache_miss += other.cache_miss;
1691        self.metadata_load_cost += other.metadata_load_cost;
1692        self.num_reads += other.num_reads;
1693        self.bytes_read += other.bytes_read;
1694    }
1695}
1696
1697/// Parquet reader metrics.
1698#[derive(Debug, Default, Clone)]
1699pub struct ReaderMetrics {
1700    /// Filtered row groups and rows metrics.
1701    pub(crate) filter_metrics: ReaderFilterMetrics,
1702    /// Duration to build the parquet reader.
1703    pub(crate) build_cost: Duration,
1704    /// Duration to scan the reader.
1705    pub(crate) scan_cost: Duration,
1706    /// Number of record batches read.
1707    pub(crate) num_record_batches: usize,
1708    /// Number of batches decoded.
1709    pub(crate) num_batches: usize,
1710    /// Number of rows read.
1711    pub(crate) num_rows: usize,
1712    /// Metrics for parquet metadata cache.
1713    pub(crate) metadata_cache_metrics: MetadataCacheMetrics,
1714    /// Optional metrics for page/row group fetch operations.
1715    pub(crate) fetch_metrics: Option<Arc<ParquetFetchMetrics>>,
1716    /// Memory size of metadata loaded for building file ranges.
1717    pub(crate) metadata_mem_size: isize,
1718    /// Number of file range builders created.
1719    pub(crate) num_range_builders: isize,
1720}
1721
1722impl ReaderMetrics {
1723    /// Adds `other` metrics to this metrics.
1724    pub(crate) fn merge_from(&mut self, other: &ReaderMetrics) {
1725        self.filter_metrics.merge_from(&other.filter_metrics);
1726        self.build_cost += other.build_cost;
1727        self.scan_cost += other.scan_cost;
1728        self.num_record_batches += other.num_record_batches;
1729        self.num_batches += other.num_batches;
1730        self.num_rows += other.num_rows;
1731        self.metadata_cache_metrics
1732            .merge_from(&other.metadata_cache_metrics);
1733        if let Some(other_fetch) = &other.fetch_metrics {
1734            if let Some(self_fetch) = &self.fetch_metrics {
1735                self_fetch.merge_from(other_fetch);
1736            } else {
1737                self.fetch_metrics = Some(other_fetch.clone());
1738            }
1739        }
1740        self.metadata_mem_size += other.metadata_mem_size;
1741        self.num_range_builders += other.num_range_builders;
1742    }
1743
1744    /// Reports total rows.
1745    pub(crate) fn observe_rows(&self, read_type: &str) {
1746        READ_ROWS_TOTAL
1747            .with_label_values(&[read_type])
1748            .inc_by(self.num_rows as u64);
1749    }
1750}
1751
1752/// Builder to build a parquet record batch stream for a row group.
1753pub(crate) struct RowGroupReaderBuilder {
1754    /// SST file to read.
1755    ///
1756    /// Holds the file handle to avoid the file purge it.
1757    file_handle: FileHandle,
1758    /// Path of the file.
1759    file_path: String,
1760    /// Metadata of the parquet file.
1761    parquet_meta: Arc<ParquetMetaData>,
1762    /// Immutable metadata size, computed once when the footer is decoded.
1763    parquet_metadata_size: usize,
1764    /// Arrow reader metadata for building async stream.
1765    arrow_metadata: ArrowReaderMetadata,
1766    /// Projected output schema aligned with `projection.projected_root_presence`.
1767    output_schema: SchemaRef,
1768    /// JSON2 columns that must be semantically rewritten into the projected target layout.
1769    json2_rewrite_targets: HashMap<String, Json2TargetLayout>,
1770    /// Object store as an Operator.
1771    object_store: ObjectStore,
1772    /// Projection mask.
1773    projection: ProjectionMaskPlan,
1774    /// Whether projected read columns include nested paths.
1775    has_nested_projection: bool,
1776    /// Cache.
1777    cache_strategy: CacheStrategy,
1778    /// Pre-built prefilter state. `None` if prefiltering is not applicable.
1779    prefilter_builder: Option<PrefilterContextBuilder>,
1780    /// Hint for rows in a decoded batch.
1781    batch_size: usize,
1782    /// Optional shared quota for projected row-group decoding.
1783    memory_tracker: Option<QueryMemoryTracker>,
1784}
1785
1786/// Context passed to [RowGroupReaderBuilder::build()] carrying all information
1787/// needed for prefiltering decisions.
1788pub(crate) struct RowGroupBuildContext<'a> {
1789    /// Index of the row group to read.
1790    pub(crate) row_group_idx: usize,
1791    /// Row selection for the row group. `None` means all rows.
1792    pub(crate) row_selection: Option<RowSelection>,
1793    /// Metrics for tracking fetch operations.
1794    pub(crate) fetch_metrics: Option<&'a ParquetFetchMetrics>,
1795}
1796
1797impl RowGroupReaderBuilder {
1798    /// Path of the file to read.
1799    pub(crate) fn file_path(&self) -> &str {
1800        &self.file_path
1801    }
1802
1803    /// Handle of the file to read.
1804    pub(crate) fn file_handle(&self) -> &FileHandle {
1805        &self.file_handle
1806    }
1807
1808    pub(crate) fn parquet_metadata(&self) -> &Arc<ParquetMetaData> {
1809        &self.parquet_meta
1810    }
1811
1812    pub(crate) fn parquet_metadata_size(&self) -> usize {
1813        self.parquet_metadata_size
1814    }
1815
1816    pub(crate) fn cache_strategy(&self) -> &CacheStrategy {
1817        &self.cache_strategy
1818    }
1819
1820    pub(crate) fn has_predicate_prefilter(&self) -> bool {
1821        self.prefilter_builder.is_some()
1822    }
1823
1824    /// Builds a parquet record batch stream to read the row group at `row_group_idx`.
1825    ///
1826    /// If prefiltering is applicable (based on `build_ctx`), this performs a two-phase read:
1827    /// 1. Reads only the prefilter columns (e.g. PK column), applies filters to get a refined row selection
1828    /// 2. Reads the full projection with the refined row selection
1829    ///
1830    /// The prefilter pass is *best-effort pruning*, not the precise filter for the query.
1831    /// Predicates that cannot be lowered to prefilter columns (column not projected,
1832    /// expression not supported, etc.) are silently skipped. Correctness rests on the
1833    /// DataFusion `FilterExec` above this reader, which always re-applies the original
1834    /// predicate. With predicate prefiltering enabled, tag and timestamp predicates that
1835    /// flow through [`SimpleFilterEvaluator`] are enforced precisely in this pass. See
1836    /// [`build_reader_filter_plan`] for the bucketing rules and disabled mode.
1837    ///
1838    /// When the prefilter result selects no rows, the second read still issues but
1839    /// parquet-rs short-circuits before any column-chunk IO: the row-group state machine
1840    /// jumps to `Finished` once it sees `num_rows_selected() == 0`, so no fast path is
1841    /// added here.
1842    pub(crate) async fn build(
1843        &self,
1844        build_ctx: RowGroupBuildContext<'_>,
1845    ) -> Result<ProjectedRecordBatchStream> {
1846        let prefilter_ctx = self
1847            .prefilter_builder
1848            .as_ref()
1849            .map(|b| b.build(build_ctx.row_group_idx));
1850
1851        let Some(mut prefilter_ctx) = prefilter_ctx else {
1852            // No prefilter applicable, build stream with full projection.
1853            let stream = self
1854                .build_with_projection(
1855                    build_ctx.row_group_idx,
1856                    build_ctx.row_selection,
1857                    self.projection.mask.clone(),
1858                    build_ctx.fetch_metrics,
1859                )
1860                .await?;
1861            return self.make_projected_stream(stream);
1862        };
1863
1864        let prefilter_start = Instant::now();
1865        let prefilter_result = execute_prefilter(&mut prefilter_ctx, self, &build_ctx).await?;
1866        if let Some(metrics) = build_ctx.fetch_metrics {
1867            let mut data = metrics.data.lock().unwrap();
1868            data.prefilter_cost += prefilter_start.elapsed();
1869            data.prefilter_filtered_rows += prefilter_result.filtered_rows;
1870        }
1871
1872        let refined_selection = Some(prefilter_result.refined_selection);
1873
1874        let stream = self
1875            .build_with_projection(
1876                build_ctx.row_group_idx,
1877                refined_selection,
1878                self.projection.mask.clone(),
1879                build_ctx.fetch_metrics,
1880            )
1881            .await?;
1882        self.make_projected_stream(stream)
1883    }
1884
1885    /// Builds the normal projection without running the generic predicate prefilter.
1886    ///
1887    /// The series reader uses this after computing its own primary-key-only row
1888    /// selection.
1889    pub(crate) async fn build_without_prefilter(
1890        &self,
1891        build_ctx: RowGroupBuildContext<'_>,
1892    ) -> Result<ProjectedRecordBatchStream> {
1893        let stream = self
1894            .build_with_projection(
1895                build_ctx.row_group_idx,
1896                build_ctx.row_selection,
1897                self.projection.mask.clone(),
1898                build_ctx.fetch_metrics,
1899            )
1900            .await?;
1901        self.make_projected_stream(stream)
1902    }
1903
1904    /// Builds a stream that reads only the encoded primary-key column.
1905    ///
1906    /// It preserves the normal reader's binary-or-dictionary decision. This path deliberately
1907    /// skips the normal prefilter pass: the caller reads `__primary_key` once and applies all
1908    /// encoded-primary-key filters to the returned batches.
1909    pub(crate) async fn build_primary_key(
1910        &self,
1911        build_ctx: RowGroupBuildContext<'_>,
1912    ) -> Result<ProjectedRecordBatchStream> {
1913        let parquet_schema = self.parquet_meta.file_metadata().schema_descr();
1914        let primary_key_index = parquet_schema
1915            .columns()
1916            .iter()
1917            .position(|column| column.name() == PRIMARY_KEY_COLUMN_NAME)
1918            .context(UnexpectedSnafu {
1919                reason: "SST does not contain __primary_key",
1920            })?;
1921        let projection = ProjectionMask::leaves(parquet_schema, [primary_key_index]);
1922
1923        self.build_with_projection(
1924            build_ctx.row_group_idx,
1925            build_ctx.row_selection,
1926            projection,
1927            build_ctx.fetch_metrics,
1928        )
1929        .await
1930    }
1931
1932    fn make_projected_stream(
1933        &self,
1934        stream: ProjectedRecordBatchStream,
1935    ) -> Result<ProjectedRecordBatchStream> {
1936        if !self.has_nested_projection && self.json2_rewrite_targets.is_empty() {
1937            return Ok(stream);
1938        }
1939
1940        let mode = if self.json2_rewrite_targets.is_empty() {
1941            AlignMode::AlignToSchema
1942        } else {
1943            AlignMode::Rewrite {
1944                columns: self.json2_rewrite_targets.clone(),
1945            }
1946        };
1947
1948        Ok(JsonSchemaAligner::new(
1949            stream,
1950            self.projection.projected_root_presence.clone(),
1951            self.output_schema.clone(),
1952            mode,
1953        )?
1954        .boxed())
1955    }
1956
1957    /// Builds a parquet record batch stream with a custom projection mask.
1958    pub(crate) async fn build_with_projection(
1959        &self,
1960        row_group_idx: usize,
1961        row_selection: Option<RowSelection>,
1962        projection: ProjectionMask,
1963        fetch_metrics: Option<&ParquetFetchMetrics>,
1964    ) -> Result<ProjectedRecordBatchStream> {
1965        let range_fetcher = SstParquetRangeFetcher::new(
1966            self.file_handle.file_id(),
1967            self.file_path.clone(),
1968            self.object_store.clone(),
1969            self.cache_strategy.clone(),
1970            row_group_idx,
1971            fetch_metrics.cloned(),
1972        )
1973        .with_memory_tracker(self.memory_tracker.clone());
1974
1975        build_sst_parquet_record_batch_stream(
1976            self.arrow_metadata.clone(),
1977            row_group_idx,
1978            row_selection,
1979            projection,
1980            range_fetcher,
1981            self.file_path.clone(),
1982            self.batch_size,
1983        )
1984    }
1985}
1986
1987#[derive(Clone)]
1988/// The filter to evaluate or the prune result of the default value.
1989pub(crate) enum MaybeFilter {
1990    /// The filter to evaluate.
1991    Filter(SimpleFilterEvaluator),
1992    /// The filter matches the default value.
1993    Matched,
1994    /// The filter is pruned.
1995    Pruned,
1996}
1997
1998impl MaybeFilter {
1999    /// Returns the inner filter when it is available.
2000    pub(crate) fn as_filter(&self) -> Option<&SimpleFilterEvaluator> {
2001        match self {
2002            MaybeFilter::Filter(filter) => Some(filter),
2003            MaybeFilter::Matched | MaybeFilter::Pruned => None,
2004        }
2005    }
2006}
2007
2008#[derive(Clone)]
2009/// Context to evaluate the column filter for a parquet file.
2010pub(crate) struct SimpleFilterContext {
2011    /// Filter to evaluate.
2012    filter: MaybeFilter,
2013    /// Debug string of the original logical expression.
2014    expr_str: String,
2015    /// Id of the column to evaluate.
2016    column_id: ColumnId,
2017    /// Semantic type of the column.
2018    semantic_type: SemanticType,
2019}
2020
2021impl SimpleFilterContext {
2022    /// Creates a context for the `expr`.
2023    ///
2024    /// Returns None if the column to filter doesn't exist in the SST metadata or the
2025    /// expected metadata.
2026    pub(crate) fn new_opt(
2027        sst_meta: &RegionMetadataRef,
2028        expected_meta: Option<&RegionMetadata>,
2029        expr: &Expr,
2030    ) -> Option<Self> {
2031        let filter = SimpleFilterEvaluator::try_new_with_column_type(expr, &|name| {
2032            expected_meta
2033                .unwrap_or(sst_meta)
2034                .column_by_name(name)
2035                .map(|column| column.column_schema.data_type.as_arrow_type())
2036        })?;
2037        let expr_str = format!("{expr:?}");
2038        let (column_metadata, maybe_filter) = match expected_meta {
2039            Some(meta) => {
2040                // Gets the column metadata from the expected metadata.
2041                let column = meta.column_by_name(filter.column_name())?;
2042                // Checks if the column is present in the SST metadata. We still uses the
2043                // column from the expected metadata.
2044                match sst_meta.column_by_id(column.column_id) {
2045                    Some(sst_column) => {
2046                        debug_assert_eq!(column.semantic_type, sst_column.semantic_type);
2047                        let maybe_filter = if sst_column.column_schema.data_type
2048                            == column.column_schema.data_type
2049                        {
2050                            MaybeFilter::Filter(filter)
2051                        } else {
2052                            // Schema evolution (altered field columns, or a
2053                            // widened time index unit) can make columns with
2054                            // the same id have different types across SSTs;
2055                            // evaluating the original filter may raise an
2056                            // invalid cross-type comparison. Timestamp
2057                            // predicates are not re-applied above this scan,
2058                            // so cast them into the file's unit instead of
2059                            // dropping them; other mismatches keep the
2060                            // conservative skip (the query layer re-applies
2061                            // those predicates).
2062                            match filter.cast_timestamp_unit(&sst_column.column_schema.data_type) {
2063                                Some(TimestampUnitCast::Filter(filter)) => {
2064                                    MaybeFilter::Filter(filter)
2065                                }
2066                                Some(TimestampUnitCast::Pruned) => MaybeFilter::Pruned,
2067                                Some(TimestampUnitCast::Matched) => MaybeFilter::Matched,
2068                                None => return None,
2069                            }
2070                        };
2071                        (column, maybe_filter)
2072                    }
2073                    None => {
2074                        // If the column is not present in the SST metadata, we evaluate the filter
2075                        // against the default value of the column.
2076                        // If we can't evaluate the filter, we return None.
2077                        if pruned_by_default(&filter, column)? {
2078                            (column, MaybeFilter::Pruned)
2079                        } else {
2080                            (column, MaybeFilter::Matched)
2081                        }
2082                    }
2083                }
2084            }
2085            None => {
2086                let column = sst_meta.column_by_name(filter.column_name())?;
2087                (column, MaybeFilter::Filter(filter))
2088            }
2089        };
2090
2091        Some(Self {
2092            filter: maybe_filter,
2093            expr_str,
2094            column_id: column_metadata.column_id,
2095            semantic_type: column_metadata.semantic_type,
2096        })
2097    }
2098
2099    /// Returns the filter to evaluate.
2100    pub(crate) fn filter(&self) -> &MaybeFilter {
2101        &self.filter
2102    }
2103
2104    /// Returns the original logical expression string.
2105    pub(crate) fn expr_str(&self) -> &str {
2106        &self.expr_str
2107    }
2108
2109    /// Returns the column id.
2110    pub(crate) fn column_id(&self) -> ColumnId {
2111        self.column_id
2112    }
2113
2114    /// Returns the semantic type of the column.
2115    pub(crate) fn semantic_type(&self) -> SemanticType {
2116        self.semantic_type
2117    }
2118}
2119
2120/// Context to evaluate a physical expression for a parquet file.
2121#[derive(Clone)]
2122pub(crate) struct PhysicalFilterContext {
2123    /// Filter to evaluate.
2124    filter: Arc<dyn PhysicalExpr>,
2125    /// Debug string of the original logical expression.
2126    expr_str: String,
2127    /// Id of the column to evaluate.
2128    column_id: ColumnId,
2129    /// Name of the column to evaluate.
2130    column_name: String,
2131    /// Semantic type of the column.
2132    semantic_type: SemanticType,
2133    /// Schema containing only the referenced column.
2134    schema: SchemaRef,
2135    /// Whether the original logical expression is immutable across queries.
2136    immutable: bool,
2137}
2138
2139impl PhysicalFilterContext {
2140    /// Creates a context for the `expr`.
2141    ///
2142    /// Returns None if the expression doesn't reference exactly one column or the
2143    /// column to filter doesn't exist in the SST metadata or the expected metadata.
2144    pub(crate) fn new_opt(
2145        sst_meta: &RegionMetadataRef,
2146        expected_meta: Option<&RegionMetadata>,
2147        read_format: &FlatReadFormat,
2148        expr: &Expr,
2149    ) -> Option<Self> {
2150        if !Self::is_prefilter_candidate(expr) {
2151            return None;
2152        }
2153        let expr_str = format!("{expr:?}");
2154        let column_name = Self::single_column_name(expr)?;
2155        let column_metadata = match expected_meta {
2156            Some(meta) => {
2157                let column = meta.column_by_name(&column_name)?;
2158                let sst_column = sst_meta.column_by_id(column.column_id)?;
2159                // Physical expr requires the column name to match the SST column name.
2160                if sst_column.column_schema.name != column_name {
2161                    return None;
2162                }
2163                // Schema evolution (an altered field column, or a widened time
2164                // index unit) can make the column's file type differ from the
2165                // expected type; evaluating the original expr against the
2166                // file's column would raise a cross-type comparison error
2167                // (e.g. Timestamp(ms) >= Timestamp(µs)). Drop the prefilter;
2168                // the query layer re-applies the predicate above the scan.
2169                if sst_column.column_schema.data_type != column.column_schema.data_type {
2170                    return None;
2171                }
2172                column
2173            }
2174            None => sst_meta.column_by_name(&column_name)?,
2175        };
2176
2177        // The column must be present in the projected arrow schema for the
2178        // prefilter to be able to read it.
2179        let (_, field) = read_format.arrow_schema().column_with_name(&column_name)?;
2180        let field = field.clone();
2181        let schema = Arc::new(ArrowSchema::new(vec![field]));
2182        let physical_expr = Predicate::to_physical_expr(expr, &schema)
2183            .inspect_err(|e| {
2184                error!(e; "Unable to build physical filter for {expr}, schema: {schema:?}");
2185            })
2186            .ok()?;
2187        let immutable = expr_is_immutable(expr);
2188
2189        Some(Self {
2190            filter: physical_expr,
2191            expr_str,
2192            column_id: column_metadata.column_id,
2193            column_name,
2194            semantic_type: column_metadata.semantic_type,
2195            schema,
2196            immutable,
2197        })
2198    }
2199
2200    /// Returns true if the expression is a variant we want to evaluate as a
2201    /// physical prefilter. Binary exprs are intentionally excluded because
2202    /// [`SimpleFilterEvaluator`] already handles them.
2203    // TODO(yingwen): extend more expressions if necessary. For example, allow some cheap scalar functions (e.g. `lower`, `length`, date truncations)
2204    fn is_prefilter_candidate(expr: &Expr) -> bool {
2205        if !matches!(
2206            expr,
2207            Expr::InList(_) | Expr::IsNull(_) | Expr::IsNotNull(_) | Expr::Between(_)
2208        ) {
2209            return false;
2210        }
2211
2212        // If any functions are found in the expr, it will be not considered as worthy enough to
2213        // be evaluated in the prefilter. At last, prefilter reads the Parquet files one more time.
2214        !expr
2215            .exists(|e| Ok(matches!(e, Expr::ScalarFunction(_))))
2216            .unwrap_or(false)
2217    }
2218
2219    fn single_column_name(expr: &Expr) -> Option<String> {
2220        let mut columns = HashSet::new();
2221        if expr_to_columns(expr, &mut columns).is_err() {
2222            return None;
2223        }
2224        if columns.len() != 1 {
2225            return None;
2226        }
2227        columns.iter().next().map(|column| column.name.clone())
2228    }
2229
2230    /// Returns the filter to evaluate.
2231    pub(crate) fn filter(&self) -> &Arc<dyn PhysicalExpr> {
2232        &self.filter
2233    }
2234
2235    /// Returns the original logical expression string.
2236    pub(crate) fn expr_str(&self) -> &str {
2237        &self.expr_str
2238    }
2239
2240    /// Returns the column id.
2241    pub(crate) fn column_id(&self) -> ColumnId {
2242        self.column_id
2243    }
2244
2245    /// Returns the column name.
2246    pub(crate) fn column_name(&self) -> &str {
2247        &self.column_name
2248    }
2249
2250    /// Returns the semantic type of the column.
2251    pub(crate) fn semantic_type(&self) -> SemanticType {
2252        self.semantic_type
2253    }
2254
2255    /// Returns the schema containing only the referenced column.
2256    pub(crate) fn schema(&self) -> &SchemaRef {
2257        &self.schema
2258    }
2259
2260    /// Returns true if the original logical expression is immutable across queries.
2261    pub(crate) fn is_immutable(&self) -> bool {
2262        self.immutable
2263    }
2264}
2265
2266fn expr_is_immutable(expr: &Expr) -> bool {
2267    let mut is_immutable = true;
2268    let _ = expr.apply(|expr| match expr {
2269        Expr::ScalarFunction(function)
2270            if function.func.signature().volatility != Volatility::Immutable =>
2271        {
2272            is_immutable = false;
2273            Ok(TreeNodeRecursion::Stop)
2274        }
2275        Expr::ScalarVariable(_, _) => {
2276            is_immutable = false;
2277            Ok(TreeNodeRecursion::Stop)
2278        }
2279        _ => Ok(TreeNodeRecursion::Continue),
2280    });
2281    is_immutable
2282}
2283
2284/// Prune a column by its default value.
2285/// Returns false if we can't create the default value or evaluate the filter.
2286fn pruned_by_default(filter: &SimpleFilterEvaluator, column: &ColumnMetadata) -> Option<bool> {
2287    let value = column.column_schema.create_default().ok().flatten()?;
2288    let scalar_value = value
2289        .try_to_scalar_value(&column.column_schema.data_type)
2290        .ok()?;
2291    let matches = filter.evaluate_scalar(&scalar_value).ok()?;
2292    Some(!matches)
2293}
2294
2295/// Parquet batch reader to read our SST format.
2296pub struct ParquetReader {
2297    /// File range context.
2298    context: FileRangeContextRef,
2299    /// Row group selection to read.
2300    selection: RowGroupSelection,
2301    /// Reader of current row group.
2302    reader: Option<FlatPruneReader>,
2303    /// Metrics for tracking row group fetch operations.
2304    fetch_metrics: ParquetFetchMetrics,
2305}
2306
2307impl ParquetReader {
2308    #[tracing::instrument(
2309        skip_all,
2310        fields(
2311            region_id = %self.context.reader_builder().file_handle.region_id(),
2312            file_id = %self.context.reader_builder().file_handle.file_id()
2313        )
2314    )]
2315    pub async fn next_record_batch(&mut self) -> Result<Option<RecordBatch>> {
2316        loop {
2317            if let Some(reader) = &mut self.reader {
2318                if let Some(batch) = reader.next_batch().await? {
2319                    return Ok(Some(batch));
2320                }
2321                self.reader = None;
2322                continue;
2323            }
2324
2325            let Some((row_group_idx, row_selection)) = self.selection.pop_first() else {
2326                return Ok(None);
2327            };
2328
2329            let skip_fields = self.context.pre_filter_mode().skip_fields();
2330            let parquet_reader = self
2331                .context
2332                .reader_builder()
2333                .build(self.context.build_context(
2334                    row_group_idx,
2335                    Some(row_selection),
2336                    Some(&self.fetch_metrics),
2337                ))
2338                .await?;
2339            self.reader = Some(FlatPruneReader::new_with_row_group_reader(
2340                self.context.clone(),
2341                FlatRowGroupReader::new(self.context.clone(), parquet_reader),
2342                skip_fields,
2343            ));
2344        }
2345    }
2346    /// Creates a new reader.
2347    #[tracing::instrument(
2348        skip_all,
2349        fields(
2350            region_id = %context.reader_builder().file_handle.region_id(),
2351            file_id = %context.reader_builder().file_handle.file_id()
2352        )
2353    )]
2354    pub(crate) async fn new(
2355        context: FileRangeContextRef,
2356        mut selection: RowGroupSelection,
2357    ) -> Result<Self> {
2358        let fetch_metrics = ParquetFetchMetrics::default();
2359        let reader = if let Some((row_group_idx, row_selection)) = selection.pop_first() {
2360            let skip_fields = context.pre_filter_mode().skip_fields();
2361            let parquet_reader = context
2362                .reader_builder()
2363                .build(context.build_context(
2364                    row_group_idx,
2365                    Some(row_selection),
2366                    Some(&fetch_metrics),
2367                ))
2368                .await?;
2369            Some(FlatPruneReader::new_with_row_group_reader(
2370                context.clone(),
2371                FlatRowGroupReader::new(context.clone(), parquet_reader),
2372                skip_fields,
2373            ))
2374        } else {
2375            None
2376        };
2377
2378        Ok(ParquetReader {
2379            context,
2380            selection,
2381            reader,
2382            fetch_metrics,
2383        })
2384    }
2385
2386    /// Returns the metadata of the SST.
2387    pub fn metadata(&self) -> &RegionMetadataRef {
2388        self.context.read_format().metadata()
2389    }
2390
2391    pub fn parquet_metadata(&self) -> Arc<ParquetMetaData> {
2392        self.context.reader_builder().parquet_meta.clone()
2393    }
2394}
2395
2396/// Reader to read a row group of a parquet file in flat format, returning RecordBatch.
2397pub(crate) struct FlatRowGroupReader {
2398    /// Context for file ranges.
2399    context: FileRangeContextRef,
2400    /// Inner parquet record batch stream.
2401    stream: ProjectedRecordBatchStream,
2402    /// Cached sequence array to override sequences.
2403    override_sequence: Option<ArrayRef>,
2404}
2405
2406impl FlatRowGroupReader {
2407    /// Creates a new flat reader from file range.
2408    pub(crate) fn new(context: FileRangeContextRef, stream: ProjectedRecordBatchStream) -> Self {
2409        // The batch length from the reader should be less than or equal to DEFAULT_READ_BATCH_SIZE.
2410        let override_sequence = context
2411            .read_format()
2412            .new_override_sequence_array(DEFAULT_READ_BATCH_SIZE);
2413
2414        Self {
2415            context,
2416            stream,
2417            override_sequence,
2418        }
2419    }
2420
2421    /// Returns the next RecordBatch.
2422    pub(crate) async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
2423        match self.stream.next().await {
2424            Some(batch_result) => {
2425                let record_batch = batch_result?;
2426
2427                let record_batch = self
2428                    .context
2429                    .read_format()
2430                    .convert_batch(record_batch, self.override_sequence.as_ref())?;
2431                Ok(Some(record_batch))
2432            }
2433            None => Ok(None),
2434        }
2435    }
2436}
2437
2438#[cfg(test)]
2439mod tests {
2440    use std::collections::HashMap;
2441    use std::fmt::{Debug, Formatter};
2442    use std::sync::{Arc, LazyLock};
2443
2444    use api::v1::OpType;
2445    use common_error::ext::WhateverResult;
2446    use common_function::scalars::json::json_get::JsonGetWithType;
2447    use common_function::scalars::udf::create_udf;
2448    use common_recordbatch::ext::RecordBatchExt;
2449    use datafusion::arrow::datatypes::DataType;
2450    use datafusion_common::ScalarValue;
2451    use datafusion_expr::expr::ScalarFunction;
2452    use datafusion_expr::{
2453        ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
2454        col, lit,
2455    };
2456    use datatypes::arrow::array::{
2457        ArrayRef, BinaryDictionaryBuilder, Int64Array, StringArray, StructArray,
2458        TimestampMillisecondArray, UInt8Array, UInt64Array,
2459    };
2460    use datatypes::arrow::datatypes::{Fields, Schema, UInt32Type};
2461    use datatypes::arrow::record_batch::RecordBatch;
2462    use datatypes::extension::json::Json2ExtensionType;
2463    use datatypes::prelude::ConcreteDataType;
2464    use datatypes::schema::ColumnSchema;
2465    use object_store::services::Memory;
2466    use parquet::arrow::ArrowWriter;
2467    use parquet::arrow::arrow_reader::RowSelector;
2468    use parquet::file::properties::WriterProperties;
2469    use store_api::codec::PrimaryKeyEncoding;
2470    use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
2471    use store_api::region_request::PathType;
2472    use store_api::storage::RegionId;
2473    use table::predicate::Predicate;
2474
2475    use super::*;
2476    use crate::cache::CacheManager;
2477    use crate::sst::parquet::metadata::MetadataLoader;
2478    use crate::sst::parquet::prefilter::{build_reader_filter_plan, execute_prefilter};
2479    use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns};
2480    use crate::sst::parquet::row_group::ParquetFetchMetrics;
2481    use crate::test_util::sst_util::{sst_file_handle, sst_region_metadata};
2482
2483    async fn prefilter_test_builder(
2484        object_store: ObjectStore,
2485        predicate: Predicate,
2486        cache_strategy: CacheStrategy,
2487    ) -> (RowGroupReaderBuilder, Arc<RegionMetadata>) {
2488        let metadata = Arc::new(
2489            crate::test_util::sst_util::sst_region_metadata_with_encoding(
2490                PrimaryKeyEncoding::Sparse,
2491            ),
2492        );
2493        let batch = |start: i64, end: i64| {
2494            let mut primary_key = BinaryDictionaryBuilder::<UInt32Type>::new();
2495            let mut fields = Vec::new();
2496            let mut timestamps = Vec::new();
2497            for value in start..end {
2498                let tag = if value == 4 { "b" } else { "a" };
2499                primary_key
2500                    .append(crate::test_util::sst_util::new_sparse_primary_key(
2501                        &[tag, "x"],
2502                        &metadata,
2503                        1,
2504                        100,
2505                    ))
2506                    .unwrap();
2507                fields.push(value as u64);
2508                timestamps.push(value);
2509            }
2510            RecordBatch::try_new(
2511                crate::sst::to_flat_sst_arrow_schema(
2512                    &metadata,
2513                    &crate::sst::FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse),
2514                ),
2515                vec![
2516                    Arc::new(UInt64Array::from(fields)) as ArrayRef,
2517                    Arc::new(TimestampMillisecondArray::from(timestamps)) as ArrayRef,
2518                    Arc::new(primary_key.finish()) as ArrayRef,
2519                    Arc::new(UInt64Array::from_value(1, (end - start) as usize)) as ArrayRef,
2520                    Arc::new(UInt8Array::from_value(
2521                        OpType::Put as u8,
2522                        (end - start) as usize,
2523                    )) as ArrayRef,
2524                ],
2525            )
2526            .unwrap()
2527        };
2528        let first_batch = batch(0, 3);
2529        let second_batch = batch(3, 6);
2530        let mut bytes = Vec::new();
2531        let mut writer = ArrowWriter::try_new(&mut bytes, first_batch.schema(), None).unwrap();
2532        writer.write(&first_batch).unwrap();
2533        writer.flush().unwrap();
2534        writer.write(&second_batch).unwrap();
2535        writer.close().unwrap();
2536
2537        let file_handle = sst_file_handle(0, 6);
2538        let file_path = file_handle.file_path("prefilter_test", PathType::Bare);
2539        let file_size = bytes.len() as u64;
2540        object_store.write(&file_path, bytes).await.unwrap();
2541
2542        let mut cache_metrics = MetadataCacheMetrics::default();
2543        let parquet_meta = Arc::new(
2544            MetadataLoader::new(object_store.clone(), &file_path, file_size)
2545                .load(&mut cache_metrics)
2546                .await
2547                .unwrap(),
2548        );
2549        let read_format = FlatReadFormat::new(
2550            metadata.clone(),
2551            ReadColumns::new(
2552                metadata
2553                    .column_metadatas
2554                    .iter()
2555                    .map(|column| column.column_id),
2556            ),
2557            None,
2558            &file_path,
2559            false,
2560        )
2561        .unwrap();
2562        let codec = build_primary_key_codec(metadata.as_ref());
2563        let filter_plan = build_reader_filter_plan(
2564            Some(&predicate),
2565            None,
2566            PreFilterMode::All,
2567            true,
2568            false,
2569            &read_format,
2570            &codec,
2571            &parquet_meta,
2572        );
2573        assert!(filter_plan.prefilter_builder.is_some());
2574
2575        let output_schema = read_format.arrow_schema().clone();
2576        let parquet_schema = parquet_meta.file_metadata().schema_descr();
2577        let source_schema = parquet_to_arrow_schema(
2578            parquet_schema,
2579            parquet_meta.file_metadata().key_value_metadata(),
2580        )
2581        .unwrap();
2582        let projection = build_projection_plan(
2583            read_format.parquet_read_columns(),
2584            parquet_schema,
2585            &source_schema,
2586        )
2587        .unwrap();
2588        let arrow_metadata =
2589            ArrowReaderMetadata::try_new(parquet_meta.clone(), ArrowReaderOptions::new()).unwrap();
2590        (
2591            RowGroupReaderBuilder {
2592                memory_tracker: None,
2593                file_handle: file_handle.clone(),
2594                file_path,
2595                parquet_meta,
2596                parquet_metadata_size: 0,
2597                arrow_metadata,
2598                output_schema,
2599                json2_rewrite_targets: HashMap::new(),
2600                object_store,
2601                projection,
2602                has_nested_projection: false,
2603                cache_strategy,
2604                prefilter_builder: filter_plan.prefilter_builder,
2605                batch_size: DEFAULT_READ_BATCH_SIZE,
2606            },
2607            metadata,
2608        )
2609    }
2610
2611    #[tokio::test(flavor = "current_thread")]
2612    async fn test_execute_prefilter_proven_filters_preserve_selection_without_fetching() {
2613        let object_store = ObjectStore::new(Memory::default()).unwrap();
2614        let predicate = Predicate::new(vec![col("field_0").gt_eq(lit(0_u64))]);
2615        let (reader_builder, _) =
2616            prefilter_test_builder(object_store, predicate, CacheStrategy::Disabled).await;
2617        let prefilter_builder = reader_builder.prefilter_builder.as_ref().unwrap();
2618
2619        for original_selection in [
2620            None,
2621            Some(RowSelection::from(vec![
2622                RowSelector::skip(1),
2623                RowSelector::select(1),
2624                RowSelector::skip(1),
2625            ])),
2626            Some(RowSelection::from(vec![])),
2627        ] {
2628            let mut prefilter_ctx = prefilter_builder.build(0);
2629            let fetch_metrics = ParquetFetchMetrics::default();
2630            let result = execute_prefilter(
2631                &mut prefilter_ctx,
2632                &reader_builder,
2633                &RowGroupBuildContext {
2634                    row_group_idx: 0,
2635                    row_selection: original_selection.clone(),
2636                    fetch_metrics: Some(&fetch_metrics),
2637                },
2638            )
2639            .await
2640            .unwrap();
2641
2642            let expected = original_selection.unwrap_or_else(|| {
2643                RowSelection::from(vec![RowSelector::select(
2644                    reader_builder.parquet_meta.row_group(0).num_rows() as usize,
2645                )])
2646            });
2647            assert_eq!(result.refined_selection, expected);
2648            assert_eq!(result.filtered_rows, 0);
2649            let metrics = fetch_metrics.data.lock().unwrap();
2650            assert_eq!(metrics.pages_to_fetch_store, 0);
2651            assert_eq!(metrics.pages_to_fetch_mem, 0);
2652            assert_eq!(metrics.pages_to_fetch_write_cache, 0);
2653        }
2654    }
2655
2656    #[tokio::test(flavor = "current_thread")]
2657    async fn test_execute_prefilter_mixed_filters_use_nonzero_row_group_and_cache() {
2658        let object_store = ObjectStore::new(Memory::default()).unwrap();
2659        let predicate = Predicate::new(vec![
2660            col("field_0").gt_eq(lit(3_u64)),
2661            col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(6), None))),
2662            col("field_0").in_list(vec![lit(3_u64), lit(4_u64)], false),
2663            col("tag_0").eq(lit("a")),
2664        ]);
2665        let cache = CacheStrategy::EnableAll(Arc::new(
2666            CacheManager::builder()
2667                .prefilter_result_cache_size(1024)
2668                .build(),
2669        ));
2670        let (reader_builder, _) =
2671            prefilter_test_builder(object_store, predicate, cache.clone()).await;
2672        let prefilter_builder = reader_builder.prefilter_builder.as_ref().unwrap();
2673
2674        for pass in 0..2 {
2675            let mut prefilter_ctx = prefilter_builder.build(1);
2676            let fetch_metrics = ParquetFetchMetrics::default();
2677            let fetch_metrics_ref = (pass == 1).then_some(&fetch_metrics);
2678            let result = execute_prefilter(
2679                &mut prefilter_ctx,
2680                &reader_builder,
2681                &RowGroupBuildContext {
2682                    row_group_idx: 1,
2683                    row_selection: Some(RowSelection::from(vec![
2684                        RowSelector::select(2),
2685                        RowSelector::skip(1),
2686                    ])),
2687                    fetch_metrics: fetch_metrics_ref,
2688                },
2689            )
2690            .await
2691            .unwrap();
2692            assert_eq!(result.filtered_rows, 1);
2693            assert_eq!(
2694                result.refined_selection,
2695                RowSelection::from(vec![RowSelector::select(1), RowSelector::skip(2),])
2696            );
2697            if pass == 1 {
2698                assert_eq!(fetch_metrics.data.lock().unwrap().pages_to_fetch_store, 0);
2699            }
2700        }
2701
2702        let disabled = prefilter_test_builder(
2703            ObjectStore::new(Memory::default()).unwrap(),
2704            Predicate::new(vec![
2705                col("field_0").gt_eq(lit(3_u64)),
2706                col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(6), None))),
2707                col("field_0").in_list(vec![lit(3_u64), lit(4_u64)], false),
2708                col("tag_0").eq(lit("a")),
2709            ]),
2710            CacheStrategy::Disabled,
2711        )
2712        .await;
2713        let mut prefilter_ctx = disabled.0.prefilter_builder.as_ref().unwrap().build(1);
2714        let result = execute_prefilter(
2715            &mut prefilter_ctx,
2716            &disabled.0,
2717            &RowGroupBuildContext {
2718                row_group_idx: 1,
2719                row_selection: None,
2720                fetch_metrics: None,
2721            },
2722        )
2723        .await
2724        .unwrap();
2725        assert_eq!(result.filtered_rows, 2);
2726        assert_eq!(
2727            result.refined_selection,
2728            RowSelection::from(vec![RowSelector::select(1), RowSelector::skip(2)])
2729        );
2730    }
2731
2732    #[test]
2733    fn test_skip_prefilter_for_json_get() -> WhateverResult<()> {
2734        fn json_get_expr(base: Expr, path: &str) -> Expr {
2735            let json_get = Arc::new(create_udf(Arc::new(JsonGetWithType::default())));
2736            Expr::ScalarFunction(ScalarFunction::new_udf(json_get, vec![base, lit(path)]))
2737        }
2738
2739        let metadata = Arc::new(sst_region_metadata());
2740        let format = FlatReadFormat::new(
2741            metadata.clone(),
2742            ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
2743            None,
2744            "test",
2745            true,
2746        )?;
2747        let new_filter =
2748            |expr: Expr| PhysicalFilterContext::new_opt(&metadata, None, &format, &expr);
2749
2750        let json_get = || json_get_expr(col("field_0"), "a.b");
2751
2752        let regular_expr = col("field_0").is_null();
2753        assert!(new_filter(regular_expr).is_some());
2754
2755        let is_null = json_get().is_null();
2756        assert!(new_filter(is_null).is_none());
2757
2758        let is_not_null = json_get().is_not_null();
2759        assert!(new_filter(is_not_null).is_none());
2760
2761        let in_list = json_get().in_list(vec![lit("value")], false);
2762        assert!(new_filter(in_list).is_none());
2763
2764        let in_list_nested = col("field_0").in_list(vec![json_get()], false);
2765        assert!(new_filter(in_list_nested).is_none());
2766
2767        let between = json_get().between(lit(1_u64), lit(10_u64));
2768        assert!(new_filter(between).is_none());
2769
2770        let between_nested = col("field_0").between(json_get(), lit(10_u64));
2771        assert!(new_filter(between_nested).is_none());
2772
2773        Ok(())
2774    }
2775
2776    /// A physical prefilter expr must be dropped when the column's file type
2777    /// differs from the expected type (e.g. an old-unit SST after the time
2778    /// index unit was widened): evaluating the expected-unit literals against
2779    /// the file's column raises a cross-unit comparison error.
2780    #[test]
2781    fn test_physical_filter_dropped_on_file_type_mismatch() {
2782        let file_metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
2783        // The region metadata after widening the `ts` unit to microsecond.
2784        let mut expected = (*file_metadata).clone();
2785        for column in expected.column_metadatas.iter_mut() {
2786            if column.column_schema.name == "ts" {
2787                column.column_schema.data_type = ConcreteDataType::timestamp_microsecond_datatype();
2788            }
2789        }
2790        let expected = Arc::new(expected);
2791
2792        let format = FlatReadFormat::new(
2793            file_metadata.clone(),
2794            ReadColumns::new(file_metadata.column_metadatas.iter().map(|c| c.column_id)),
2795            None,
2796            "test",
2797            true,
2798        )
2799        .unwrap();
2800
2801        let between = col("ts").between(
2802            lit(ScalarValue::TimestampMicrosecond(Some(1_000_000), None)),
2803            lit(ScalarValue::TimestampMicrosecond(Some(2_000_000), None)),
2804        );
2805        // Same type: the prefilter is kept.
2806        assert!(
2807            PhysicalFilterContext::new_opt(&file_metadata, Some(&file_metadata), &format, &between)
2808                .is_some()
2809        );
2810        // Widened unit: the prefilter is dropped.
2811        assert!(
2812            PhysicalFilterContext::new_opt(&file_metadata, Some(&expected), &format, &between)
2813                .is_none()
2814        );
2815    }
2816
2817    #[tokio::test]
2818    async fn test_nested_projection_reads_partial_json2_physical_fields() -> WhateverResult<()> {
2819        // Write a full JSON2-like Arrow struct:
2820        // j: { a: { x: int, y: string }, b: string }.
2821        // The test later requests only j.a.x and verifies that the physical Parquet projection
2822        // does not materialize j.a.y or j.b.
2823
2824        let xy_fields = Fields::from(vec![
2825            Arc::new(Field::new("x", DataType::Int64, true)),
2826            Arc::new(Field::new("y", DataType::Utf8, true)),
2827        ]);
2828        let a_field = Arc::new(Field::new("a", DataType::Struct(xy_fields.clone()), true));
2829        let b_field = Arc::new(Field::new("b", DataType::Utf8, true));
2830        let json_fields = Fields::from(vec![a_field, b_field]);
2831        let json_field = Field::new("j", DataType::Struct(json_fields.clone()), true)
2832            .with_extension_type(Json2ExtensionType::default());
2833        let schema = Arc::new(Schema::new(vec![json_field]));
2834
2835        let a_array = Arc::new(StructArray::new(
2836            xy_fields,
2837            vec![
2838                Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef,
2839                Arc::new(StringArray::from_iter_values(["x1", "x2", "x3"])) as ArrayRef,
2840            ],
2841            None,
2842        )) as ArrayRef;
2843        let b_array = Arc::new(StringArray::from_iter_values(["b1", "b2", "b3"])) as ArrayRef;
2844        let j_array =
2845            Arc::new(StructArray::new(json_fields, vec![a_array, b_array], None)) as ArrayRef;
2846        let columns = vec![j_array];
2847
2848        let batch = RecordBatch::try_new(schema, columns).map_err(|e| e.to_string())?;
2849
2850        // Persist the complete nested schema to an in-memory Parquet file so the projection is
2851        // exercised through parquet-rs rather than a mock.
2852
2853        let object_store = ObjectStore::new(Memory::default()).map_err(|e| e.to_string())?;
2854        let file_handle = sst_file_handle(0, 1);
2855        let file_path = file_handle.file_path("test_table", PathType::Bare);
2856
2857        let mut parquet_bytes = Vec::new();
2858        ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), None)
2859            .and_then(|mut w| {
2860                w.write(&batch)?;
2861                Ok(w)
2862            })
2863            .and_then(|w| w.close())
2864            .map_err(|e| e.to_string())?;
2865        let file_size = parquet_bytes.len() as u64;
2866        object_store
2867            .write(&file_path, parquet_bytes)
2868            .await
2869            .map_err(|e| e.to_string())?;
2870
2871        let mut cache_metrics = MetadataCacheMetrics::default();
2872        let loader = MetadataLoader::new(object_store.clone(), &file_path, file_size);
2873        let parquet_meta = loader.load(&mut cache_metrics).await?;
2874        let parquet_schema = parquet_meta.file_metadata().schema_descr();
2875        assert_eq!(3, parquet_schema.num_columns());
2876
2877        // Ask Parquet to read only the deepest requested JSON2 path. This should select the single
2878        // leaf j.a.x and avoid both sibling leaves j.a.y and j.b.
2879
2880        let projection =
2881            ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths(
2882                vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]],
2883            )]);
2884        let projection_plan =
2885            build_projection_plan(&projection, parquet_schema, batch.schema_ref()).unwrap();
2886        assert_eq!(vec![true], projection_plan.projected_root_presence);
2887        assert_eq!(
2888            projection_plan.mask,
2889            ProjectionMask::leaves(parquet_schema, vec![0])
2890        );
2891
2892        // Read through the low-level stream directly.
2893
2894        let arrow_metadata =
2895            ArrowReaderMetadata::try_new(Arc::new(parquet_meta), ArrowReaderOptions::new())
2896                .map_err(|e| e.to_string())?;
2897        let fetcher = SstParquetRangeFetcher::new(
2898            file_handle.file_id(),
2899            file_path.clone(),
2900            object_store,
2901            CacheStrategy::Disabled,
2902            0,
2903            None,
2904        );
2905        let mut stream = build_sst_parquet_record_batch_stream(
2906            arrow_metadata,
2907            0,
2908            None,
2909            projection_plan.mask,
2910            fetcher,
2911            file_path,
2912            1024,
2913        )?;
2914
2915        let Some(batch) = stream.next().await.transpose()? else {
2916            unreachable!()
2917        };
2918        let expected = r#"
2919+-------------+
2920| j           |
2921+-------------+
2922| {a: {x: 1}} |
2923| {a: {x: 2}} |
2924| {a: {x: 3}} |
2925+-------------+
2926"#;
2927        assert_eq!(batch.pretty_print(), expected.trim());
2928        Ok(())
2929    }
2930
2931    #[tokio::test(flavor = "current_thread")]
2932    async fn test_minmax_predicate_key_not_built_when_index_result_cache_disabled() {
2933        #[derive(Eq, PartialEq, Hash)]
2934        struct PanicDebugUdf;
2935
2936        impl Debug for PanicDebugUdf {
2937            fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
2938                panic!("minmax predicate key should not format exprs when cache is disabled");
2939            }
2940        }
2941
2942        impl ScalarUDFImpl for PanicDebugUdf {
2943            fn name(&self) -> &str {
2944                "panic_debug_udf"
2945            }
2946
2947            fn signature(&self) -> &Signature {
2948                static SIGNATURE: LazyLock<Signature> =
2949                    LazyLock::new(|| Signature::variadic_any(Volatility::Immutable));
2950                &SIGNATURE
2951            }
2952
2953            fn return_type(&self, _arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
2954                Ok(DataType::Int64)
2955            }
2956
2957            fn invoke_with_args(
2958                &self,
2959                _args: ScalarFunctionArgs,
2960            ) -> datafusion_common::Result<ColumnarValue> {
2961                Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(1))))
2962            }
2963        }
2964
2965        let object_store = ObjectStore::new(Memory::default()).unwrap();
2966        let file_handle = sst_file_handle(0, 1);
2967        let table_dir = "test_table".to_string();
2968        let path_type = PathType::Bare;
2969        let file_path = file_handle.file_path(&table_dir, path_type);
2970
2971        let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
2972        let batch = RecordBatch::try_from_iter([("col", col)]).unwrap();
2973        let mut parquet_bytes = Vec::new();
2974        let mut writer = ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), None).unwrap();
2975        writer.write(&batch).unwrap();
2976        writer.close().unwrap();
2977        let file_size = parquet_bytes.len() as u64;
2978        object_store.write(&file_path, parquet_bytes).await.unwrap();
2979
2980        let region_metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
2981        let read_format = FlatReadFormat::new(
2982            region_metadata.clone(),
2983            ReadColumns::new(
2984                region_metadata
2985                    .column_metadatas
2986                    .iter()
2987                    .map(|column| column.column_id),
2988            ),
2989            None,
2990            &file_path,
2991            false,
2992        )
2993        .unwrap();
2994
2995        let mut cache_metrics = MetadataCacheMetrics::default();
2996        let loader = MetadataLoader::new(object_store.clone(), &file_path, file_size);
2997        let parquet_meta = loader.load(&mut cache_metrics).await.unwrap();
2998
2999        let udf = Arc::new(ScalarUDF::new_from_impl(PanicDebugUdf));
3000        let predicate = Predicate::new(vec![Expr::ScalarFunction(ScalarFunction::new_udf(
3001            udf,
3002            vec![],
3003        ))]);
3004        let builder = ParquetReaderBuilder::new(table_dir, path_type, file_handle, object_store)
3005            .predicate(Some(predicate))
3006            .cache(CacheStrategy::Disabled);
3007
3008        let row_group_size = parquet_meta.row_group(0).num_rows() as usize;
3009        let total_row_count = parquet_meta.file_metadata().num_rows() as usize;
3010        let mut metrics = ReaderFilterMetrics::default();
3011        let selection = builder.row_groups_by_minmax(
3012            &read_format,
3013            &parquet_meta,
3014            row_group_size,
3015            total_row_count,
3016            &mut metrics,
3017            false,
3018        );
3019
3020        assert!(!selection.is_empty());
3021    }
3022
3023    #[test]
3024    fn test_expr_is_immutable_checks_scalar_function_volatility() {
3025        #[derive(Debug, PartialEq, Eq, Hash)]
3026        struct TestVolatilityUdf {
3027            name: String,
3028            signature: Signature,
3029        }
3030
3031        impl TestVolatilityUdf {
3032            fn new(name: &str, volatility: Volatility) -> Self {
3033                Self {
3034                    name: name.to_string(),
3035                    signature: Signature::variadic_any(volatility),
3036                }
3037            }
3038        }
3039
3040        impl ScalarUDFImpl for TestVolatilityUdf {
3041            fn name(&self) -> &str {
3042                &self.name
3043            }
3044
3045            fn signature(&self) -> &Signature {
3046                &self.signature
3047            }
3048
3049            fn return_type(&self, _arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
3050                Ok(DataType::Int64)
3051            }
3052
3053            fn invoke_with_args(
3054                &self,
3055                _args: ScalarFunctionArgs,
3056            ) -> datafusion_common::Result<ColumnarValue> {
3057                Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(1))))
3058            }
3059        }
3060
3061        let expr = |name: &str, volatility| {
3062            Expr::ScalarFunction(ScalarFunction::new_udf(
3063                Arc::new(ScalarUDF::new_from_impl(TestVolatilityUdf::new(
3064                    name, volatility,
3065                ))),
3066                vec![],
3067            ))
3068        };
3069
3070        assert!(expr_is_immutable(&expr(
3071            "immutable_udf",
3072            Volatility::Immutable
3073        )));
3074        assert!(!expr_is_immutable(&expr("stable_udf", Volatility::Stable)));
3075        assert!(!expr_is_immutable(&expr(
3076            "volatile_udf",
3077            Volatility::Volatile
3078        )));
3079
3080        let scalar_variable = Expr::ScalarVariable(
3081            Arc::new(Field::new("@@version", DataType::Utf8, false)),
3082            vec!["@@version".to_string()],
3083        );
3084        assert!(!expr_is_immutable(&scalar_variable));
3085    }
3086
3087    #[tokio::test(flavor = "current_thread")]
3088    async fn test_has_row_level_selection() {
3089        let object_store = ObjectStore::new(Memory::default()).unwrap();
3090        let file_path = "row_level_selection.parquet";
3091
3092        let col = Arc::new(Int64Array::from_iter_values([1, 2, 3, 4, 5])) as ArrayRef;
3093        let batch = RecordBatch::try_from_iter([("col", col)]).unwrap();
3094        let props = WriterProperties::builder()
3095            .set_max_row_group_row_count(Some(3))
3096            .build();
3097        let mut parquet_bytes = Vec::new();
3098        let mut writer =
3099            ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap();
3100        writer.write(&batch).unwrap();
3101        writer.close().unwrap();
3102        let file_size = parquet_bytes.len() as u64;
3103        object_store.write(file_path, parquet_bytes).await.unwrap();
3104
3105        let mut cache_metrics = MetadataCacheMetrics::default();
3106        let loader = MetadataLoader::new(object_store, file_path, file_size);
3107        let parquet_meta = loader.load(&mut cache_metrics).await.unwrap();
3108        assert_eq!(2, parquet_meta.num_row_groups());
3109
3110        let full_row_groups = RowGroupSelection::from_full_row_group_ids([0, 1], 3, 5);
3111        assert!(!has_row_level_selection(&full_row_groups, &parquet_meta));
3112
3113        let prefix_selection = RowGroupSelection::from_row_ranges(vec![(0, vec![0..1, 1..2])], 3);
3114        assert!(has_row_level_selection(&prefix_selection, &parquet_meta));
3115
3116        let interior_selection = RowGroupSelection::from_row_ranges(vec![(0, vec![1..2, 2..3])], 3);
3117        assert!(has_row_level_selection(&interior_selection, &parquet_meta));
3118    }
3119
3120    fn expected_metadata_with_reused_tag_name(
3121        old_metadata: &RegionMetadata,
3122    ) -> Arc<RegionMetadata> {
3123        let mut builder = RegionMetadataBuilder::new(old_metadata.region_id);
3124        builder
3125            .push_column_metadata(ColumnMetadata {
3126                column_schema: ColumnSchema::new(
3127                    "tag_0".to_string(),
3128                    ConcreteDataType::string_datatype(),
3129                    true,
3130                ),
3131                semantic_type: SemanticType::Tag,
3132                column_id: 10,
3133            })
3134            .push_column_metadata(ColumnMetadata {
3135                column_schema: ColumnSchema::new(
3136                    "tag_1".to_string(),
3137                    ConcreteDataType::string_datatype(),
3138                    true,
3139                ),
3140                semantic_type: SemanticType::Tag,
3141                column_id: 1,
3142            })
3143            .push_column_metadata(ColumnMetadata {
3144                column_schema: ColumnSchema::new(
3145                    "field_0".to_string(),
3146                    ConcreteDataType::uint64_datatype(),
3147                    true,
3148                ),
3149                semantic_type: SemanticType::Field,
3150                column_id: 2,
3151            })
3152            .push_column_metadata(ColumnMetadata {
3153                column_schema: ColumnSchema::new(
3154                    "ts".to_string(),
3155                    ConcreteDataType::timestamp_millisecond_datatype(),
3156                    false,
3157                ),
3158                semantic_type: SemanticType::Timestamp,
3159                column_id: 3,
3160            })
3161            .primary_key(vec![10, 1]);
3162
3163        Arc::new(builder.build().unwrap())
3164    }
3165
3166    #[test]
3167    fn test_simple_filter_context_uses_default_value_for_mismatched_expected_metadata() {
3168        let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
3169        let expected_metadata = expected_metadata_with_reused_tag_name(metadata.as_ref());
3170        let ctx = SimpleFilterContext::new_opt(
3171            &metadata,
3172            Some(expected_metadata.as_ref()),
3173            &col("tag_0").eq(lit("a")),
3174        )
3175        .unwrap();
3176        assert!(matches!(
3177            ctx.filter(),
3178            MaybeFilter::Matched | MaybeFilter::Pruned
3179        ));
3180    }
3181
3182    #[test]
3183    fn test_nullable_label_filter_uses_missing_sst_column_default() {
3184        let metadata = Arc::new(sst_region_metadata());
3185        let expected = expected_metadata_with_reused_tag_name(&metadata);
3186        let expr = col("tag_0").is_null().or(col("tag_0").not_eq(lit("lo")));
3187        let ctx = SimpleFilterContext::new_opt(&metadata, Some(&expected), &expr).unwrap();
3188        assert!(matches!(ctx.filter(), MaybeFilter::Matched));
3189    }
3190
3191    #[test]
3192    fn test_simple_filter_context_drops_mismatched_field_filter() {
3193        let (sst_metadata, latest_metadata) = mock_metadata();
3194        let ctx = SimpleFilterContext::new_opt(
3195            &sst_metadata,
3196            Some(latest_metadata.as_ref()),
3197            &col("field_0").eq(lit(1_i64)),
3198        );
3199
3200        assert!(ctx.is_none());
3201    }
3202
3203    fn mock_metadata() -> (RegionMetadataRef, RegionMetadataRef) {
3204        let region_id = RegionId::new(1, 1);
3205        let make_tag_0 = || ColumnMetadata {
3206            column_schema: ColumnSchema::new(
3207                "tag_0".to_string(),
3208                ConcreteDataType::string_datatype(),
3209                true,
3210            ),
3211            semantic_type: SemanticType::Tag,
3212            column_id: 0,
3213        };
3214        let make_ts = || ColumnMetadata {
3215            column_schema: ColumnSchema::new(
3216                "ts".to_string(),
3217                ConcreteDataType::timestamp_millisecond_datatype(),
3218                false,
3219            ),
3220            semantic_type: SemanticType::Timestamp,
3221            column_id: 2,
3222        };
3223        let make_field_0 = |data_type| ColumnMetadata {
3224            column_schema: ColumnSchema::new("field_0".to_string(), data_type, true),
3225            semantic_type: SemanticType::Field,
3226            column_id: 1,
3227        };
3228
3229        let mut sst_builder = RegionMetadataBuilder::new(region_id);
3230        sst_builder
3231            .push_column_metadata(make_tag_0())
3232            .push_column_metadata(make_field_0(ConcreteDataType::uint64_datatype()))
3233            .push_column_metadata(make_ts())
3234            .primary_key(vec![0]);
3235        let sst_metadata = Arc::new(sst_builder.build().unwrap());
3236
3237        let mut expected_builder = RegionMetadataBuilder::new(region_id);
3238        expected_builder
3239            .push_column_metadata(make_tag_0())
3240            .push_column_metadata(make_field_0(ConcreteDataType::int64_datatype()))
3241            .push_column_metadata(make_ts())
3242            .primary_key(vec![0]);
3243
3244        let expected_metadata = Arc::new(expected_builder.build().unwrap());
3245
3246        (sst_metadata, expected_metadata)
3247    }
3248
3249    #[test]
3250    fn test_physical_filter_context_skips_renamed_column() {
3251        let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
3252        let expected_metadata = expected_metadata_with_reused_tag_name(metadata.as_ref());
3253        let read_format = FlatReadFormat::new(
3254            metadata.clone(),
3255            ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
3256            None,
3257            "test",
3258            true,
3259        )
3260        .unwrap();
3261
3262        let ctx = PhysicalFilterContext::new_opt(
3263            &metadata,
3264            Some(expected_metadata.as_ref()),
3265            &read_format,
3266            &col("tag_0").in_list(vec![lit("a"), lit("b")], false),
3267        );
3268
3269        assert!(ctx.is_none());
3270    }
3271
3272    #[test]
3273    fn test_physical_filter_context_only_accepts_prefilter_candidates() {
3274        let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
3275        let read_format = FlatReadFormat::new(
3276            metadata.clone(),
3277            ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
3278            None,
3279            "test",
3280            true,
3281        )
3282        .unwrap();
3283
3284        // InList is on the allowlist — should build a context.
3285        let in_list = col("tag_0").in_list(vec![lit("a"), lit("b")], false);
3286        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &in_list).is_some());
3287
3288        // NOT IN uses the same variant with `negated: true` — also accepted.
3289        let not_in = col("tag_0").in_list(vec![lit("a"), lit("b")], true);
3290        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &not_in).is_some());
3291
3292        // IS NULL / IS NOT NULL are accepted.
3293        let is_null = col("tag_0").is_null();
3294        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &is_null).is_some());
3295        let is_not_null = col("tag_0").is_not_null();
3296        assert!(
3297            PhysicalFilterContext::new_opt(&metadata, None, &read_format, &is_not_null).is_some()
3298        );
3299
3300        // BETWEEN is accepted.
3301        let between = col("field_0").between(lit(1_u64), lit(10_u64));
3302        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &between).is_some());
3303
3304        // Binary expr is handled by SimpleFilterEvaluator — rejected here.
3305        let binary = col("tag_0").eq(lit("a"));
3306        assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, &binary).is_none());
3307    }
3308
3309    fn write_test_parquet_with_pk_column(values: &[&[u8]]) -> bytes::Bytes {
3310        use datatypes::arrow::array::{
3311            BinaryArray, TimestampMillisecondArray, UInt8Array, UInt64Array,
3312        };
3313        use datatypes::arrow::datatypes::{Field as ArrowField, Schema as ArrowSchema, TimeUnit};
3314        use store_api::storage::consts::{
3315            OP_TYPE_COLUMN_NAME, PRIMARY_KEY_COLUMN_NAME, SEQUENCE_COLUMN_NAME,
3316        };
3317
3318        let n = values.len();
3319        let schema = Arc::new(ArrowSchema::new(vec![
3320            ArrowField::new(
3321                "ts",
3322                DataType::Timestamp(TimeUnit::Millisecond, None),
3323                false,
3324            ),
3325            ArrowField::new(PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false),
3326            ArrowField::new(SEQUENCE_COLUMN_NAME, DataType::UInt64, false),
3327            ArrowField::new(OP_TYPE_COLUMN_NAME, DataType::UInt8, false),
3328        ]));
3329        let ts: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![0_i64; n]));
3330        let pk: ArrayRef = Arc::new(BinaryArray::from_iter_values(values.iter().copied()));
3331        let seq: ArrayRef = Arc::new(UInt64Array::from(vec![0_u64; n]));
3332        let op: ArrayRef = Arc::new(UInt8Array::from(vec![0_u8; n]));
3333        let batch = RecordBatch::try_new(schema.clone(), vec![ts, pk, seq, op]).unwrap();
3334
3335        let mut bytes = Vec::new();
3336        let mut writer = ArrowWriter::try_new(&mut bytes, schema, None).unwrap();
3337        writer.write(&batch).unwrap();
3338        writer.close().unwrap();
3339        bytes::Bytes::from(bytes)
3340    }
3341
3342    fn load_parquet_meta(bytes: bytes::Bytes) -> Arc<ParquetMetaData> {
3343        let builder =
3344            parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
3345        builder.metadata().clone()
3346    }
3347
3348    #[test]
3349    fn test_should_read_pk_as_binary_small_chunk_returns_false() {
3350        let bytes = write_test_parquet_with_pk_column(&[b"a", b"b", b"c"]);
3351        let meta = load_parquet_meta(bytes);
3352
3353        assert!(!should_read_pk_as_binary_with_limit(&meta, 1024));
3354    }
3355
3356    #[test]
3357    fn test_should_read_pk_as_binary_large_chunk_returns_true() {
3358        let owned: Vec<Vec<u8>> = (0..512u32)
3359            .map(|i| {
3360                let mut v = vec![0u8; 16];
3361                v[..4].copy_from_slice(&i.to_le_bytes());
3362                v
3363            })
3364            .collect();
3365        let refs: Vec<&[u8]> = owned.iter().map(|v| v.as_slice()).collect();
3366        let bytes = write_test_parquet_with_pk_column(&refs);
3367        let meta = load_parquet_meta(bytes);
3368
3369        assert!(should_read_pk_as_binary_with_limit(&meta, 1024));
3370    }
3371}