From eaeb0bfbc118bc0701cbd6b72c4b28372f252c1a Mon Sep 17 00:00:00 2001 From: evenyag Date: Thu, 2 Jul 2026 19:35:22 +0800 Subject: [PATCH] feat: collect bulk memtable column stats Signed-off-by: evenyag --- src/mito2/src/memtable/bulk.rs | 31 +- src/mito2/src/memtable/bulk/part.rs | 368 ++++++++------------- src/mito2/src/memtable/bulk/stats.rs | 471 +++++++++++++++++++++++++++ 3 files changed, 639 insertions(+), 231 deletions(-) create mode 100644 src/mito2/src/memtable/bulk/stats.rs diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index bba9cea5bc..9d39518927 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -20,6 +20,7 @@ pub(crate) mod json_align; pub mod part; pub mod part_reader; mod row_group_reader; +mod stats; use std::collections::{BTreeMap, HashSet}; use std::sync::atomic::{AtomicI64, AtomicU64, AtomicUsize, Ordering}; @@ -53,6 +54,7 @@ use crate::memtable::bulk::part::{ should_prune_bulk_part, }; use crate::memtable::bulk::part_reader::BulkPartBatchIter; +use crate::memtable::bulk::stats::BatchStats; use crate::memtable::stats::WriteMetrics; use crate::memtable::{ AllocTracker, BoxedBatchIterator, BoxedRecordBatchIterator, EncodedBulkPart, EncodedRange, @@ -447,20 +449,26 @@ impl Memtable for BulkMemtable { if bulk_parts.should_compact_unordered_part() && let Some(bulk_part) = bulk_parts.unordered_part.to_bulk_part()? { + let batch_stats = + BatchStats::compute(std::slice::from_ref(&bulk_part.batch), &self.metadata); bulk_parts.parts.push(BulkPartWrapper { part: PartToMerge::Bulk { part: bulk_part, file_id: FileId::random(), + batch_stats, }, merging: false, }); bulk_parts.unordered_part.clear(); } } else { + let batch_stats = + BatchStats::compute(std::slice::from_ref(&fragment.batch), &self.metadata); bulk_parts.parts.push(BulkPartWrapper { part: PartToMerge::Bulk { part: fragment, file_id: FileId::random(), + batch_stats, }, merging: false, }); @@ -507,12 +515,17 @@ impl Memtable for BulkMemtable { if !bulk_parts.unordered_part.is_empty() && let Some(unordered_bulk_part) = bulk_parts.unordered_part.to_bulk_part()? { + let batch_stats = BatchStats::compute( + std::slice::from_ref(&unordered_bulk_part.batch), + &self.metadata, + ); let part_stats = unordered_bulk_part.to_memtable_stats(&self.metadata); let range = MemtableRange::new( Arc::new(MemtableRangeContext::new( self.id, Box::new(BulkRangeIterBuilder { part: unordered_bulk_part, + batch_stats, context: context.clone(), sequence, }), @@ -533,8 +546,11 @@ impl Memtable for BulkMemtable { let part_stats = part_wrapper.part.to_memtable_stats(&self.metadata); let iter_builder: Box = match &part_wrapper.part { - PartToMerge::Bulk { part, .. } => Box::new(BulkRangeIterBuilder { + PartToMerge::Bulk { + part, batch_stats, .. + } => Box::new(BulkRangeIterBuilder { part: part.clone(), + batch_stats: batch_stats.clone(), context: context.clone(), sequence, }), @@ -771,6 +787,7 @@ impl BulkMemtable { /// Iterator builder for bulk range pub struct BulkRangeIterBuilder { pub part: BulkPart, + pub(crate) batch_stats: BatchStats, pub context: Arc, pub sequence: Option, } @@ -799,8 +816,7 @@ impl IterBuilder for BulkRangeIterBuilder { _time_range: Option<(Timestamp, Timestamp)>, metrics: Option, ) -> Result { - let metadata = self.context.read_format().metadata(); - if should_prune_bulk_part(&self.part.batch, &self.context, metadata) { + if should_prune_bulk_part(&self.batch_stats, &self.context) { return Ok(Box::new(std::iter::empty())); } @@ -927,7 +943,11 @@ impl BulkPartWrapper { #[derive(Clone)] enum PartToMerge { /// Raw bulk part. - Bulk { part: BulkPart, file_id: FileId }, + Bulk { + part: BulkPart, + file_id: FileId, + batch_stats: BatchStats, + }, /// Multiple bulk parts. Multi { part: MultiBulkPart, @@ -2222,10 +2242,13 @@ mod tests { /// Helper to create a BulkPartWrapper from a BulkPart. fn create_bulk_part_wrapper(part: BulkPart) -> BulkPartWrapper { + let metadata = metadata_for_test(); + let batch_stats = BatchStats::compute(std::slice::from_ref(&part.batch), &metadata); BulkPartWrapper { part: PartToMerge::Bulk { part, file_id: FileId::random(), + batch_stats, }, merging: false, } diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 92cf2666df..1ef69d1fd1 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -25,13 +25,8 @@ use bytes::Bytes; use common_grpc::flight::{FlightDecoder, FlightEncoder, FlightMessage}; use common_recordbatch::DfRecordBatch as RecordBatch; use common_time::Timestamp; -use datafusion_common::Column; -use datafusion_common::pruning::PruningStatistics; -use datafusion_expr::utils::expr_to_columns; use datatypes::arrow; -use datatypes::arrow::array::{ - Array, ArrayRef, BinaryArray, BooleanArray, StringDictionaryBuilder, UInt8Array, UInt64Array, -}; +use datatypes::arrow::array::{Array, ArrayRef, StringDictionaryBuilder, UInt8Array, UInt64Array}; use datatypes::arrow::compute::{SortColumn, SortOptions, concat_batches}; use datatypes::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef, UInt32Type, @@ -42,7 +37,7 @@ use datatypes::prelude::{MutableVector, Vector}; use datatypes::value::ValueRef; use datatypes::vectors::Helper; use mito_codec::key_values::{KeyValue, KeyValues}; -use mito_codec::row_converter::{PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields}; +use mito_codec::row_converter::PrimaryKeyCodec; use parquet::arrow::ArrowWriter; use parquet::basic::{Compression, ZstdLevel}; use parquet::file::metadata::ParquetMetaData; @@ -52,7 +47,7 @@ use snafu::{OptionExt, ResultExt}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; -use store_api::storage::{ColumnId, FileId, SequenceNumber, SequenceRange}; +use store_api::storage::{FileId, SequenceNumber, SequenceRange}; use crate::error::{ self, ColumnNotFoundSnafu, ComputeArrowSnafu, CreateDefaultSnafu, DataTypeMismatchSnafu, @@ -62,6 +57,7 @@ use crate::error::{ use crate::memtable::bulk::context::{BulkIterContext, BulkIterContextRef}; use crate::memtable::bulk::json_align::Json2Aligner; use crate::memtable::bulk::part_reader::EncodedBulkPartIter; +use crate::memtable::bulk::stats::{BatchPruningStats, BatchStats}; use crate::memtable::time_series::{ValueBuilder, Values}; use crate::memtable::{BoxedRecordBatchIterator, MemScanMetrics, MemtableStats}; use crate::sst::SeriesEstimator; @@ -1265,190 +1261,16 @@ impl BulkPartEncoder { } } -/// Per-batch min/max statistics for the first tag column in a `MultiBulkPart`. -/// -/// Since batches are sorted by primary key, we can extract the min/max of the first tag -/// from the first/last row's encoded primary key in each batch. These statistics enable -/// batch-level pruning using predicates, analogous to row-group pruning in parquet. -#[derive(Debug, Clone)] -struct BatchStats { - /// Number of batches. - num_batches: usize, - /// Column id of the first tag. - first_tag_id: ColumnId, - /// Min values of the first tag, one element per batch. - min_values: ArrayRef, - /// Max values of the first tag, one element per batch. - max_values: ArrayRef, -} - -impl BatchStats { - /// Computes batch statistics from a slice of record batches. - /// - /// Returns `None` if there is no primary key (no first tag to collect stats for) - /// or if extracting statistics fails. - fn compute(batches: &[RecordBatch], metadata: &RegionMetadata) -> Option { - // `primary_key.first()` is correct for both dense and sparse encodings. - // For dense, values follow the order of `metadata.primary_key`. - // For sparse, `decode_leftmost` decodes the first value which also - // corresponds to `primary_key.first()`. See `SparsePrimaryKeyCodec` for format details. - let first_tag_id = *metadata.primary_key.first()?; - let first_tag_column = metadata.column_by_id(first_tag_id)?; - let data_type = &first_tag_column.column_schema.data_type; - - let converter = build_primary_key_codec_with_fields( - metadata.primary_key_encoding, - [(first_tag_id, SortField::new(data_type.clone()))].into_iter(), - ); - let pk_index = primary_key_column_index(batches.first()?.num_columns()); - - let mut min_builder = data_type.create_mutable_vector(batches.len()); - let mut max_builder = data_type.create_mutable_vector(batches.len()); - - for batch in batches { - match Self::extract_first_tag_bounds(batch, pk_index, &*converter) { - Some((min_val, max_val)) => { - min_builder.push_value_ref(&min_val.as_value_ref()); - max_builder.push_value_ref(&max_val.as_value_ref()); - } - None => { - min_builder.push_null(); - max_builder.push_null(); - } - } - } - - Some(Self { - num_batches: batches.len(), - first_tag_id, - min_values: min_builder.to_vector().to_arrow_array(), - max_values: max_builder.to_vector().to_arrow_array(), - }) - } - - /// Extracts the first tag value from the first and last rows of a batch. - fn extract_first_tag_bounds( - batch: &RecordBatch, - pk_index: usize, - converter: &dyn PrimaryKeyCodec, - ) -> Option<(datatypes::value::Value, datatypes::value::Value)> { - if batch.num_rows() == 0 { - return None; - } - - let pk_dict = batch - .column(pk_index) - .as_any() - .downcast_ref::()?; - let pk_values = pk_dict.values().as_any().downcast_ref::()?; - - let keys = pk_dict.keys(); - let min_key = keys.value(0); - let max_key = keys.value(batch.num_rows() - 1); - let min_bytes = pk_values.value(min_key as usize); - let max_bytes = pk_values.value(max_key as usize); - - Some(( - converter.decode_leftmost(min_bytes).ok()??, - converter.decode_leftmost(max_bytes).ok()??, - )) - } -} - -/// Adapter implementing `PruningStatistics` for `BatchStats`. -/// -/// Used with `Predicate::prune_with_stats()` to skip batches whose first-tag -/// min/max range does not match the query predicate. -struct BatchPruningStats<'a> { - stats: &'a BatchStats, - metadata: &'a RegionMetadataRef, -} - -impl PruningStatistics for BatchPruningStats<'_> { - fn min_values(&self, column: &Column) -> Option { - let col = self.metadata.column_by_name(&column.name)?; - if col.column_id == self.stats.first_tag_id { - Some(self.stats.min_values.clone()) - } else { - None - } - } - - fn max_values(&self, column: &Column) -> Option { - let col = self.metadata.column_by_name(&column.name)?; - if col.column_id == self.stats.first_tag_id { - Some(self.stats.max_values.clone()) - } else { - None - } - } - - fn num_containers(&self) -> usize { - self.stats.num_batches - } - - fn null_counts(&self, _column: &Column) -> Option { - None - } - - fn row_counts(&self, _column: &Column) -> Option { - None - } - - fn contained( - &self, - _column: &Column, - _values: &std::collections::HashSet, - ) -> Option { - None - } -} - -/// Returns true if the predicate references the given column name. -fn predicate_references_column(predicate: &table::predicate::Predicate, column_name: &str) -> bool { - let mut columns = HashSet::new(); - for expr in predicate.exprs() { - let _ = expr_to_columns(expr, &mut columns); - } - columns.iter().any(|col| col.name == column_name) -} - -/// Returns true if the batch should be pruned (skipped) based on the first-tag min/max +/// Returns true if the batch should be pruned (skipped) based on min/max /// statistics and the predicate in the context. Returns false if no pruning is possible -/// (no primary key, no predicate, or the batch matches the predicate). -pub(crate) fn should_prune_bulk_part( - batch: &RecordBatch, - context: &BulkIterContext, - metadata: &RegionMetadata, -) -> bool { +/// (no predicate, no stats, or the batch matches the predicate). +pub(crate) fn should_prune_bulk_part(stats: &BatchStats, context: &BulkIterContext) -> bool { let predicate = match &context.predicate { Some(p) => p, None => return false, }; - // Check if the predicate references the first tag column to avoid computing - // expensive batch statistics when they won't help with pruning. - let first_tag_id = match metadata.primary_key.first() { - Some(id) => *id, - None => return false, - }; - // Safety: `first_tag_id` comes from `metadata.primary_key` so the column always exists. - let first_tag_name = &metadata - .column_by_id(first_tag_id) - .unwrap() - .column_schema - .name; - if !predicate_references_column(predicate, first_tag_name) { - return false; - } - let stats = match BatchStats::compute(std::slice::from_ref(batch), metadata) { - Some(s) => s, - None => return false, - }; let region_meta = context.read_format().metadata(); - let pruning_stats = BatchPruningStats { - stats: &stats, - metadata: region_meta, - }; + let pruning_stats = BatchPruningStats::new(stats, region_meta); let mask = predicate.prune_with_stats(&pruning_stats, region_meta.schema.arrow_schema()); !mask.first().copied().unwrap_or(true) } @@ -1472,9 +1294,8 @@ pub struct MultiBulkPart { max_sequence: SequenceNumber, /// Number of series. series_count: usize, - /// Pre-computed per-batch statistics for the first tag column. - /// `None` if there is no primary key. - batch_stats: Option, + /// Pre-computed per-batch column statistics. + batch_stats: BatchStats, } impl MultiBulkPart { @@ -1574,9 +1395,8 @@ impl MultiBulkPart { /// Reads data from this part with the given context and filters. /// - /// If batch-level statistics are available and a predicate is set, prunes - /// batches whose first-tag min/max range doesn't match the predicate before - /// creating the iterator. + /// If a predicate is set, prunes batches whose min/max column stats don't match + /// the predicate before creating the iterator. pub(crate) fn read( &self, context: BulkIterContextRef, @@ -1603,17 +1423,12 @@ impl MultiBulkPart { Ok(Some(Box::new(iter) as BoxedRecordBatchIterator)) } - /// Prunes batches using the first-tag min/max statistics and the predicate. - /// Returns all batches if no stats or no predicate is available. + /// Prunes batches using min/max statistics and the predicate. + /// Returns all batches if no predicate is available. fn prune_batches(&self, context: &BulkIterContextRef) -> Vec { - if let Some(stats) = &self.batch_stats - && let Some(predicate) = &context.predicate - { + if let Some(predicate) = &context.predicate { let region_meta = context.read_format().metadata(); - let pruning_stats = BatchPruningStats { - stats, - metadata: region_meta, - }; + let pruning_stats = BatchPruningStats::new(&self.batch_stats, region_meta); let mask = predicate.prune_with_stats(&pruning_stats, region_meta.schema.arrow_schema()); self.batches @@ -1813,6 +1628,20 @@ mod tests { assert_eq!(expected_rows, total_rows_read); } + fn read_multi_rows( + multi: &MultiBulkPart, + metadata: RegionMetadataRef, + predicate: Predicate, + ) -> usize { + let context = + Arc::new(BulkIterContext::new(metadata, None, Some(predicate), false).unwrap()); + multi + .read(context, None, None) + .unwrap() + .map(|reader| reader.map(|r| r.unwrap().num_rows()).sum()) + .unwrap_or(0) + } + #[test] fn test_prune_row_groups() { let part = prepare(vec![ @@ -2775,27 +2604,27 @@ mod tests { #[test] fn test_multi_bulk_part_prune_batches() { - // Three batches with distinct k0 ranges: ["a"], ["m"], ["z"]. + // Three batches with distinct column ranges. let (multi, metadata) = build_multi_bulk_part(&[ &[MutationInput { k0: "a", - k1: 0, + k1: 10, timestamps: &[1, 2], - v1: &[Some(1.0), Some(2.0)], + v1: &[Some(10.0), Some(11.0)], sequence: 0, }], &[MutationInput { k0: "m", - k1: 0, - timestamps: &[3, 4], - v1: &[Some(3.0), Some(4.0)], + k1: 20, + timestamps: &[101, 102], + v1: &[Some(20.0), Some(21.0)], sequence: 1, }], &[MutationInput { k0: "z", - k1: 0, - timestamps: &[5, 6], - v1: &[Some(5.0), Some(6.0)], + k1: 30, + timestamps: &[201, 202], + v1: &[Some(30.0), Some(31.0)], sequence: 2, }], ]); @@ -2803,31 +2632,61 @@ mod tests { assert_eq!(multi.num_batches(), 3); // k0 = "m" => only middle batch (2 rows). - let context = Arc::new( - BulkIterContext::new( + assert_eq!( + read_multi_rows( + &multi, metadata.clone(), - None, - Some(Predicate::new(vec![ - datafusion_expr::col("k0").eq(datafusion_expr::lit("m")), - ])), - false, - ) - .unwrap(), + Predicate::new(vec![ + datafusion_expr::col("k0").eq(datafusion_expr::lit("m")) + ]) + ), + 2 ); - let reader = multi - .read(context, None, None) - .unwrap() - .expect("should have results"); - let total_rows: usize = reader.map(|r| r.unwrap().num_rows()).sum(); - assert_eq!(total_rows, 2); - // k0 = "nonexistent" => all pruned, returns None. + // k1 = 20 => only middle batch (2 rows). This is not the first tag. + assert_eq!( + read_multi_rows( + &multi, + metadata.clone(), + Predicate::new(vec![ + datafusion_expr::col("k1").eq(datafusion_expr::lit(20u32)) + ]) + ), + 2 + ); + + // ts range intersects only the last batch. + assert_eq!( + read_multi_rows( + &multi, + metadata.clone(), + Predicate::new(vec![datafusion_expr::col("ts").gt(datafusion_expr::lit( + ScalarValue::TimestampMillisecond(Some(200), None) + ))]) + ), + 2 + ); + + // v1 = 20.0 => only one row in the middle batch; min/max should prune + // the first and last batches before row filtering applies. + assert_eq!( + read_multi_rows( + &multi, + metadata.clone(), + Predicate::new(vec![ + datafusion_expr::col("v1").eq(datafusion_expr::lit(20.0f64)) + ]) + ), + 1 + ); + + // k1 = 999 => all pruned, returns None. let context = Arc::new( BulkIterContext::new( metadata.clone(), None, Some(Predicate::new(vec![ - datafusion_expr::col("k0").eq(datafusion_expr::lit("nonexistent")), + datafusion_expr::col("k1").eq(datafusion_expr::lit(999u32)), ])), false, ) @@ -2844,4 +2703,59 @@ mod tests { let total_rows: usize = reader.map(|r| r.unwrap().num_rows()).sum(); assert_eq!(total_rows, 6); } + + #[test] + fn test_bulk_part_minmax_prune_non_first_tag_and_field() { + let metadata = metadata_for_test(); + let part = build_converted_bulk_part(&[MutationInput { + k0: "a", + k1: 10, + timestamps: &[1, 2], + v1: &[Some(10.0), Some(11.0)], + sequence: 0, + }]); + let context = BulkIterContext::new( + metadata.clone(), + None, + Some(Predicate::new(vec![ + datafusion_expr::col("k1").eq(datafusion_expr::lit(999u32)), + ])), + false, + ) + .unwrap(); + let batch_stats = BatchStats::compute(std::slice::from_ref(&part.batch), &metadata); + assert!(should_prune_bulk_part(&batch_stats, &context)); + + let context = BulkIterContext::new( + metadata.clone(), + None, + Some(Predicate::new(vec![ + datafusion_expr::col("v1").gt(datafusion_expr::lit(100.0f64)), + ])), + false, + ) + .unwrap(); + assert!(should_prune_bulk_part(&batch_stats, &context)); + } + + #[test] + fn test_bulk_part_minmax_all_null_field_keeps_batch() { + let metadata = metadata_for_test(); + let part = build_converted_bulk_part(&[MutationInput { + k0: "a", + k1: 10, + timestamps: &[1, 2], + v1: &[None, None], + sequence: 0, + }]); + let context = BulkIterContext::new( + metadata.clone(), + None, + Some(Predicate::new(vec![datafusion_expr::col("v1").is_null()])), + false, + ) + .unwrap(); + let batch_stats = BatchStats::compute(std::slice::from_ref(&part.batch), &metadata); + assert!(!should_prune_bulk_part(&batch_stats, &context)); + } } diff --git a/src/mito2/src/memtable/bulk/stats.rs b/src/mito2/src/memtable/bulk/stats.rs new file mode 100644 index 0000000000..e481c4198d --- /dev/null +++ b/src/mito2/src/memtable/bulk/stats.rs @@ -0,0 +1,471 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Batch-level statistics for bulk memtable pruning. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use datafusion_common::pruning::PruningStatistics; +use datafusion_common::{Column, ScalarValue}; +use datatypes::arrow; +use datatypes::arrow::array::{ + Array, ArrayRef, AsArray, BinaryViewArray, BooleanArray, FixedSizeBinaryArray, + GenericBinaryArray, GenericStringArray, PrimitiveArray, StringViewArray, +}; +use datatypes::arrow::compute::kernels::aggregate; +use datatypes::arrow::datatypes::{ + ArrowPrimitiveType, ArrowTimestampType, DataType as ArrowDataType, Date32Type, Date64Type, + Decimal128Type, Decimal256Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, + Int64Type, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, + Time64NanosecondType, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, + i256, +}; +use datatypes::data_type::{ConcreteDataType, DataType}; +use snafu::ResultExt; +use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef}; +use store_api::storage::ColumnId; + +use crate::error::{ComputeArrowSnafu, Result}; + +type ScalarPair = (ScalarValue, ScalarValue); + +/// Per-batch min/max statistics for columns in a [`MultiBulkPart`](crate::memtable::bulk::part::MultiBulkPart). +#[derive(Debug, Clone)] +pub(crate) struct BatchStats { + /// Number of batches. + num_batches: usize, + /// Column stats by region column id. + columns: HashMap, +} + +#[derive(Debug, Clone)] +struct ColumnStats { + min_values: ArrayRef, + max_values: ArrayRef, +} + +impl BatchStats { + /// Computes batch statistics from a slice of record batches. + pub(crate) fn compute( + batches: &[common_recordbatch::DfRecordBatch], + metadata: &RegionMetadata, + ) -> Self { + let mut columns = HashMap::with_capacity(metadata.column_metadatas.len()); + + for column in &metadata.column_metadatas { + let Some(stats) = compute_column_stats(batches, column) else { + continue; + }; + columns.insert(column.column_id, stats); + } + + Self { + num_batches: batches.len(), + columns, + } + } + + fn min_values(&self, column_id: ColumnId) -> Option { + self.columns + .get(&column_id) + .map(|stats| stats.min_values.clone()) + } + + fn max_values(&self, column_id: ColumnId) -> Option { + self.columns + .get(&column_id) + .map(|stats| stats.max_values.clone()) + } +} + +fn compute_column_stats( + batches: &[common_recordbatch::DfRecordBatch], + column: &ColumnMetadata, +) -> Option { + if matches!( + &column.column_schema.data_type, + ConcreteDataType::Json(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) + | ConcreteDataType::Vector(_) + ) { + return None; + } + + let arrow_type = column.column_schema.data_type.as_arrow_type(); + if !is_supported_arrow_type(&arrow_type) { + return None; + } + + let column_idx = batches + .first()? + .schema() + .index_of(&column.column_schema.name) + .ok()?; + let null_scalar = ScalarValue::try_new_null(&arrow_type).ok()?; + let mut mins = Vec::with_capacity(batches.len()); + let mut maxes = Vec::with_capacity(batches.len()); + + for batch in batches { + let Some((min, max)) = min_max_scalar(batch.column(column_idx), &arrow_type).ok()? else { + mins.push(null_scalar.clone()); + maxes.push(null_scalar.clone()); + continue; + }; + + mins.push(min); + maxes.push(max); + } + + Some(ColumnStats { + min_values: ScalarValue::iter_to_array(mins).ok()?, + max_values: ScalarValue::iter_to_array(maxes).ok()?, + }) +} + +fn is_supported_arrow_type(arrow_type: &ArrowDataType) -> bool { + match arrow_type { + ArrowDataType::Boolean + | ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + | ArrowDataType::Utf8 + | ArrowDataType::LargeUtf8 + | ArrowDataType::Utf8View + | ArrowDataType::Binary + | ArrowDataType::LargeBinary + | ArrowDataType::BinaryView + | ArrowDataType::FixedSizeBinary(_) + | ArrowDataType::Date32 + | ArrowDataType::Date64 + | ArrowDataType::Timestamp(_, _) + | ArrowDataType::Decimal128(_, _) + | ArrowDataType::Decimal256(_, _) => true, + ArrowDataType::Time32(unit) => matches!( + unit, + arrow::datatypes::TimeUnit::Second | arrow::datatypes::TimeUnit::Millisecond + ), + ArrowDataType::Time64(unit) => matches!( + unit, + arrow::datatypes::TimeUnit::Microsecond | arrow::datatypes::TimeUnit::Nanosecond + ), + ArrowDataType::Dictionary(_, value_type) => is_supported_arrow_type(value_type), + _ => false, + } +} + +/// Returns exact min/max scalars for a supported Arrow array. +/// +/// Empty/all-null arrays return `Ok(None)`. Unsupported arrays return `Ok(None)`. +fn min_max_scalar(array: &ArrayRef, logical_type: &ArrowDataType) -> Result> { + if array.is_empty() || array.null_count() == array.len() { + return Ok(None); + } + + if let ArrowDataType::Dictionary(_, value_type) = array.data_type() { + let decoded = + arrow::compute::cast(array.as_ref(), value_type).context(ComputeArrowSnafu)?; + return min_max_scalar(&decoded, value_type); + } + + let stats = match logical_type { + ArrowDataType::Boolean => boolean_min_max(array.as_boolean()), + ArrowDataType::Int8 => primitive_min_max::(array.as_primitive::()), + ArrowDataType::Int16 => primitive_min_max::(array.as_primitive::()), + ArrowDataType::Int32 => primitive_min_max::(array.as_primitive::()), + ArrowDataType::Int64 => primitive_min_max::(array.as_primitive::()), + ArrowDataType::UInt8 => primitive_min_max::(array.as_primitive::()), + ArrowDataType::UInt16 => { + primitive_min_max::(array.as_primitive::()) + } + ArrowDataType::UInt32 => { + primitive_min_max::(array.as_primitive::()) + } + ArrowDataType::UInt64 => { + primitive_min_max::(array.as_primitive::()) + } + ArrowDataType::Float32 => { + primitive_min_max::(array.as_primitive::()) + } + ArrowDataType::Float64 => { + primitive_min_max::(array.as_primitive::()) + } + ArrowDataType::Utf8 => string_min_max(array.as_string::()), + ArrowDataType::LargeUtf8 => string_min_max(array.as_string::()), + ArrowDataType::Utf8View => string_view_min_max(array.as_string_view()), + ArrowDataType::Binary => binary_min_max(array.as_binary::()), + ArrowDataType::LargeBinary => binary_min_max(array.as_binary::()), + ArrowDataType::BinaryView => binary_view_min_max(array.as_binary_view()), + ArrowDataType::FixedSizeBinary(_) => { + fixed_size_binary_min_max(array.as_fixed_size_binary()) + } + ArrowDataType::Date32 => { + primitive_min_max::(array.as_primitive::()) + } + ArrowDataType::Date64 => { + primitive_min_max::(array.as_primitive::()) + } + ArrowDataType::Timestamp(unit, timezone) => match unit { + arrow::datatypes::TimeUnit::Second => timestamp_min_max::( + array.as_primitive::(), + timezone.clone(), + ), + arrow::datatypes::TimeUnit::Millisecond => { + timestamp_min_max::( + array.as_primitive::(), + timezone.clone(), + ) + } + arrow::datatypes::TimeUnit::Microsecond => { + timestamp_min_max::( + array.as_primitive::(), + timezone.clone(), + ) + } + arrow::datatypes::TimeUnit::Nanosecond => timestamp_min_max::( + array.as_primitive::(), + timezone.clone(), + ), + }, + ArrowDataType::Time32(unit) => match unit { + arrow::datatypes::TimeUnit::Second => { + primitive_min_max::(array.as_primitive::()) + } + arrow::datatypes::TimeUnit::Millisecond => primitive_min_max::( + array.as_primitive::(), + ), + _ => None, + }, + ArrowDataType::Time64(unit) => match unit { + arrow::datatypes::TimeUnit::Microsecond => primitive_min_max::( + array.as_primitive::(), + ), + arrow::datatypes::TimeUnit::Nanosecond => primitive_min_max::( + array.as_primitive::(), + ), + _ => None, + }, + ArrowDataType::Decimal128(precision, scale) => { + decimal128_min_max(array.as_primitive::(), *precision, *scale) + } + ArrowDataType::Decimal256(precision, scale) => { + decimal256_min_max(array.as_primitive::(), *precision, *scale) + } + _ => None, + }; + + Ok(stats) +} + +trait ScalarValueFromPrimitive: ArrowPrimitiveType { + fn scalar(value: Self::Native) -> ScalarValue; +} + +macro_rules! impl_scalar_value_from_primitive { + ($arrow_type:ty, $variant:ident) => { + impl ScalarValueFromPrimitive for $arrow_type { + fn scalar(value: Self::Native) -> ScalarValue { + ScalarValue::$variant(Some(value)) + } + } + }; +} + +impl_scalar_value_from_primitive!(Int8Type, Int8); +impl_scalar_value_from_primitive!(Int16Type, Int16); +impl_scalar_value_from_primitive!(Int32Type, Int32); +impl_scalar_value_from_primitive!(Int64Type, Int64); +impl_scalar_value_from_primitive!(UInt8Type, UInt8); +impl_scalar_value_from_primitive!(UInt16Type, UInt16); +impl_scalar_value_from_primitive!(UInt32Type, UInt32); +impl_scalar_value_from_primitive!(UInt64Type, UInt64); +impl_scalar_value_from_primitive!(Float32Type, Float32); +impl_scalar_value_from_primitive!(Float64Type, Float64); +impl_scalar_value_from_primitive!(Date32Type, Date32); +impl_scalar_value_from_primitive!(Date64Type, Date64); +impl_scalar_value_from_primitive!(Time32SecondType, Time32Second); +impl_scalar_value_from_primitive!(Time32MillisecondType, Time32Millisecond); +impl_scalar_value_from_primitive!(Time64MicrosecondType, Time64Microsecond); +impl_scalar_value_from_primitive!(Time64NanosecondType, Time64Nanosecond); + +fn primitive_min_max(array: &PrimitiveArray) -> Option +where + T: ScalarValueFromPrimitive, +{ + let min = aggregate::min::(array)?; + let max = aggregate::max::(array)?; + Some((T::scalar(min), T::scalar(max))) +} + +fn timestamp_min_max(array: &PrimitiveArray, timezone: Option>) -> Option +where + T: ArrowTimestampType, +{ + let min = aggregate::min::(array)?; + let max = aggregate::max::(array)?; + Some(( + ScalarValue::new_timestamp::(Some(min), timezone.clone()), + ScalarValue::new_timestamp::(Some(max), timezone), + )) +} + +fn decimal128_min_max( + array: &PrimitiveArray, + precision: u8, + scale: i8, +) -> Option { + let min = aggregate::min::(array)?; + let max = aggregate::max::(array)?; + Some(( + ScalarValue::Decimal128(Some(min), precision, scale), + ScalarValue::Decimal128(Some(max), precision, scale), + )) +} + +fn decimal256_min_max( + array: &PrimitiveArray, + precision: u8, + scale: i8, +) -> Option { + let min: i256 = aggregate::min::(array)?; + let max: i256 = aggregate::max::(array)?; + Some(( + ScalarValue::Decimal256(Some(min), precision, scale), + ScalarValue::Decimal256(Some(max), precision, scale), + )) +} + +fn boolean_min_max(array: &BooleanArray) -> Option { + let min = aggregate::min_boolean(array)?; + let max = aggregate::max_boolean(array)?; + Some(( + ScalarValue::Boolean(Some(min)), + ScalarValue::Boolean(Some(max)), + )) +} + +fn string_min_max(array: &GenericStringArray) -> Option +where + O: arrow::array::OffsetSizeTrait, +{ + let min = aggregate::min_string(array)?; + let max = aggregate::max_string(array)?; + if O::IS_LARGE { + Some(( + ScalarValue::LargeUtf8(Some(min.to_string())), + ScalarValue::LargeUtf8(Some(max.to_string())), + )) + } else { + Some(( + ScalarValue::Utf8(Some(min.to_string())), + ScalarValue::Utf8(Some(max.to_string())), + )) + } +} + +fn string_view_min_max(array: &StringViewArray) -> Option { + let min = aggregate::min_string_view(array)?; + let max = aggregate::max_string_view(array)?; + Some(( + ScalarValue::Utf8View(Some(min.to_string())), + ScalarValue::Utf8View(Some(max.to_string())), + )) +} + +fn binary_min_max(array: &GenericBinaryArray) -> Option +where + O: arrow::array::OffsetSizeTrait, +{ + let min = aggregate::min_binary(array)?; + let max = aggregate::max_binary(array)?; + if O::IS_LARGE { + Some(( + ScalarValue::LargeBinary(Some(min.to_vec())), + ScalarValue::LargeBinary(Some(max.to_vec())), + )) + } else { + Some(( + ScalarValue::Binary(Some(min.to_vec())), + ScalarValue::Binary(Some(max.to_vec())), + )) + } +} + +fn binary_view_min_max(array: &BinaryViewArray) -> Option { + let min = aggregate::min_binary_view(array)?; + let max = aggregate::max_binary_view(array)?; + Some(( + ScalarValue::BinaryView(Some(min.to_vec())), + ScalarValue::BinaryView(Some(max.to_vec())), + )) +} + +fn fixed_size_binary_min_max(array: &FixedSizeBinaryArray) -> Option { + let min = aggregate::min_fixed_size_binary(array)?; + let max = aggregate::max_fixed_size_binary(array)?; + Some(( + ScalarValue::FixedSizeBinary(array.value_length(), Some(min.to_vec())), + ScalarValue::FixedSizeBinary(array.value_length(), Some(max.to_vec())), + )) +} + +/// Adapter implementing [`PruningStatistics`] for [`BatchStats`]. +pub(crate) struct BatchPruningStats<'a> { + stats: &'a BatchStats, + metadata: &'a RegionMetadataRef, +} + +impl<'a> BatchPruningStats<'a> { + /// Creates a new [`BatchPruningStats`]. + pub(crate) fn new(stats: &'a BatchStats, metadata: &'a RegionMetadataRef) -> Self { + Self { stats, metadata } + } +} + +impl PruningStatistics for BatchPruningStats<'_> { + fn min_values(&self, column: &Column) -> Option { + let col = self.metadata.column_by_name(&column.name)?; + self.stats.min_values(col.column_id) + } + + fn max_values(&self, column: &Column) -> Option { + let col = self.metadata.column_by_name(&column.name)?; + self.stats.max_values(col.column_id) + } + + fn num_containers(&self) -> usize { + self.stats.num_batches + } + + fn null_counts(&self, _column: &Column) -> Option { + None + } + + fn row_counts(&self, _column: &Column) -> Option { + None + } + + fn contained(&self, _column: &Column, _values: &HashSet) -> Option { + None + } +}