1use 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
101const MAX_ROW_GROUPS_TO_CHECK_PK: usize = 4;
103
104fn 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
187pub struct ParquetReaderBuilder {
189 series_index: Option<SeriesIndexReadContext>,
191 table_dir: String,
193 path_type: PathType,
195 file_handle: FileHandle,
196 object_store: ObjectStore,
197 predicate: Option<Predicate>,
199 read_cols: Option<ReadColumns>,
204 json2_rewrite_targets: Json2RewriteTargets,
206 cache_strategy: CacheStrategy,
208 inverted_index_appliers: [Option<InvertedIndexApplierRef>; 2],
210 bloom_filter_index_appliers: [Option<BloomFilterIndexApplierRef>; 2],
211 fulltext_index_appliers: [Option<FulltextIndexApplierRef>; 2],
212 expected_metadata: Option<RegionMetadataRef>,
216 compaction: bool,
218 pre_filter_mode: PreFilterMode,
220 enable_predicate_prefilter: bool,
222 postpone_time_index_filter: bool,
224 decode_primary_key_values: bool,
226 page_index_policy: PageIndexPolicy,
227 defer_optional_page_index: bool,
228 batch_size: usize,
230 memory_tracker: Option<QueryMemoryTracker>,
232}
233
234impl ParquetReaderBuilder {
235 pub(crate) fn memory_tracker(mut self, tracker: Option<QueryMemoryTracker>) -> Self {
237 self.memory_tracker = tracker;
238 self
239 }
240
241 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 #[must_use]
276 pub(crate) fn series_index(mut self, context: Option<SeriesIndexReadContext>) -> Self {
277 self.series_index = context;
278 self
279 }
280
281 #[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 #[must_use]
290 pub fn predicate(mut self, predicate: Option<Predicate>) -> ParquetReaderBuilder {
291 self.predicate = predicate;
292 self
293 }
294
295 #[must_use]
299 pub fn projection(mut self, read_cols: Option<ReadColumns>) -> ParquetReaderBuilder {
300 self.read_cols = read_cols;
301 self
302 }
303
304 #[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 #[must_use]
313 pub fn cache(mut self, cache: CacheStrategy) -> ParquetReaderBuilder {
314 self.cache_strategy = cache;
315 self
316 }
317
318 #[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 #[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 #[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 #[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 #[must_use]
357 pub fn compaction(mut self, compaction: bool) -> Self {
358 self.compaction = compaction;
359 self
360 }
361
362 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 let skip_auto_convert = self.compaction && is_same_region_partition;
483
484 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(®ion_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(®ion_meta);
505 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 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 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 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 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 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(®ion_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 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 let mut referenced_columns = HashSet::new();
763 region_partition_expr.collect_column_names(&mut referenced_columns);
764
765 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 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 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 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 #[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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 !required_row_groups.contains_non_empty_row_group(*rg_id)
1460 || cached_row_groups.contains_row_group(*rg_id)
1462 })
1463}
1464
1465#[derive(Debug, Default, Clone)]
1467pub(crate) struct ReaderFilterMetrics {
1468 pub(crate) rg_total: usize,
1470 pub(crate) rg_fulltext_filtered: usize,
1472 pub(crate) rg_inverted_filtered: usize,
1474 pub(crate) rg_minmax_filtered: usize,
1476 pub(crate) rg_bloom_filtered: usize,
1478
1479 pub(crate) rows_total: usize,
1481 pub(crate) rows_fulltext_filtered: usize,
1483 pub(crate) rows_inverted_filtered: usize,
1485 pub(crate) rows_bloom_filtered: usize,
1487 pub(crate) rows_precise_filtered: usize,
1489
1490 pub(crate) fulltext_index_cache_hit: usize,
1492 pub(crate) fulltext_index_cache_miss: usize,
1494 pub(crate) inverted_index_cache_hit: usize,
1496 pub(crate) inverted_index_cache_miss: usize,
1498 pub(crate) bloom_filter_cache_hit: usize,
1500 pub(crate) bloom_filter_cache_miss: usize,
1502 pub(crate) minmax_cache_hit: usize,
1504 pub(crate) minmax_cache_miss: usize,
1506
1507 pub(crate) inverted_index_apply_metrics: Option<InvertedIndexApplyMetrics>,
1509 pub(crate) bloom_filter_apply_metrics: Option<BloomFilterIndexApplyMetrics>,
1511 pub(crate) fulltext_index_apply_metrics: Option<FulltextIndexApplyMetrics>,
1513
1514 pub(crate) pruner_cache_hit: usize,
1516 pub(crate) pruner_cache_miss: usize,
1518 pub(crate) pruner_prune_cost: Duration,
1520 pub(crate) files_time_range_pruned: usize,
1522}
1523
1524impl ReaderFilterMetrics {
1525 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 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 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#[derive(Default, Clone, Copy)]
1627pub struct MetadataCacheMetrics {
1628 pub mem_cache_hit: usize,
1630 pub file_cache_hit: usize,
1632 pub cache_miss: usize,
1634 pub metadata_load_cost: Duration,
1636 pub num_reads: usize,
1638 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 pub(crate) fn is_empty(&self) -> bool {
1683 self.metadata_load_cost.is_zero()
1684 }
1685
1686 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#[derive(Debug, Default, Clone)]
1699pub struct ReaderMetrics {
1700 pub(crate) filter_metrics: ReaderFilterMetrics,
1702 pub(crate) build_cost: Duration,
1704 pub(crate) scan_cost: Duration,
1706 pub(crate) num_record_batches: usize,
1708 pub(crate) num_batches: usize,
1710 pub(crate) num_rows: usize,
1712 pub(crate) metadata_cache_metrics: MetadataCacheMetrics,
1714 pub(crate) fetch_metrics: Option<Arc<ParquetFetchMetrics>>,
1716 pub(crate) metadata_mem_size: isize,
1718 pub(crate) num_range_builders: isize,
1720}
1721
1722impl ReaderMetrics {
1723 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 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
1752pub(crate) struct RowGroupReaderBuilder {
1754 file_handle: FileHandle,
1758 file_path: String,
1760 parquet_meta: Arc<ParquetMetaData>,
1762 parquet_metadata_size: usize,
1764 arrow_metadata: ArrowReaderMetadata,
1766 output_schema: SchemaRef,
1768 json2_rewrite_targets: HashMap<String, Json2TargetLayout>,
1770 object_store: ObjectStore,
1772 projection: ProjectionMaskPlan,
1774 has_nested_projection: bool,
1776 cache_strategy: CacheStrategy,
1778 prefilter_builder: Option<PrefilterContextBuilder>,
1780 batch_size: usize,
1782 memory_tracker: Option<QueryMemoryTracker>,
1784}
1785
1786pub(crate) struct RowGroupBuildContext<'a> {
1789 pub(crate) row_group_idx: usize,
1791 pub(crate) row_selection: Option<RowSelection>,
1793 pub(crate) fetch_metrics: Option<&'a ParquetFetchMetrics>,
1795}
1796
1797impl RowGroupReaderBuilder {
1798 pub(crate) fn file_path(&self) -> &str {
1800 &self.file_path
1801 }
1802
1803 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 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 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 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 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 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)]
1988pub(crate) enum MaybeFilter {
1990 Filter(SimpleFilterEvaluator),
1992 Matched,
1994 Pruned,
1996}
1997
1998impl MaybeFilter {
1999 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)]
2009pub(crate) struct SimpleFilterContext {
2011 filter: MaybeFilter,
2013 expr_str: String,
2015 column_id: ColumnId,
2017 semantic_type: SemanticType,
2019}
2020
2021impl SimpleFilterContext {
2022 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 let column = meta.column_by_name(filter.column_name())?;
2042 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 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 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 pub(crate) fn filter(&self) -> &MaybeFilter {
2101 &self.filter
2102 }
2103
2104 pub(crate) fn expr_str(&self) -> &str {
2106 &self.expr_str
2107 }
2108
2109 pub(crate) fn column_id(&self) -> ColumnId {
2111 self.column_id
2112 }
2113
2114 pub(crate) fn semantic_type(&self) -> SemanticType {
2116 self.semantic_type
2117 }
2118}
2119
2120#[derive(Clone)]
2122pub(crate) struct PhysicalFilterContext {
2123 filter: Arc<dyn PhysicalExpr>,
2125 expr_str: String,
2127 column_id: ColumnId,
2129 column_name: String,
2131 semantic_type: SemanticType,
2133 schema: SchemaRef,
2135 immutable: bool,
2137}
2138
2139impl PhysicalFilterContext {
2140 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 if sst_column.column_schema.name != column_name {
2161 return None;
2162 }
2163 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 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 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 !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 pub(crate) fn filter(&self) -> &Arc<dyn PhysicalExpr> {
2232 &self.filter
2233 }
2234
2235 pub(crate) fn expr_str(&self) -> &str {
2237 &self.expr_str
2238 }
2239
2240 pub(crate) fn column_id(&self) -> ColumnId {
2242 self.column_id
2243 }
2244
2245 pub(crate) fn column_name(&self) -> &str {
2247 &self.column_name
2248 }
2249
2250 pub(crate) fn semantic_type(&self) -> SemanticType {
2252 self.semantic_type
2253 }
2254
2255 pub(crate) fn schema(&self) -> &SchemaRef {
2257 &self.schema
2258 }
2259
2260 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
2284fn 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
2295pub struct ParquetReader {
2297 context: FileRangeContextRef,
2299 selection: RowGroupSelection,
2301 reader: Option<FlatPruneReader>,
2303 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 #[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 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
2396pub(crate) struct FlatRowGroupReader {
2398 context: FileRangeContextRef,
2400 stream: ProjectedRecordBatchStream,
2402 override_sequence: Option<ArrayRef>,
2404}
2405
2406impl FlatRowGroupReader {
2407 pub(crate) fn new(context: FileRangeContextRef, stream: ProjectedRecordBatchStream) -> Self {
2409 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 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 #[test]
2781 fn test_physical_filter_dropped_on_file_type_mismatch() {
2782 let file_metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
2783 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 assert!(
2807 PhysicalFilterContext::new_opt(&file_metadata, Some(&file_metadata), &format, &between)
2808 .is_some()
2809 );
2810 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 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 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 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 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 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 let not_in = col("tag_0").in_list(vec![lit("a"), lit("b")], true);
3290 assert!(PhysicalFilterContext::new_opt(&metadata, None, &read_format, ¬_in).is_some());
3291
3292 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 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 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}