rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-08-18 20:18:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: collect bulk memtable column stats

Browse Source 
Signed-off-by: evenyag <realevenyag@gmail.com>
This commit is contained in:
evenyag
2026-07-02 19:35:22 +08:00
parent 0a002f4d65
commit eaeb0bfbc1
3 changed files with 639 additions and 231 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mito2/src/memtable/bulk.rs
+27 -4
View File
@@ -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<dyn IterBuilder> = 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<BulkIterContext>,
pub sequence: Option<SequenceRange>,
}
@@ -799,8 +816,7 @@ impl IterBuilder for BulkRangeIterBuilder {
_time_range: Option<(Timestamp, Timestamp)>,
metrics: Option<MemScanMetrics>,
) -> Result<BoxedRecordBatchIterator> {
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,
}
src/mito2/src/memtable/bulk/part.rs
+141 -227
View File
@@ -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<Self> {
// `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::<PrimaryKeyArray>()?;
let pk_values = pk_dict.values().as_any().downcast_ref::<BinaryArray>()?;
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<ArrayRef> {
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<ArrayRef> {
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<ArrayRef> {
None
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
None
}
fn contained(
&self,
_column: &Column,
_values: &std::collections::HashSet<datafusion_common::ScalarValue>,
) -> Option<BooleanArray> {
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<BatchStats>,
/// 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<RecordBatch> {
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));
}
}
src/mito2/src/memtable/bulk/stats.rs
+471
View File
@@ -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<ColumnId, ColumnStats>,
}
#[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<ArrayRef> {
self.columns
.get(&column_id)
.map(|stats| stats.min_values.clone())
}
fn max_values(&self, column_id: ColumnId) -> Option<ArrayRef> {
self.columns
.get(&column_id)
.map(|stats| stats.max_values.clone())
}
}
fn compute_column_stats(
batches: &[common_recordbatch::DfRecordBatch],
column: &ColumnMetadata,
) -> Option<ColumnStats> {
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<Option<ScalarPair>> {
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::<Int8Type>(array.as_primitive::<Int8Type>()),
ArrowDataType::Int16 => primitive_min_max::<Int16Type>(array.as_primitive::<Int16Type>()),
ArrowDataType::Int32 => primitive_min_max::<Int32Type>(array.as_primitive::<Int32Type>()),
ArrowDataType::Int64 => primitive_min_max::<Int64Type>(array.as_primitive::<Int64Type>()),
ArrowDataType::UInt8 => primitive_min_max::<UInt8Type>(array.as_primitive::<UInt8Type>()),
ArrowDataType::UInt16 => {
primitive_min_max::<UInt16Type>(array.as_primitive::<UInt16Type>())
}
ArrowDataType::UInt32 => {
primitive_min_max::<UInt32Type>(array.as_primitive::<UInt32Type>())
}
ArrowDataType::UInt64 => {
primitive_min_max::<UInt64Type>(array.as_primitive::<UInt64Type>())
}
ArrowDataType::Float32 => {
primitive_min_max::<Float32Type>(array.as_primitive::<Float32Type>())
}
ArrowDataType::Float64 => {
primitive_min_max::<Float64Type>(array.as_primitive::<Float64Type>())
}
ArrowDataType::Utf8 => string_min_max(array.as_string::<i32>()),
ArrowDataType::LargeUtf8 => string_min_max(array.as_string::<i64>()),
ArrowDataType::Utf8View => string_view_min_max(array.as_string_view()),
ArrowDataType::Binary => binary_min_max(array.as_binary::<i32>()),
ArrowDataType::LargeBinary => binary_min_max(array.as_binary::<i64>()),
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::<Date32Type>(array.as_primitive::<Date32Type>())
}
ArrowDataType::Date64 => {
primitive_min_max::<Date64Type>(array.as_primitive::<Date64Type>())
}
ArrowDataType::Timestamp(unit, timezone) => match unit {
arrow::datatypes::TimeUnit::Second => timestamp_min_max::<TimestampSecondType>(
array.as_primitive::<TimestampSecondType>(),
timezone.clone(),
),
arrow::datatypes::TimeUnit::Millisecond => {
timestamp_min_max::<TimestampMillisecondType>(
array.as_primitive::<TimestampMillisecondType>(),
timezone.clone(),
)
}
arrow::datatypes::TimeUnit::Microsecond => {
timestamp_min_max::<TimestampMicrosecondType>(
array.as_primitive::<TimestampMicrosecondType>(),
timezone.clone(),
)
}
arrow::datatypes::TimeUnit::Nanosecond => timestamp_min_max::<TimestampNanosecondType>(
array.as_primitive::<TimestampNanosecondType>(),
timezone.clone(),
),
},
ArrowDataType::Time32(unit) => match unit {
arrow::datatypes::TimeUnit::Second => {
primitive_min_max::<Time32SecondType>(array.as_primitive::<Time32SecondType>())
}
arrow::datatypes::TimeUnit::Millisecond => primitive_min_max::<Time32MillisecondType>(
array.as_primitive::<Time32MillisecondType>(),
),
_ => None,
},
ArrowDataType::Time64(unit) => match unit {
arrow::datatypes::TimeUnit::Microsecond => primitive_min_max::<Time64MicrosecondType>(
array.as_primitive::<Time64MicrosecondType>(),
),
arrow::datatypes::TimeUnit::Nanosecond => primitive_min_max::<Time64NanosecondType>(
array.as_primitive::<Time64NanosecondType>(),
),
_ => None,
},
ArrowDataType::Decimal128(precision, scale) => {
decimal128_min_max(array.as_primitive::<Decimal128Type>(), *precision, *scale)
}
ArrowDataType::Decimal256(precision, scale) => {
decimal256_min_max(array.as_primitive::<Decimal256Type>(), *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<T>(array: &PrimitiveArray<T>) -> Option<ScalarPair>
where
T: ScalarValueFromPrimitive,
{
let min = aggregate::min::<T>(array)?;
let max = aggregate::max::<T>(array)?;
Some((T::scalar(min), T::scalar(max)))
}
fn timestamp_min_max<T>(array: &PrimitiveArray<T>, timezone: Option<Arc<str>>) -> Option<ScalarPair>
where
T: ArrowTimestampType,
{
let min = aggregate::min::<T>(array)?;
let max = aggregate::max::<T>(array)?;
Some((
ScalarValue::new_timestamp::<T>(Some(min), timezone.clone()),
ScalarValue::new_timestamp::<T>(Some(max), timezone),
))
}
fn decimal128_min_max(
array: &PrimitiveArray<Decimal128Type>,
precision: u8,
scale: i8,
) -> Option<ScalarPair> {
let min = aggregate::min::<Decimal128Type>(array)?;
let max = aggregate::max::<Decimal128Type>(array)?;
Some((
ScalarValue::Decimal128(Some(min), precision, scale),
ScalarValue::Decimal128(Some(max), precision, scale),
))
}
fn decimal256_min_max(
array: &PrimitiveArray<Decimal256Type>,
precision: u8,
scale: i8,
) -> Option<ScalarPair> {
let min: i256 = aggregate::min::<Decimal256Type>(array)?;
let max: i256 = aggregate::max::<Decimal256Type>(array)?;
Some((
ScalarValue::Decimal256(Some(min), precision, scale),
ScalarValue::Decimal256(Some(max), precision, scale),
))
}
fn boolean_min_max(array: &BooleanArray) -> Option<ScalarPair> {
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<O>(array: &GenericStringArray<O>) -> Option<ScalarPair>
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<ScalarPair> {
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<O>(array: &GenericBinaryArray<O>) -> Option<ScalarPair>
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<ScalarPair> {
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<ScalarPair> {
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<ArrayRef> {
let col = self.metadata.column_by_name(&column.name)?;
self.stats.min_values(col.column_id)
}
fn max_values(&self, column: &Column) -> Option<ArrayRef> {
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<ArrayRef> {
None
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
None
}
fn contained(&self, _column: &Column, _values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
None
}
}