test(mito2): cover regex inverted index pruning

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

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::<usize>()
+        );
+        assert_eq!(
+            100,
+            baseline_batches
+                .iter()
+                .map(RecordBatch::num_rows)
+                .sum::<usize>()
+        );
+    }
+
     #[tokio::test]
     async fn test_write_flat_read_with_bloom_filter() {
         let mut env = TestEnv::new().await;