From 97648525cfabbeff592c9b42c40e1d860bbdb5bd Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 11 Sep 2026 09:35:57 +0000 Subject: [PATCH] perf(mito2): skip proven all-match prefilters (#9066) * perf(mito2): skip proven all-match prefilters Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(perf): add manual all-match prefilter reproduction Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(mito2): cover all-match prefilter execution paths Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(mito2): match prefilter fixture to sparse SST schema Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore(mito2): address all-match prefilter lint findings Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(perf): cover all-match prefilters in default regressions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .github/scripts/query-regression-run.py | 1 + src/mito2/src/sst/parquet/prefilter.rs | 554 +++++++++++++++++- src/mito2/src/sst/parquet/reader.rs | 258 +++++++- tests/perf/README.md | 22 +- .../mito_prefilter_all_match/case.toml | 213 +++++++ .../test_query_regression_case_selection.py | 5 + 6 files changed, 1039 insertions(+), 14 deletions(-) create mode 100644 tests/perf/query_cases/mito_prefilter_all_match/case.toml diff --git a/.github/scripts/query-regression-run.py b/.github/scripts/query-regression-run.py index 268cb0bae6..a99706c04a 100644 --- a/.github/scripts/query-regression-run.py +++ b/.github/scripts/query-regression-run.py @@ -39,6 +39,7 @@ DEFAULT_CASES = [ "tests/perf/query_cases/prom_remote_write_integer_counter/case.toml", "tests/perf/query_cases/promql_range_boundary/case.toml", "tests/perf/query_cases/promql_instant_last_row_9034/case.toml", + "tests/perf/query_cases/mito_prefilter_all_match/case.toml", ] HEAVY_CASES = [ diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index d39f540c18..4362440e41 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -24,15 +24,19 @@ use std::sync::Arc; use api::v1::SemanticType; use common_recordbatch::filter::SimpleFilterEvaluator; +use datafusion_common::ScalarValue; use datafusion_expr::Expr; use datatypes::arrow::array::{Array, BinaryArray, BooleanArray, BooleanBufferBuilder}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::datatypes::SchemaRef; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use datatypes::value::Value; use futures::StreamExt; use mito_codec::row_converter::{PrimaryKeyCodec, PrimaryKeyFilter, build_primary_key_codec}; use parquet::arrow::ProjectionMask; -use parquet::arrow::arrow_reader::RowSelection; +use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +use parquet::file::metadata::ParquetMetaData; use parquet::schema::types::SchemaDescriptor; use smallvec::{SmallVec, smallvec}; use snafu::{OptionExt, ResultExt}; @@ -47,7 +51,7 @@ use crate::error::{ }; use crate::sst::parquet::file_range::PreFilterMode; use crate::sst::parquet::flat_format::FlatReadFormat; -use crate::sst::parquet::format::PrimaryKeyArray; +use crate::sst::parquet::format::{PrimaryKeyArray, StatValues}; use crate::sst::parquet::reader::{ MaybeFilter, PhysicalFilterContext, RowGroupBuildContext, RowGroupReaderBuilder, SimpleFilterContext, @@ -398,6 +402,7 @@ pub(crate) fn build_bulk_filter_plan( /// the prefilter pass. A caller can postpone simple timestamp filters to the normal /// precise-filter path when the scan time range covers the SST. When predicate /// prefiltering is disabled, all simple filters remain on the normal path instead. +#[allow(clippy::too_many_arguments)] pub(crate) fn build_reader_filter_plan( predicate: Option<&Predicate>, expected_metadata: Option<&RegionMetadata>, @@ -406,6 +411,7 @@ pub(crate) fn build_reader_filter_plan( postpone_time_index_filter: bool, read_format: &FlatReadFormat, codec: &Arc, + parquet_metadata: &ParquetMetaData, ) -> ReaderFilterPlan { let Some(predicate) = predicate else { return ReaderFilterPlan { @@ -533,6 +539,7 @@ pub(crate) fn build_reader_filter_plan( prefilter_simple_filters.clone(), prefilter_physical_filters, schema_version, + parquet_metadata, ); if prefilter_builder.is_some() { @@ -567,6 +574,8 @@ pub(crate) struct PrefilterContext { pk_filter_expr_strs: Option>, /// Arrow schema used to build narrowed prefilter projections. arrow_schema: SchemaRef, + /// Simple filters already proven SQL-true by this row group's statistics. + proven_simple_filters: Vec, } /// Pre-built state for constructing [PrefilterContext] per row group. @@ -583,6 +592,8 @@ pub(crate) struct PrefilterContextBuilder { metadata: RegionMetadataRef, schema_version: u64, arrow_schema: SchemaRef, + /// Per-row-group simple filters already proven SQL-true by column statistics. + proven_simple_filters: Vec>, } impl PrefilterContextBuilder { @@ -592,6 +603,7 @@ impl PrefilterContextBuilder { /// - The read format doesn't use flat layout /// - No prefilter columns are selected /// - Prefilter would read the full projection without any PK filter + #[allow(clippy::too_many_arguments)] pub(crate) fn new( read_format: &FlatReadFormat, codec: &Arc, @@ -600,6 +612,7 @@ impl PrefilterContextBuilder { filters: Vec, physical_filters: Vec, schema_version: u64, + parquet_metadata: &ParquetMetaData, ) -> Option { let metadata = read_format.metadata(); let use_raw_tag_columns = read_format.batch_has_raw_pk_columns(); @@ -646,6 +659,9 @@ impl PrefilterContextBuilder { return None; } + let proven_simple_filters = + simple_filter_stats_proofs(read_format, parquet_metadata.row_groups(), &filters); + Some(Self { pk_filters, pk_filter_expr_strs, @@ -655,11 +671,12 @@ impl PrefilterContextBuilder { metadata: metadata.clone(), schema_version, arrow_schema: read_format.arrow_schema().clone(), + proven_simple_filters, }) } /// Builds a [PrefilterContext] for a specific row group. - pub(crate) fn build(&self) -> PrefilterContext { + pub(crate) fn build(&self, row_group_idx: usize) -> PrefilterContext { let pk_filter = self .build_primary_key_filter() .map(|filter| Box::new(filter) as Box); @@ -670,6 +687,11 @@ impl PrefilterContextBuilder { schema_version: self.schema_version, pk_filter_expr_strs: self.pk_filter_expr_strs.clone(), arrow_schema: self.arrow_schema.clone(), + proven_simple_filters: self + .proven_simple_filters + .get(row_group_idx) + .cloned() + .unwrap_or_else(|| vec![false; self.filters.len()]), } } @@ -687,6 +709,114 @@ impl PrefilterContextBuilder { const PREFILTER_COLUMN_RATIO_THRESHOLD: f64 = 0.5; const PREFILTER_MIN_REMAINING_COLUMNS: usize = 2; +/// Returns row-group-major proof bits for simple filters. Statistics are +/// extracted once per eligible filter across all row groups. +fn simple_filter_stats_proofs( + read_format: &FlatReadFormat, + row_groups: &[parquet::file::metadata::RowGroupMetaData], + filters: &[SimpleFilterContext], +) -> Vec> { + let mut proofs = vec![vec![false; filters.len()]; row_groups.len()]; + for (filter_idx, filter_ctx) in filters.iter().enumerate() { + let Some((filter, literal)) = eligible_simple_filter(read_format, filter_ctx) else { + continue; + }; + let (StatValues::Values(mins), StatValues::Values(maxs), StatValues::Values(null_counts)) = ( + read_format.min_values(row_groups, filter_ctx.column_id()), + read_format.max_values(row_groups, filter_ctx.column_id()), + read_format.null_counts(row_groups, filter_ctx.column_id()), + ) else { + continue; + }; + for (row_group_idx, proof) in proofs.iter_mut().enumerate() { + proof[filter_idx] = simple_filter_is_true_by_values( + filter, + &literal, + stat_value_at(&mins, row_group_idx), + stat_value_at(&maxs, row_group_idx), + stat_value_at(&null_counts, row_group_idx), + ); + } + } + proofs +} + +fn eligible_simple_filter<'a>( + read_format: &FlatReadFormat, + filter_ctx: &'a SimpleFilterContext, +) -> Option<(&'a SimpleFilterEvaluator, Value)> { + if filter_ctx.semantic_type() != SemanticType::Field { + return None; + } + let filter = filter_ctx.filter().as_filter()?; + let literal = filter.literal_value()?; + let column = read_format + .metadata() + .column_by_id(filter_ctx.column_id())?; + column_type_matches_literal(&column.column_schema.data_type, &literal) + .then_some((filter, literal)) +} + +fn simple_filter_is_true_by_values( + filter: &SimpleFilterEvaluator, + literal: &Value, + min: Option, + max: Option, + null_count: Option, +) -> bool { + let (Some(min), Some(max), Some(null_count)) = (min, max, null_count) else { + return false; + }; + if null_count != Value::UInt64(0) + || !same_supported_value_type(&min, literal) + || !same_supported_value_type(&max, literal) + || min > max + { + return false; + } + + if filter.is_gt() { + min > *literal + } else if filter.is_gt_eq() { + min >= *literal + } else if filter.is_lt() { + max < *literal + } else if filter.is_lt_eq() { + max <= *literal + } else if filter.is_eq() { + min == *literal && max == *literal + } else if filter.is_not_eq() { + max < *literal || min > *literal + } else { + false + } +} + +fn column_type_matches_literal(data_type: &ConcreteDataType, literal: &Value) -> bool { + matches!( + (data_type, literal), + (ConcreteDataType::Int32(_), Value::Int32(_)) + | (ConcreteDataType::UInt32(_), Value::UInt32(_)) + | (ConcreteDataType::Int64(_), Value::Int64(_)) + | (ConcreteDataType::UInt64(_), Value::UInt64(_)) + ) +} + +fn same_supported_value_type(left: &Value, right: &Value) -> bool { + matches!( + (left, right), + (Value::Int32(_), Value::Int32(_)) + | (Value::UInt32(_), Value::UInt32(_)) + | (Value::Int64(_), Value::Int64(_)) + | (Value::UInt64(_), Value::UInt64(_)) + ) +} + +fn stat_value_at(values: &datatypes::arrow::array::ArrayRef, index: usize) -> Option { + let scalar = ScalarValue::try_from_array(values, index).ok()?; + Value::try_from(scalar).ok() +} + /// Result of prefiltering a row group. pub(crate) struct PrefilterResult { /// Refined row selection after prefiltering. @@ -941,6 +1071,9 @@ async fn execute_prefilter_by_reading_columns( build_ctx: &RowGroupBuildContext<'_>, ) -> Result { let entries = all_prefilter_entries(prefilter_ctx); + if entries.is_empty() { + return Ok(identity_prefilter_result(reader_builder, build_ctx)); + } let (mask, rows_before_filter) = build_prefilter_masks(prefilter_ctx, reader_builder, build_ctx, &entries).await?; @@ -965,6 +1098,13 @@ fn all_prefilter_entries(prefilter_ctx: &PrefilterContext) -> Vec, +) -> PrefilterResult { + let row_count = reader_builder + .parquet_metadata() + .row_group(build_ctx.row_group_idx) + .num_rows() as usize; + PrefilterResult { + refined_selection: identity_row_selection(&build_ctx.row_selection, row_count), + filtered_rows: 0, + } +} + +fn identity_row_selection( + original_selection: &Option, + row_count: usize, +) -> RowSelection { + original_selection + .clone() + .unwrap_or_else(|| RowSelection::from(vec![RowSelector::select(row_count)])) +} + fn rows_before_filter( reader_builder: &RowGroupReaderBuilder, build_ctx: &RowGroupBuildContext<'_>, @@ -1200,15 +1371,24 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + use bytes::Bytes; use common_recordbatch::filter::SimpleFilterEvaluator; use datafusion_common::ScalarValue; use datafusion_expr::{col, lit}; use datatypes::arrow::array::{ - ArrayRef, DictionaryArray, TimestampMillisecondArray, UInt8Array, UInt32Array, UInt64Array, + ArrayRef, DictionaryArray, Int32Array, TimestampMillisecondArray, UInt8Array, UInt32Array, + UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; + use datatypes::arrow::record_batch::RecordBatch; + use datatypes::prelude::ConcreteDataType; + use datatypes::value::Value; use mito_codec::row_converter::{PrimaryKeyFilter, build_primary_key_codec}; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::{ParquetRecordBatchReaderBuilder, RowSelector}; use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::RegionMetadataBuilder; + use store_api::region_request::{AlterKind, ModifyColumnType}; use super::*; use crate::read::read_columns::ReadColumns; @@ -1268,6 +1448,83 @@ mod tests { .collect() } + fn stats_metadata_with_options( + values: &[Vec], + nulls: Option<&[bool]>, + writer_options: Option, + ) -> Arc { + let first = new_record_batch_with_custom_sequence(&["a", "x"], 0, values[0].len(), 1); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, first.schema(), writer_options).unwrap(); + for (idx, values) in values.iter().enumerate() { + let has_null = nulls.and_then(|nulls| nulls.get(idx)).copied() == Some(true); + let batch = new_record_batch_with_custom_sequence( + &["a", "x"], + 0, + if has_null { + values.len().max(2) + } else { + values.len() + }, + 1, + ); + let mut columns = batch.columns().to_vec(); + columns[2] = Arc::new(if has_null { + UInt64Array::from(vec![Some(values[0]), None]) + } else { + UInt64Array::from_iter_values(values.iter().copied()) + }); + writer + .write(&RecordBatch::try_new(batch.schema(), columns).unwrap()) + .unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes)) + .unwrap() + .metadata() + .clone() + } + + fn stats_metadata(values: &[Vec], nulls: Option<&[bool]>) -> Arc { + stats_metadata_with_options(values, nulls, None) + } + + fn int32_stats_metadata(values: &[i32]) -> Arc { + let batch = new_record_batch_with_custom_sequence(&["a", "x"], 0, values.len(), 1); + let mut fields = batch + .schema() + .fields() + .iter() + .map(|field| field.as_ref().clone()) + .collect::>(); + fields[2].set_data_type(DataType::Int32); + let mut columns = batch.columns().to_vec(); + columns[2] = Arc::new(Int32Array::from(values.to_vec())); + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes)) + .unwrap() + .metadata() + .clone() + } + + fn metadata_with_field_type(data_type: ConcreteDataType) -> RegionMetadataRef { + let mut builder = RegionMetadataBuilder::from_existing(sst_region_metadata()); + builder + .alter(AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnType { + column_name: "field_0".to_string(), + target_type: data_type, + }], + }) + .unwrap(); + Arc::new(builder.build().unwrap()) + } + fn new_physical_filter_contexts( metadata: &RegionMetadataRef, read_format: &FlatReadFormat, @@ -1502,10 +1759,292 @@ mod tests { Vec::new(), Vec::new(), metadata.schema_version, + &stats_metadata(&[vec![1]], None), ); assert!(builder.is_none()); } + #[test] + fn test_simple_filter_stats_uses_real_metadata() { + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let read_format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new( + metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + "test", + true, + ) + .unwrap(); + let parquet_metadata = stats_metadata(&[vec![i64::MAX as u64 + 2]], None); + let filters = new_simple_filter_contexts( + &metadata, + &[ + col("field_0").gt(lit(i64::MAX as u64 + 1)), + col("field_0").gt(lit(i64::MAX as u64 + 2)), + lit(i64::MAX as u64 + 1).lt(col("field_0")), + ], + ); + + let proofs = + simple_filter_stats_proofs(&read_format, parquet_metadata.row_groups(), &filters); + assert_eq!(proofs, vec![vec![true, false, true]]); + } + + #[test] + fn test_simple_filter_stats_retain_narrow_or_incomplete_metadata() { + let parquet_metadata = stats_metadata(&[vec![2]], Some(&[true])); + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let read_format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new( + metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + "test", + true, + ) + .unwrap(); + let filter = new_simple_filter_contexts(&metadata, &[col("field_0").gt(lit(1_u64))]); + assert_eq!( + simple_filter_stats_proofs(&read_format, parquet_metadata.row_groups(), &filter), + vec![vec![false]], + ); + let missing_stats = stats_metadata_with_options( + &[vec![2]], + None, + Some( + parquet::file::properties::WriterProperties::builder() + .set_statistics_enabled(parquet::file::properties::EnabledStatistics::None) + .build(), + ), + ); + assert_eq!( + simple_filter_stats_proofs(&read_format, missing_stats.row_groups(), &filter), + vec![vec![false]], + ); + + // Parquet INT32 stats decode as Value::Int32; without the actual SST + // type gate this Int8 field would be incorrectly proven true. + let narrow_metadata = metadata_with_field_type(ConcreteDataType::int8_datatype()); + let narrow_read_format = FlatReadFormat::new( + narrow_metadata.clone(), + ReadColumns::new( + narrow_metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + "test", + true, + ) + .unwrap(); + let narrow_filter = + new_simple_filter_contexts(&narrow_metadata, &[col("field_0").gt(lit(1_i32))]); + let narrow_stats = int32_stats_metadata(&[2]); + assert_eq!( + simple_filter_stats_proofs( + &narrow_read_format, + narrow_stats.row_groups(), + &narrow_filter, + ), + vec![vec![false]], + ); + } + + #[test] + fn test_prefilter_builder_uses_row_group_proofs_by_index() { + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let read_format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new( + metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + "test", + true, + ) + .unwrap(); + let codec = build_primary_key_codec(metadata.as_ref()); + let builder = PrefilterContextBuilder::new( + &read_format, + &codec, + None, + None, + new_simple_filter_contexts(&metadata, &[col("field_0").gt(lit(1_u64))]), + Vec::new(), + metadata.schema_version, + &stats_metadata(&[vec![2], vec![1]], None), + ) + .unwrap(); + + assert!(builder.build(0).proven_simple_filters[0]); + assert!(!builder.build(1).proven_simple_filters[0]); + } + + #[test] + fn test_simple_filter_stats_prove_integer_predicates() { + macro_rules! assert_all_operators { + ($literal:expr, $base:expr, $below:expr, $above:expr) => {{ + for (expr, min, max) in [ + (col("x").gt(lit($literal)), $above.clone(), $above.clone()), + (col("x").gt_eq(lit($literal)), $base.clone(), $above.clone()), + (col("x").lt(lit($literal)), $below.clone(), $below.clone()), + (col("x").lt_eq(lit($literal)), $below.clone(), $base.clone()), + (col("x").eq(lit($literal)), $base.clone(), $base.clone()), + ( + col("x").not_eq(lit($literal)), + $below.clone(), + $below.clone(), + ), + ] { + let filter = SimpleFilterEvaluator::try_new(&expr).unwrap(); + assert!(simple_filter_is_true_by_values( + &filter, + &$base, + Some(min), + Some(max), + Some(Value::UInt64(0)), + )); + } + }}; + } + + assert_all_operators!(5_i32, Value::Int32(5), Value::Int32(-1), Value::Int32(6)); + assert_all_operators!(5_u32, Value::UInt32(5), Value::UInt32(4), Value::UInt32(6)); + assert_all_operators!(5_i64, Value::Int64(5), Value::Int64(-1), Value::Int64(6)); + assert_all_operators!( + i64::MAX as u64 + 2, + Value::UInt64(i64::MAX as u64 + 2), + Value::UInt64(i64::MAX as u64 + 1), + Value::UInt64(i64::MAX as u64 + 3) + ); + + // Boundary cases must remain unproven: the min/max interval still + // admits a row that violates the predicate. + for (expr, min, max) in [ + (col("x").lt(lit(5_i64)), 1, 5), + (col("x").gt(lit(5_i64)), 5, 9), + (col("x").lt_eq(lit(5_i64)), 1, 6), + (col("x").gt_eq(lit(5_i64)), 4, 9), + (col("x").eq(lit(5_i64)), 5, 6), + (col("x").not_eq(lit(5_i64)), 5, 9), + (col("x").not_eq(lit(5_i64)), 1, 5), + (col("x").not_eq(lit(5_i64)), 1, 9), + ] { + let filter = SimpleFilterEvaluator::try_new(&expr).unwrap(); + assert!( + !simple_filter_is_true_by_values( + &filter, + &Value::Int64(5), + Some(Value::Int64(min)), + Some(Value::Int64(max)), + Some(Value::UInt64(0)), + ), + "{expr:?} must not be proven by stats {min}..={max}", + ); + } + } + + #[test] + fn test_simple_filter_stats_retain_unknown_or_unsupported_values() { + let filter = SimpleFilterEvaluator::try_new(&col("x").gt(lit(1_i32))).unwrap(); + let literal = Value::Int32(1); + for (min, max, null_count) in [ + ( + Some(Value::Int32(2)), + Some(Value::Int32(3)), + Some(Value::UInt64(1)), + ), + (Some(Value::Int32(2)), Some(Value::Int32(3)), None), + ( + Some(Value::Null), + Some(Value::Int32(3)), + Some(Value::UInt64(0)), + ), + (None, Some(Value::Int32(3)), Some(Value::UInt64(0))), + (Some(Value::Int32(2)), None, Some(Value::UInt64(0))), + ( + Some(Value::Int64(2)), + Some(Value::Int64(3)), + Some(Value::UInt64(0)), + ), + ( + Some(Value::Int32(3)), + Some(Value::Int32(2)), + Some(Value::UInt64(0)), + ), + ] { + assert!(!simple_filter_is_true_by_values( + &filter, &literal, min, max, null_count + )); + } + + let float_filter = SimpleFilterEvaluator::try_new(&col("x").gt(lit(1.0_f64))).unwrap(); + assert!(!simple_filter_is_true_by_values( + &float_filter, + &Value::Float64(1.0.into()), + Some(Value::Float64(2.0.into())), + Some(Value::Float64(3.0.into())), + Some(Value::UInt64(0)), + )); + } + + #[test] + fn test_prefilter_entries_keep_only_unproven_simple_filter_indices() { + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let filters = new_simple_filter_contexts( + &metadata, + &[col("field_0").gt(lit(1_u64)), col("field_0").lt(lit(9_u64))], + ); + let context = PrefilterContext { + pk_filter: None, + filters, + physical_filters: Vec::new(), + schema_version: metadata.schema_version, + pk_filter_expr_strs: None, + arrow_schema: metadata.schema.arrow_schema().clone(), + proven_simple_filters: vec![true, false], + }; + + let entries = all_prefilter_entries(&context); + assert!(matches!( + entries.as_slice(), + [PrefilterEntry { + kind: PrefilterEntryKind::Simple(1), + .. + }] + )); + } + + #[test] + fn test_identity_row_selection_preserves_input() { + let sparse = RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(3), + RowSelector::skip(1), + ]); + assert_eq!(identity_row_selection(&Some(sparse.clone()), 6), sparse); + + let empty = RowSelection::from(vec![]); + assert_eq!(identity_row_selection(&Some(empty.clone()), 6), empty); + assert_eq!( + identity_row_selection(&None, 6), + RowSelection::from(vec![RowSelector::select(6)]) + ); + } + #[test] fn test_should_use_prefilter() { assert!(should_use_prefilter(1, 5, 6)); @@ -1602,6 +2141,7 @@ mod tests { false, &full_read_format, &codec, + &stats_metadata(&[vec![1]], None), ); assert!(skip_fields_plan.prefilter_builder.is_some()); assert_eq!( @@ -1621,6 +2161,7 @@ mod tests { true, &full_read_format, &codec, + &stats_metadata(&[vec![1]], None), ); assert!(postponed_time_plan.prefilter_builder.is_some()); assert_eq!( @@ -1638,6 +2179,7 @@ mod tests { true, &full_read_format, &codec, + &stats_metadata(&[vec![1]], None), ); assert!(postponed_time_only_plan.prefilter_builder.is_none()); assert_eq!( @@ -1667,6 +2209,7 @@ mod tests { false, &projected_read_format, &metric_codec, + &stats_metadata(&[vec![1]], None), ); assert!(pk_prefilter_plan.prefilter_builder.is_some()); assert!( @@ -1691,6 +2234,7 @@ mod tests { true, &projected_read_format, &metric_codec, + &stats_metadata(&[vec![1]], None), ); assert!(disabled_plan.prefilter_builder.is_none()); assert_eq!( @@ -1724,6 +2268,7 @@ mod tests { false, &read_format, &codec, + &stats_metadata(&[vec![1]], None), ); let plan_b_a = build_reader_filter_plan( Some(&Predicate::new(vec![expr_b, expr_a])), @@ -1733,6 +2278,7 @@ mod tests { false, &read_format, &codec, + &stats_metadata(&[vec![1]], None), ); let exprs_ab = plan_ab.prefilter_builder.unwrap().pk_filter_expr_strs; diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index ff58d920b1..b8588bd36c 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -598,6 +598,7 @@ impl ParquetReaderBuilder { self.postpone_time_index_filter, &read_format, &codec, + &parquet_meta, ); if self.defer_optional_page_index @@ -1975,7 +1976,10 @@ impl RowGroupReaderBuilder { &self, build_ctx: RowGroupBuildContext<'_>, ) -> Result { - let prefilter_ctx = self.prefilter_builder.as_ref().map(|b| b.build()); + let prefilter_ctx = self + .prefilter_builder + .as_ref() + .map(|b| b.build(build_ctx.row_group_idx)); let Some(mut prefilter_ctx) = prefilter_ctx else { // No prefilter applicable, build stream with full projection. @@ -2553,9 +2557,11 @@ impl FlatRowGroupReader { #[cfg(test)] mod tests { use std::any::Any; + use std::collections::HashMap; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, LazyLock}; + use api::v1::OpType; use common_error::ext::WhateverResult; use common_function::scalars::json::json_get::JsonGetWithType; use common_function::scalars::udf::create_udf; @@ -2567,25 +2573,271 @@ mod tests { ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, col, lit, }; - use datatypes::arrow::array::{ArrayRef, Int64Array, StringArray, StructArray}; - use datatypes::arrow::datatypes::{Fields, Schema}; + use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, Int64Array, StringArray, StructArray, + TimestampMillisecondArray, UInt8Array, UInt64Array, + }; + use datatypes::arrow::datatypes::{Fields, Schema, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::Json2ExtensionType; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use object_store::services::Memory; use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::RowSelector; use parquet::file::properties::WriterProperties; + use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_request::PathType; use store_api::storage::RegionId; use table::predicate::Predicate; use super::*; + use crate::cache::CacheManager; use crate::sst::parquet::metadata::MetadataLoader; + use crate::sst::parquet::prefilter::{build_reader_filter_plan, execute_prefilter}; use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns}; + use crate::sst::parquet::row_group::ParquetFetchMetrics; use crate::test_util::sst_util::{sst_file_handle, sst_region_metadata}; + async fn prefilter_test_builder( + object_store: ObjectStore, + predicate: Predicate, + cache_strategy: CacheStrategy, + ) -> (RowGroupReaderBuilder, Arc) { + let metadata = Arc::new( + crate::test_util::sst_util::sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + ), + ); + let batch = |start: i64, end: i64| { + let mut primary_key = BinaryDictionaryBuilder::::new(); + let mut fields = Vec::new(); + let mut timestamps = Vec::new(); + for value in start..end { + let tag = if value == 4 { "b" } else { "a" }; + primary_key + .append(crate::test_util::sst_util::new_sparse_primary_key( + &[tag, "x"], + &metadata, + 1, + 100, + )) + .unwrap(); + fields.push(value as u64); + timestamps.push(value); + } + RecordBatch::try_new( + crate::sst::to_flat_sst_arrow_schema( + &metadata, + &crate::sst::FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ), + vec![ + Arc::new(UInt64Array::from(fields)) as ArrayRef, + Arc::new(TimestampMillisecondArray::from(timestamps)) as ArrayRef, + Arc::new(primary_key.finish()) as ArrayRef, + Arc::new(UInt64Array::from_value(1, (end - start) as usize)) as ArrayRef, + Arc::new(UInt8Array::from_value( + OpType::Put as u8, + (end - start) as usize, + )) as ArrayRef, + ], + ) + .unwrap() + }; + let first_batch = batch(0, 3); + let second_batch = batch(3, 6); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, first_batch.schema(), None).unwrap(); + writer.write(&first_batch).unwrap(); + writer.flush().unwrap(); + writer.write(&second_batch).unwrap(); + writer.close().unwrap(); + + let file_handle = sst_file_handle(0, 6); + let file_path = file_handle.file_path("prefilter_test", PathType::Bare); + let file_size = bytes.len() as u64; + object_store.write(&file_path, bytes).await.unwrap(); + + let mut cache_metrics = MetadataCacheMetrics::default(); + let parquet_meta = Arc::new( + MetadataLoader::new(object_store.clone(), &file_path, file_size) + .load(&mut cache_metrics) + .await + .unwrap(), + ); + let read_format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new( + metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + &file_path, + false, + ) + .unwrap(); + let codec = build_primary_key_codec(metadata.as_ref()); + let filter_plan = build_reader_filter_plan( + Some(&predicate), + None, + PreFilterMode::All, + true, + false, + &read_format, + &codec, + &parquet_meta, + ); + assert!(filter_plan.prefilter_builder.is_some()); + + let output_schema = read_format.arrow_schema().clone(); + let parquet_schema = parquet_meta.file_metadata().schema_descr(); + let projection = build_projection_plan(read_format.parquet_read_columns(), parquet_schema); + let arrow_metadata = + ArrowReaderMetadata::try_new(parquet_meta.clone(), ArrowReaderOptions::new()).unwrap(); + ( + RowGroupReaderBuilder { + file_handle: file_handle.clone(), + file_path, + parquet_meta, + parquet_metadata_size: 0, + arrow_metadata, + output_schema, + json2_rewrite_targets: HashMap::new(), + object_store, + projection, + has_nested_projection: false, + cache_strategy, + prefilter_builder: filter_plan.prefilter_builder, + batch_size: DEFAULT_READ_BATCH_SIZE, + }, + metadata, + ) + } + + #[tokio::test(flavor = "current_thread")] + async fn test_execute_prefilter_proven_filters_preserve_selection_without_fetching() { + let object_store = ObjectStore::new(Memory::default()).unwrap(); + let predicate = Predicate::new(vec![col("field_0").gt_eq(lit(0_u64))]); + let (reader_builder, _) = + prefilter_test_builder(object_store, predicate, CacheStrategy::Disabled).await; + let prefilter_builder = reader_builder.prefilter_builder.as_ref().unwrap(); + + for original_selection in [ + None, + Some(RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(1), + ])), + Some(RowSelection::from(vec![])), + ] { + let mut prefilter_ctx = prefilter_builder.build(0); + let fetch_metrics = ParquetFetchMetrics::default(); + let result = execute_prefilter( + &mut prefilter_ctx, + &reader_builder, + &RowGroupBuildContext { + row_group_idx: 0, + row_selection: original_selection.clone(), + fetch_metrics: Some(&fetch_metrics), + }, + ) + .await + .unwrap(); + + let expected = original_selection.unwrap_or_else(|| { + RowSelection::from(vec![RowSelector::select( + reader_builder.parquet_meta.row_group(0).num_rows() as usize, + )]) + }); + assert_eq!(result.refined_selection, expected); + assert_eq!(result.filtered_rows, 0); + let metrics = fetch_metrics.data.lock().unwrap(); + assert_eq!(metrics.pages_to_fetch_store, 0); + assert_eq!(metrics.pages_to_fetch_mem, 0); + assert_eq!(metrics.pages_to_fetch_write_cache, 0); + } + } + + #[tokio::test(flavor = "current_thread")] + async fn test_execute_prefilter_mixed_filters_use_nonzero_row_group_and_cache() { + let object_store = ObjectStore::new(Memory::default()).unwrap(); + let predicate = Predicate::new(vec![ + col("field_0").gt_eq(lit(3_u64)), + col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(6), None))), + col("field_0").in_list(vec![lit(3_u64), lit(4_u64)], false), + col("tag_0").eq(lit("a")), + ]); + let cache = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .prefilter_result_cache_size(1024) + .build(), + )); + let (reader_builder, _) = + prefilter_test_builder(object_store, predicate, cache.clone()).await; + let prefilter_builder = reader_builder.prefilter_builder.as_ref().unwrap(); + + for pass in 0..2 { + let mut prefilter_ctx = prefilter_builder.build(1); + let fetch_metrics = ParquetFetchMetrics::default(); + let fetch_metrics_ref = (pass == 1).then_some(&fetch_metrics); + let result = execute_prefilter( + &mut prefilter_ctx, + &reader_builder, + &RowGroupBuildContext { + row_group_idx: 1, + row_selection: Some(RowSelection::from(vec![ + RowSelector::select(2), + RowSelector::skip(1), + ])), + fetch_metrics: fetch_metrics_ref, + }, + ) + .await + .unwrap(); + assert_eq!(result.filtered_rows, 1); + assert_eq!( + result.refined_selection, + RowSelection::from(vec![RowSelector::select(1), RowSelector::skip(2),]) + ); + if pass == 1 { + assert_eq!(fetch_metrics.data.lock().unwrap().pages_to_fetch_store, 0); + } + } + + let disabled = prefilter_test_builder( + ObjectStore::new(Memory::default()).unwrap(), + Predicate::new(vec![ + col("field_0").gt_eq(lit(3_u64)), + col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(6), None))), + col("field_0").in_list(vec![lit(3_u64), lit(4_u64)], false), + col("tag_0").eq(lit("a")), + ]), + CacheStrategy::Disabled, + ) + .await; + let mut prefilter_ctx = disabled.0.prefilter_builder.as_ref().unwrap().build(1); + let result = execute_prefilter( + &mut prefilter_ctx, + &disabled.0, + &RowGroupBuildContext { + row_group_idx: 1, + row_selection: None, + fetch_metrics: None, + }, + ) + .await + .unwrap(); + assert_eq!(result.filtered_rows, 2); + assert_eq!( + result.refined_selection, + RowSelection::from(vec![RowSelector::select(1), RowSelector::skip(2)]) + ); + } + #[test] fn test_skip_prefilter_for_json_get() -> WhateverResult<()> { fn json_get_expr(base: Expr, path: &str) -> Expr { diff --git a/tests/perf/README.md b/tests/perf/README.md index d518fec996..6f7893ae66 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -204,9 +204,14 @@ case for issue #7913. It writes 8192 series × 20160 samples through remote-writ in 1440-sample daily time chunks, flushing after each chunk before running 1d/7d/14d TQL selectors. It is not included in the default `all` case set because ingestion cost dominates routine CI validation. Commenting `/query-regression heavy` runs -only this case; `/query-regression` runs the seven routine default cases. Manual +only this case; `/query-regression` runs the eight routine default cases. Manual workflow dispatch accepts the `heavy` token to select this case. +The routine default set also includes +`tests/perf/query_cases/mito_prefilter_all_match/case.toml`. It covers the +Mito prefilter all-match path, a mixed integer/float filter, and integer and +float selective controls over a 262144-row direct SST fixture. + ## OTLP trace load scenario `scenario.kind = "otlp_trace_load"` runs a bounded native `otelgen` process @@ -366,6 +371,8 @@ uv run --no-project python .github/scripts/query-regression-run.py \ --work-dir /tmp/query-regression-work ``` +For a focused manual reproduction of the Mito prefilter all-match optimization, use the existing lifecycle command above with `--cases tests/perf/query_cases/mito_prefilter_all_match/case.toml`. Its four count probes require explicit result inspection rather than automatic validation: expect `262144`, `131072`, `16896`, and `16896` in query order. The default lifecycle retains caches, so the optimization remains exercised but its timing includes warm-cache effects; use an explicitly configured cold environment when a cold comparison is required. + The Rust runner subcommands are also useful for focused diagnostics: ```bash @@ -393,12 +400,13 @@ parquetbench/scanbench` as the read-bench tool against each target's data direct The workflow runs when an allowlisted repository admin comments `/query-regression` on a non-draft PR. It does not rerun on pushes, -ready-for-review, or reopen events. `/query-regression` runs the seven routine -default cases, including `promql_instant_last_row_9034`; -`/query-regression heavy` runs only the high-cardinality remote-write #7913 case. -PR runs build base/candidate once and use `--allow-large-fixture`. Manual -`workflow_dispatch` runs can pass `all`, `heavy`, one case path, or a -comma/whitespace-separated list of case paths, and can override refs. +ready-for-review, or reopen events. `/query-regression` runs the eight routine +default cases, including `promql_instant_last_row_9034` and +`mito_prefilter_all_match`; `/query-regression heavy` runs only the +high-cardinality remote-write #7913 case. PR runs build base/candidate once and +use `--allow-large-fixture`. Manual `workflow_dispatch` runs can pass `all`, +`heavy`, one case path, or a comma/whitespace-separated list of case paths, and +can override refs. Comment admission is two workflows. `slash-command-dispatch.yml` uses [peter-evans/slash-command-dispatch](https://github.com/peter-evans/slash-command-dispatch) diff --git a/tests/perf/query_cases/mito_prefilter_all_match/case.toml b/tests/perf/query_cases/mito_prefilter_all_match/case.toml new file mode 100644 index 0000000000..9db69eb446 --- /dev/null +++ b/tests/perf/query_cases/mito_prefilter_all_match/case.toml @@ -0,0 +1,213 @@ +# Performance regression coverage for the Mito prefilter all-match optimization. +# +# The count probes are deliberately not automatic assertions: inspect their +# results manually. Each should return 262144, 131072, 16896, and 16896 +# respectively. + +[case] +name = "mito_prefilter_all_match" +description = "Mito prefilter all-match regression coverage with integer and float selective controls" + +[scenario] +kind = "direct_readable_sst" +seed = 9402 + +[[scenario.tables]] +database = "public" +name = "mito_prefilter_all_match" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["tag_a", "tag_b"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "tag_a" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 1, prefix = "tag_a" } + +[[scenario.tables.columns]] +name = "tag_b" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 1, prefix = "tag_b" } + +[[scenario.tables.columns]] +name = "field_01" +type = "UINT64" +semantic = "field" + +[[scenario.tables.columns]] +name = "field_02" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 2.0 } + +[[scenario.tables.columns]] +name = "field_03" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 3.0 } + +[[scenario.tables.columns]] +name = "field_04" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 4.0 } + +[[scenario.tables.columns]] +name = "field_05" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 5.0 } + +[[scenario.tables.columns]] +name = "field_06" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 6.0 } + +[[scenario.tables.columns]] +name = "field_07" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 7.0 } + +[[scenario.tables.columns]] +name = "field_08" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 8.0 } + +[[scenario.tables.columns]] +name = "field_09" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 9.0 } + +[[scenario.tables.columns]] +name = "field_10" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 10.0 } + +[[scenario.tables.columns]] +name = "field_11" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 11.0 } + +[[scenario.tables.columns]] +name = "field_12" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 12.0 } + +[[scenario.tables.columns]] +name = "field_13" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 13.0 } + +[[scenario.tables.columns]] +name = "field_14" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 14.0 } + +[[scenario.tables.columns]] +name = "field_15" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 15.0 } + +[[scenario.tables.columns]] +name = "field_16" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 16.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +[scenario.layout] +regions = 1 +sst_count = 1 +rows_per_sst = 262144 +row_group_size = 262144 +series_count = 1 +start_unix_nanos = 1704067200000000000 +step_nanos = 1000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "timestamp_major" + +[[scenario.queries]] +name = "all_match_field_01" +kind = "sql" +query = "EXPLAIN ANALYZE VERBOSE SELECT * FROM mito_prefilter_all_match WHERE field_01 >= 0" +warmup = 2 +iterations = 10 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 30 + +[[scenario.queries]] +name = "all_match_field_01_and_field_02" +kind = "sql" +query = "EXPLAIN ANALYZE VERBOSE SELECT * FROM mito_prefilter_all_match WHERE field_01 >= 0 AND field_02 > 1.98" +warmup = 2 +iterations = 10 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 30 + +[[scenario.queries]] +name = "selective_control_field_01" +kind = "sql" +query = "EXPLAIN ANALYZE VERBOSE SELECT * FROM mito_prefilter_all_match WHERE field_01 >= 131072" +warmup = 2 +iterations = 10 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 30 + +[[scenario.queries]] +name = "selective_control_field_02" +kind = "sql" +query = "EXPLAIN ANALYZE VERBOSE SELECT * FROM mito_prefilter_all_match WHERE field_02 > 1.98" +warmup = 2 +iterations = 10 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 30 + +[[scenario.queries]] +name = "count_all_match_field_01" +kind = "sql" +query = "SELECT count(*) FROM mito_prefilter_all_match WHERE field_01 >= 0" +warmup = 0 +iterations = 1 + +[[scenario.queries]] +name = "count_selective_control_field_01" +kind = "sql" +query = "SELECT count(*) FROM mito_prefilter_all_match WHERE field_01 >= 131072" +warmup = 0 +iterations = 1 + +[[scenario.queries]] +name = "count_all_match_field_01_and_field_02" +kind = "sql" +query = "SELECT count(*) FROM mito_prefilter_all_match WHERE field_01 >= 0 AND field_02 > 1.98" +warmup = 0 +iterations = 1 + +[[scenario.queries]] +name = "count_selective_control_field_02" +kind = "sql" +query = "SELECT count(*) FROM mito_prefilter_all_match WHERE field_02 > 1.98" +warmup = 0 +iterations = 1 diff --git a/tests/perf/test_query_regression_case_selection.py b/tests/perf/test_query_regression_case_selection.py index 84b3cd44fa..6ccf18a442 100644 --- a/tests/perf/test_query_regression_case_selection.py +++ b/tests/perf/test_query_regression_case_selection.py @@ -41,9 +41,14 @@ class QueryRegressionCaseSelectionTest(unittest.TestCase): "tests/perf/query_cases/prom_remote_write_integer_counter/case.toml", "tests/perf/query_cases/promql_range_boundary/case.toml", "tests/perf/query_cases/promql_instant_last_row_9034/case.toml", + "tests/perf/query_cases/mito_prefilter_all_match/case.toml", ], ) + def test_implicit_selection_has_eight_routine_cases(self) -> None: + self.assertEqual(len(runner.DEFAULT_CASES), 8) + self.assertEqual(runner.split_cases([]), runner.DEFAULT_CASES) + def test_heavy_selects_only_remote_write_7913(self) -> None: self.assertEqual( runner.split_cases(["heavy"]),