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