diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 8a05c099e9..4a57f2a6ee 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -1897,6 +1897,190 @@ mod tests { assert_eq!(metrics.filter_metrics.rows_inverted_filtered, 50); } + #[tokio::test] + async fn test_write_flat_read_with_inverted_index_regex_match() { + // Scope: existing FST/inverted-index RegexMatch candidate filtering only; + // neither an REI n-gram index nor log planner support is implemented here. + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let file_path = RegionFilePathFactory::new(FILE_DIR.to_string(), PathType::Bare); + let metadata = Arc::new(sst_region_metadata()); + let row_group_size = 100; + + // Four physical row groups. RG 0 and RG 2 alternate matching and + // nonmatching values, while RG 1 and RG 3 contain only nonmatching + // values. The index segments intentionally contain ten rows so the + // candidate selection for RG 0 and RG 2 is coarse. + let mut flat_batches = Vec::with_capacity(4); + let mut expected_batches = Vec::with_capacity(2); + for row_group in 0..4 { + let rows: Vec<_> = (0..100) + .map(|row| { + let tag_0 = if row_group % 2 == 0 && row % 2 == 0 { + "api-match" + } else { + "zzz-nonmatch" + }; + (tag_0, "service", (row_group * 100 + row) as i64) + }) + .collect(); + flat_batches.push(new_record_batch_from_rows(&rows)); + + if row_group % 2 == 0 { + let matching_rows: Vec<_> = rows + .iter() + .copied() + .filter(|(tag_0, _, _)| tag_0.starts_with("api-")) + .collect(); + expected_batches.push(new_record_batch_from_rows(&matching_rows)); + } + } + + let flat_source = new_flat_source_from_record_batches(flat_batches); + let write_opts = WriteOptions { + row_group_size, + ..Default::default() + }; + + let puffin_manager = env + .get_puffin_manager() + .build(object_store.clone(), file_path.clone()); + let indexer_builder = IndexerBuilderImpl { + build_type: IndexBuildType::Flush, + metadata: metadata.clone(), + puffin_manager, + write_cache_enabled: false, + intermediate_manager: env.get_intermediate_manager(), + index_options: IndexOptions { + inverted_index: InvertedIndexOptions { + segment_row_count: 10, + ..Default::default() + }, + }, + inverted_index_config: Default::default(), + fulltext_index_config: Default::default(), + bloom_filter_index_config: Default::default(), + #[cfg(feature = "vector_index")] + vector_index_config: Default::default(), + }; + + let info = write_flat_sst( + object_store.clone(), + metadata.clone(), + indexer_builder, + file_path.clone(), + flat_source, + &write_opts, + ) + .await; + assert_eq!(400, info.num_rows); + assert_eq!(4, info.num_row_groups); + let handle = create_file_handle_from_sst_info(&info, &metadata); + + let preds = vec![Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("tag_0"))), + op: Operator::RegexMatch, + right: Box::new(lit("^api-.*")), + })]; + let predicate = Predicate::new(preds.clone()); + let inverted_index_applier = InvertedIndexApplierBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + object_store.clone(), + &metadata, + HashSet::from_iter([0]), + env.get_puffin_manager(), + ) + .build(&preds) + .unwrap() + .map(Arc::new); + + let indexed_builder = ParquetReaderBuilder::new( + FILE_DIR.to_string(), + PathType::Bare, + handle.clone(), + object_store.clone(), + ) + .predicate(Some(predicate.clone())) + .inverted_index_appliers([inverted_index_applier, None]); + let mut indexed_metrics = ReaderMetrics::default(); + let (indexed_context, indexed_selection) = indexed_builder + .build_reader_input(&mut indexed_metrics) + .await + .unwrap() + .unwrap(); + + // The ten-row index segments make the two selected row groups coarse: + // each contains all 100 rows, including 50 regex false positives. + assert_eq!(2, indexed_selection.row_group_count()); + assert_eq!(100, indexed_selection.get(0).unwrap().row_count()); + assert_eq!(100, indexed_selection.get(2).unwrap().row_count()); + assert_eq!(4, indexed_metrics.filter_metrics.rg_total); + assert_eq!(2, indexed_metrics.filter_metrics.rg_inverted_filtered); + assert_eq!(200, indexed_metrics.filter_metrics.rows_inverted_filtered); + + let mut indexed_reader = ParquetReader::new(Arc::new(indexed_context), indexed_selection) + .await + .unwrap(); + let mut indexed_batches = Vec::new(); + while let Some(batch) = indexed_reader.next_record_batch().await.unwrap() { + indexed_batches.push(batch); + } + + // Disable the applier while keeping the same residual predicate. This + // scans all row groups and provides the correctness baseline. + let baseline_builder = + ParquetReaderBuilder::new(FILE_DIR.to_string(), PathType::Bare, handle, object_store) + .predicate(Some(predicate)); + let mut baseline_metrics = ReaderMetrics::default(); + let (baseline_context, baseline_selection) = baseline_builder + .build_reader_input(&mut baseline_metrics) + .await + .unwrap() + .unwrap(); + assert_eq!(4, baseline_selection.row_group_count()); + assert_eq!(4, baseline_metrics.filter_metrics.rg_total); + assert_eq!(0, baseline_metrics.filter_metrics.rg_inverted_filtered); + assert_eq!(0, baseline_metrics.filter_metrics.rows_inverted_filtered); + + let mut baseline_reader = + ParquetReader::new(Arc::new(baseline_context), baseline_selection) + .await + .unwrap(); + let mut baseline_batches = Vec::new(); + while let Some(batch) = baseline_reader.next_record_batch().await.unwrap() { + baseline_batches.push(batch); + } + + let expected = pretty_format_batches(&expected_batches) + .unwrap() + .to_string(); + assert_eq!( + expected, + pretty_format_batches(&indexed_batches).unwrap().to_string() + ); + assert_eq!( + expected, + pretty_format_batches(&baseline_batches) + .unwrap() + .to_string() + ); + assert_eq!( + 100, + indexed_batches + .iter() + .map(RecordBatch::num_rows) + .sum::() + ); + assert_eq!( + 100, + baseline_batches + .iter() + .map(RecordBatch::num_rows) + .sum::() + ); + } + #[tokio::test] async fn test_write_flat_read_with_bloom_filter() { let mut env = TestEnv::new().await;