From 3ce7e7bc81b215045260e028092127ffb1221408 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Mon, 19 Jan 2026 14:41:51 +0800 Subject: [PATCH] perf: accelerate pk filter (#7584) * perf: accelerate pk filter Signed-off-by: Ruihang Xia * add benchmark Signed-off-by: Ruihang Xia * unit tests Signed-off-by: Ruihang Xia * support compare ops Signed-off-by: Ruihang Xia * refactor to reduce dense decode/encode duplication Signed-off-by: Ruihang Xia * refactor if-else chain Signed-off-by: Ruihang Xia --------- Signed-off-by: Ruihang Xia --- src/common/recordbatch/src/filter.rs | 51 +++ src/mito-codec/Cargo.toml | 5 + .../benches/bench_primary_key_filter.rs | 294 ++++++++++++++ src/mito-codec/src/primary_key_filter.rs | 377 ++++++++++++++++-- src/mito-codec/src/row_converter/dense.rs | 155 +++++-- src/mito-codec/src/row_converter/sparse.rs | 74 ++++ 6 files changed, 876 insertions(+), 80 deletions(-) create mode 100644 src/mito-codec/benches/bench_primary_key_filter.rs diff --git a/src/common/recordbatch/src/filter.rs b/src/common/recordbatch/src/filter.rs index a62b4aa3945..d7c522e656a 100644 --- a/src/common/recordbatch/src/filter.rs +++ b/src/common/recordbatch/src/filter.rs @@ -32,6 +32,7 @@ use datatypes::arrow::compute::filter_record_batch; use datatypes::arrow::datatypes::{DataType, UInt32Type}; use datatypes::arrow::error::ArrowError; use datatypes::compute::or_kleene; +use datatypes::value::Value; use datatypes::vectors::VectorRef; use regex::Regex; use snafu::ResultExt; @@ -172,6 +173,56 @@ impl SimpleFilterEvaluator { &self.column_name } + pub fn is_eq(&self) -> bool { + matches!(self.op, Operator::Eq) + } + + pub fn is_not_eq(&self) -> bool { + matches!(self.op, Operator::NotEq) + } + + pub fn is_lt(&self) -> bool { + matches!(self.op, Operator::Lt) + } + + pub fn is_lt_eq(&self) -> bool { + matches!(self.op, Operator::LtEq) + } + + pub fn is_gt(&self) -> bool { + matches!(self.op, Operator::Gt) + } + + pub fn is_gt_eq(&self) -> bool { + matches!(self.op, Operator::GtEq) + } + + /// Returns true if this filter represents an `OR` chain of equality comparisons, e.g. + /// `col = lit1 OR col = lit2 ...`. + pub fn is_or_eq_chain(&self) -> bool { + matches!(self.op, Operator::Or) + } + + /// Returns the literal as a [`Value`]. It returns `None` if the literal can't be converted. + pub fn literal_value(&self) -> Option { + let array = self.literal.get().0; + let scalar = ScalarValue::try_from_array(array, 0).ok()?; + Value::try_from(scalar).ok() + } + + /// Returns the literal list as a list of [`Value`]s. It returns `None` if any literal can't be + /// converted. + pub fn literal_list_values(&self) -> Option> { + self.literal_list + .iter() + .map(|scalar| { + let array = scalar.get().0; + let scalar = ScalarValue::try_from_array(array, 0).ok()?; + Value::try_from(scalar).ok() + }) + .collect() + } + pub fn evaluate_scalar(&self, input: &ScalarValue) -> Result { let input = input .to_scalar() diff --git a/src/mito-codec/Cargo.toml b/src/mito-codec/Cargo.toml index 81808f2714d..07d64482e0f 100644 --- a/src/mito-codec/Cargo.toml +++ b/src/mito-codec/Cargo.toml @@ -35,3 +35,8 @@ datafusion-expr.workspace = true name = "bench_sparse_encoding" harness = false required-features = ["testing"] + +[[bench]] +name = "bench_primary_key_filter" +harness = false +required-features = ["testing"] diff --git a/src/mito-codec/benches/bench_primary_key_filter.rs b/src/mito-codec/benches/bench_primary_key_filter.rs new file mode 100644 index 00000000000..47158f87f8b --- /dev/null +++ b/src/mito-codec/benches/bench_primary_key_filter.rs @@ -0,0 +1,294 @@ +// 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. + +use std::hint::black_box; +use std::sync::Arc; + +use api::v1::SemanticType; +use common_query::prelude::{greptime_timestamp, greptime_value}; +use common_recordbatch::filter::SimpleFilterEvaluator; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::Column; +use datafusion_expr::{BinaryExpr, Expr, Literal, Operator}; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use datatypes::value::{Value, ValueRef}; +use mito_codec::row_converter::{ + DensePrimaryKeyCodec, PrimaryKeyCodec, PrimaryKeyCodecExt, SparsePrimaryKeyCodec, +}; +use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; +use store_api::storage::{ColumnId, RegionId}; + +fn setup_metadata() -> RegionMetadataRef { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("pod", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "namespace", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "container", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 3, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + greptime_value(), + ConcreteDataType::float64_datatype(), + false, + ), + semantic_type: SemanticType::Field, + column_id: 4, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + greptime_timestamp(), + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 5, + }) + .primary_key(vec![1, 2, 3]); + Arc::new(builder.build().unwrap()) +} + +fn create_test_row() -> Vec<(ColumnId, ValueRef<'static>)> { + vec![ + (1, ValueRef::String("greptime-frontend-6989d9899-22222")), + (2, ValueRef::String("greptime-cluster")), + (3, ValueRef::String("greptime-frontend-6989d9899-22222")), + ] +} + +fn create_eq_filter(column_name: &str, value: &str) -> SimpleFilterEvaluator { + create_filter_with_op(column_name, Operator::Eq, value) +} + +fn create_filter_with_op( + column_name: &str, + op: Operator, + value: T, +) -> SimpleFilterEvaluator { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name(column_name))), + op, + right: Box::new(value.lit()), + }); + SimpleFilterEvaluator::try_new(&expr).unwrap() +} + +fn create_or_eq_filter(column_name: &str, values: &[&str]) -> SimpleFilterEvaluator { + assert!(!values.is_empty()); + + let mut expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name(column_name))), + op: Operator::Eq, + right: Box::new(values[0].lit()), + }); + for value in &values[1..] { + let rhs = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name(column_name))), + op: Operator::Eq, + right: Box::new(value.lit()), + }); + expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(expr), + op: Operator::Or, + right: Box::new(rhs), + }); + } + SimpleFilterEvaluator::try_new(&expr).unwrap() +} + +fn encode_dense_pk(metadata: &RegionMetadataRef, row: &[(ColumnId, ValueRef<'static>)]) -> Vec { + let codec = DensePrimaryKeyCodec::new(metadata); + let mut pk = Vec::new(); + codec + .encode_to_vec(row.iter().map(|(_, v)| v.clone()), &mut pk) + .unwrap(); + pk +} + +fn encode_sparse_pk( + metadata: &RegionMetadataRef, + row: &[(ColumnId, ValueRef<'static>)], +) -> Vec { + let codec = SparsePrimaryKeyCodec::new(metadata); + let mut pk = Vec::new(); + codec.encode_to_vec(row.iter().cloned(), &mut pk).unwrap(); + pk +} + +fn matches_dense_scalar( + metadata: &RegionMetadataRef, + codec: &DensePrimaryKeyCodec, + filters: &[SimpleFilterEvaluator], + pk: &[u8], + offsets_buf: &mut Vec, +) -> bool { + offsets_buf.clear(); + if filters.is_empty() || metadata.primary_key.is_empty() { + return true; + } + + let mut result = true; + for filter in filters { + let Some(column) = metadata.column_by_name(filter.column_name()) else { + continue; + }; + if column.semantic_type != SemanticType::Tag { + continue; + } + + let index = metadata.primary_key_index(column.column_id).unwrap(); + let value = codec.decode_value_at(pk, index, offsets_buf).unwrap(); + let scalar_value = value + .try_to_scalar_value(&column.column_schema.data_type) + .unwrap(); + result &= filter.evaluate_scalar(&scalar_value).unwrap_or(true); + } + + result +} + +fn matches_sparse_scalar( + metadata: &RegionMetadataRef, + codec: &SparsePrimaryKeyCodec, + filters: &[SimpleFilterEvaluator], + pk: &[u8], + offsets_map: &mut std::collections::HashMap, +) -> bool { + offsets_map.clear(); + if filters.is_empty() || metadata.primary_key.is_empty() { + return true; + } + + let mut result = true; + for filter in filters { + let Some(column) = metadata.column_by_name(filter.column_name()) else { + continue; + }; + if column.semantic_type != SemanticType::Tag { + continue; + } + + let value = if let Some(offset) = codec.has_column(pk, offsets_map, column.column_id) { + codec.decode_value_at(pk, offset, column.column_id).unwrap() + } else { + Value::Null + }; + let scalar_value = value + .try_to_scalar_value(&column.column_schema.data_type) + .unwrap(); + result &= filter.evaluate_scalar(&scalar_value).unwrap_or(true); + } + + result +} + +fn bench_primary_key_filter(c: &mut Criterion) { + let metadata = setup_metadata(); + let row = create_test_row(); + + let eq_filter = create_eq_filter("pod", "greptime-frontend-6989d9899-22222"); + let gt_filter = create_filter_with_op("pod", Operator::Gt, "greptime-frontend-6989d9899-22221"); + let lt_eq_filter = + create_filter_with_op("pod", Operator::LtEq, "greptime-frontend-6989d9899-22222"); + let or_filter = create_or_eq_filter( + "pod", + &[ + "a", + "b", + "c", + "d", + "e", + "f", + "g", + "greptime-frontend-6989d9899-22222", + ], + ); + + let cases = [ + ("eq", eq_filter), + ("gt", gt_filter), + ("lt_eq", lt_eq_filter), + ("or_eq", or_filter), + ]; + + for (case_name, filter) in cases { + let filters = Arc::new(vec![filter]); + + let dense_pk = encode_dense_pk(&metadata, &row); + let dense_codec = DensePrimaryKeyCodec::new(&metadata); + let mut dense_fast = dense_codec.primary_key_filter(&metadata, filters.clone()); + let mut dense_offsets = Vec::new(); + + let sparse_pk = encode_sparse_pk(&metadata, &row); + let sparse_codec = SparsePrimaryKeyCodec::new(&metadata); + let mut sparse_fast = sparse_codec.primary_key_filter(&metadata, filters.clone()); + let mut sparse_offsets = std::collections::HashMap::new(); + + let mut group = c.benchmark_group(format!("primary_key_filter/{case_name}")); + + group.bench_function("dense/fast", |b| { + b.iter(|| black_box(dense_fast.matches(black_box(&dense_pk)))) + }); + group.bench_function("dense/scalar", |b| { + b.iter(|| { + black_box(matches_dense_scalar( + &metadata, + &dense_codec, + filters.as_ref(), + black_box(&dense_pk), + &mut dense_offsets, + )) + }) + }); + + group.bench_function("sparse/fast", |b| { + b.iter(|| black_box(sparse_fast.matches(black_box(&sparse_pk)))) + }); + group.bench_function("sparse/scalar", |b| { + b.iter(|| { + black_box(matches_sparse_scalar( + &metadata, + &sparse_codec, + filters.as_ref(), + black_box(&sparse_pk), + &mut sparse_offsets, + )) + }) + }); + + group.finish(); + } +} + +criterion_group!(benches, bench_primary_key_filter); +criterion_main!(benches); diff --git a/src/mito-codec/src/primary_key_filter.rs b/src/mito-codec/src/primary_key_filter.rs index c71fafc974f..189c7a08cd2 100644 --- a/src/mito-codec/src/primary_key_filter.rs +++ b/src/mito-codec/src/primary_key_filter.rs @@ -17,13 +17,17 @@ use std::sync::Arc; use api::v1::SemanticType; use common_recordbatch::filter::SimpleFilterEvaluator; +use datatypes::data_type::ConcreteDataType; use datatypes::value::Value; +use memcomparable::Serializer; use store_api::metadata::RegionMetadataRef; use store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME; use store_api::storage::ColumnId; use crate::error::Result; -use crate::row_converter::{DensePrimaryKeyCodec, PrimaryKeyFilter, SparsePrimaryKeyCodec}; +use crate::row_converter::{ + DensePrimaryKeyCodec, PrimaryKeyFilter, SortField, SparsePrimaryKeyCodec, +}; /// Returns true if this is a partition column for metrics in the memtable. pub fn is_partition_column(name: &str) -> bool { @@ -32,52 +36,213 @@ pub fn is_partition_column(name: &str) -> bool { #[derive(Clone)] struct PrimaryKeyFilterInner { - metadata: RegionMetadataRef, filters: Arc>, + compiled_filters: Vec, } impl PrimaryKeyFilterInner { - fn evaluate_filters( - &self, - mut decode_fn: impl FnMut(ColumnId, &RegionMetadataRef) -> Result, - ) -> bool { - if self.filters.is_empty() || self.metadata.primary_key.is_empty() { - return true; + fn new(metadata: RegionMetadataRef, filters: Arc>) -> Self { + let compiled_filters = Self::compile_filters(&metadata, &filters); + Self { + filters, + compiled_filters, + } + } + + fn compile_filters( + metadata: &RegionMetadataRef, + filters: &[SimpleFilterEvaluator], + ) -> Vec { + if filters.is_empty() || metadata.primary_key.is_empty() { + return Vec::new(); } - let mut result = true; - for filter in self.filters.iter() { + let mut compiled_filters = Vec::with_capacity(filters.len()); + for (filter_idx, filter) in filters.iter().enumerate() { if is_partition_column(filter.column_name()) { continue; } - let Some(column) = self.metadata.column_by_name(filter.column_name()) else { + + let Some(column) = metadata.column_by_name(filter.column_name()) else { continue; }; - // ignore filters that are not referencing primary key columns + // Ignore filters that are not referencing primary key tag columns. if column.semantic_type != SemanticType::Tag { continue; } - let value = match decode_fn(column.column_id, &self.metadata) { - Ok(v) => v, - Err(e) => { - common_telemetry::error!(e; "Failed to decode primary key"); - return true; - } + // Safety: A tag column is always in primary key. + let Some(pk_index) = metadata.primary_key_index(column.column_id) else { + continue; }; - // TODO(yingwen): `evaluate_scalar()` creates temporary arrays to compare scalars. We - // can compare the bytes directly without allocation and matching types as we use - // comparable encoding. - // Safety: arrow schema and datatypes are constructed from the same source. - let scalar_value = value - .try_to_scalar_value(&column.column_schema.data_type) - .unwrap(); - result &= filter.evaluate_scalar(&scalar_value).unwrap_or(true); + let data_type = column.column_schema.data_type.clone(); + let fast_path = CompiledFastPath::try_new(filter, &data_type); + + compiled_filters.push(CompiledPrimaryKeyFilter { + filter_idx, + column_id: column.column_id, + pk_index, + data_type, + fast_path, + }); } - result + compiled_filters } + + fn evaluate_filters<'a>(&self, accessor: &mut impl PrimaryKeyValueAccessor<'a>) -> bool { + if self.compiled_filters.is_empty() { + return true; + } + + for compiled in &self.compiled_filters { + let filter = &self.filters[compiled.filter_idx]; + + let passed = if let Some(fast_path) = &compiled.fast_path { + let encoded_value = match accessor.encoded_value(compiled) { + Ok(v) => v, + Err(e) => { + common_telemetry::error!(e; "Failed to decode primary key"); + return true; + } + }; + fast_path.matches(encoded_value) + } else { + let value = match accessor.decode_value(compiled) { + Ok(v) => v, + Err(e) => { + common_telemetry::error!(e; "Failed to decode primary key"); + return true; + } + }; + + // Safety: arrow schema and datatypes are constructed from the same source. + let scalar_value = value.try_to_scalar_value(&compiled.data_type).unwrap(); + filter.evaluate_scalar(&scalar_value).unwrap_or(true) + }; + + if !passed { + return false; + } + } + + true + } +} + +#[derive(Clone)] +struct CompiledPrimaryKeyFilter { + filter_idx: usize, + column_id: ColumnId, + pk_index: usize, + data_type: ConcreteDataType, + fast_path: Option, +} + +#[derive(Clone)] +enum CompiledFastPath { + Eq(Vec), + NotEq(Vec), + Lt(Vec), + LtEq(Vec), + Gt(Vec), + GtEq(Vec), + InList(Vec>), +} + +impl CompiledFastPath { + fn try_new(filter: &SimpleFilterEvaluator, data_type: &ConcreteDataType) -> Option { + let field = SortField::new(data_type.clone()); + // Safety guard: memcomparable defines a total order for floats that doesn't necessarily + // match DataFusion/Arrow float comparison semantics (e.g. `-0.0 == 0.0`). + if field.encode_data_type().is_float() { + return None; + } + + let encoded = |value: &Value| -> Option> { + let mut buf = Vec::new(); + let mut serializer = Serializer::new(&mut buf); + field + .serialize(&mut serializer, &value.as_value_ref()) + .ok()?; + Some(buf) + }; + + if filter.is_eq() { + let value = filter.literal_value()?; + Some(Self::Eq(encoded(&value)?)) + } else if filter.is_not_eq() { + let value = filter.literal_value()?; + Some(Self::NotEq(encoded(&value)?)) + } else if filter.is_lt() { + let value = filter.literal_value()?; + Some(Self::Lt(encoded(&value)?)) + } else if filter.is_lt_eq() { + let value = filter.literal_value()?; + Some(Self::LtEq(encoded(&value)?)) + } else if filter.is_gt() { + let value = filter.literal_value()?; + Some(Self::Gt(encoded(&value)?)) + } else if filter.is_gt_eq() { + let value = filter.literal_value()?; + Some(Self::GtEq(encoded(&value)?)) + } else if filter.is_or_eq_chain() { + let values = filter.literal_list_values()?; + let mut list = Vec::with_capacity(values.len()); + for value in values { + let bytes = encoded(&value)?; + // `NULL` never matches in equality comparisons (`NULL = x` is NULL). + if bytes.first() == Some(&0) { + continue; + } + list.push(bytes); + } + Some(Self::InList(list)) + } else { + None + } + } + + fn matches(&self, encoded_value: Option<&[u8]>) -> bool { + let Some(encoded_value) = encoded_value else { + return false; + }; + + // `NULL` never matches in comparisons (`NULL = x` and `NULL != x` are NULL). + if encoded_value.first() == Some(&0) { + return false; + } + + match self { + CompiledFastPath::Eq(encoded_literal) => { + encoded_literal.first() != Some(&0) && encoded_value == &encoded_literal[..] + } + CompiledFastPath::NotEq(encoded_literal) => { + encoded_literal.first() != Some(&0) && encoded_value != &encoded_literal[..] + } + CompiledFastPath::Lt(encoded_literal) => { + encoded_literal.first() != Some(&0) && encoded_value < &encoded_literal[..] + } + CompiledFastPath::LtEq(encoded_literal) => { + encoded_literal.first() != Some(&0) && encoded_value <= &encoded_literal[..] + } + CompiledFastPath::Gt(encoded_literal) => { + encoded_literal.first() != Some(&0) && encoded_value > &encoded_literal[..] + } + CompiledFastPath::GtEq(encoded_literal) => { + encoded_literal.first() != Some(&0) && encoded_value >= &encoded_literal[..] + } + CompiledFastPath::InList(encoded_literals) => { + encoded_literals.iter().any(|lit| encoded_value == &lit[..]) + } + } + } +} + +trait PrimaryKeyValueAccessor<'a> { + fn encoded_value(&mut self, filter: &CompiledPrimaryKeyFilter) -> Result>; + fn decode_value(&mut self, filter: &CompiledPrimaryKeyFilter) -> Result; } /// Dense primary key filter. @@ -95,7 +260,7 @@ impl DensePrimaryKeyFilter { codec: DensePrimaryKeyCodec, ) -> Self { Self { - inner: PrimaryKeyFilterInner { metadata, filters }, + inner: PrimaryKeyFilterInner::new(metadata, filters), codec, offsets_buf: Vec::new(), } @@ -105,12 +270,31 @@ impl DensePrimaryKeyFilter { impl PrimaryKeyFilter for DensePrimaryKeyFilter { fn matches(&mut self, pk: &[u8]) -> bool { self.offsets_buf.clear(); - self.inner.evaluate_filters(|column_id, metadata| { - // index of tag column in primary key - // Safety: A tag column is always in primary key. - let index = metadata.primary_key_index(column_id).unwrap(); - self.codec.decode_value_at(pk, index, &mut self.offsets_buf) - }) + let mut accessor = DensePrimaryKeyValueAccessor { + pk, + codec: &self.codec, + offsets_buf: &mut self.offsets_buf, + }; + self.inner.evaluate_filters(&mut accessor) + } +} + +struct DensePrimaryKeyValueAccessor<'a, 'b> { + pk: &'a [u8], + codec: &'b DensePrimaryKeyCodec, + offsets_buf: &'b mut Vec, +} + +impl<'a> PrimaryKeyValueAccessor<'a> for DensePrimaryKeyValueAccessor<'a, '_> { + fn encoded_value(&mut self, filter: &CompiledPrimaryKeyFilter) -> Result> { + self.codec + .encoded_value_at(self.pk, filter.pk_index, self.offsets_buf) + .map(Some) + } + + fn decode_value(&mut self, filter: &CompiledPrimaryKeyFilter) -> Result { + self.codec + .decode_value_at(self.pk, filter.pk_index, self.offsets_buf) } } @@ -129,7 +313,7 @@ impl SparsePrimaryKeyFilter { codec: SparsePrimaryKeyCodec, ) -> Self { Self { - inner: PrimaryKeyFilterInner { metadata, filters }, + inner: PrimaryKeyFilterInner::new(metadata, filters), codec, offsets_map: HashMap::new(), } @@ -139,13 +323,37 @@ impl SparsePrimaryKeyFilter { impl PrimaryKeyFilter for SparsePrimaryKeyFilter { fn matches(&mut self, pk: &[u8]) -> bool { self.offsets_map.clear(); - self.inner.evaluate_filters(|column_id, _| { - if let Some(offset) = self.codec.has_column(pk, &mut self.offsets_map, column_id) { - self.codec.decode_value_at(pk, offset, column_id) - } else { - Ok(Value::Null) - } - }) + let mut accessor = SparsePrimaryKeyValueAccessor { + pk, + codec: &self.codec, + offsets_map: &mut self.offsets_map, + }; + self.inner.evaluate_filters(&mut accessor) + } +} + +struct SparsePrimaryKeyValueAccessor<'a, 'b> { + pk: &'a [u8], + codec: &'b SparsePrimaryKeyCodec, + offsets_map: &'b mut HashMap, +} + +impl<'a> PrimaryKeyValueAccessor<'a> for SparsePrimaryKeyValueAccessor<'a, '_> { + fn encoded_value(&mut self, filter: &CompiledPrimaryKeyFilter) -> Result> { + self.codec + .encoded_value_for_column(self.pk, self.offsets_map, filter.column_id) + } + + fn decode_value(&mut self, filter: &CompiledPrimaryKeyFilter) -> Result { + if let Some(offset) = self + .codec + .has_column(self.pk, self.offsets_map, filter.column_id) + { + self.codec + .decode_value_at(self.pk, offset, filter.column_id) + } else { + Ok(Value::Null) + } } } @@ -159,7 +367,7 @@ mod tests { use datafusion_expr::{BinaryExpr, Expr, Literal, Operator}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; - use datatypes::value::ValueRef; + use datatypes::value::{OrderedFloat, ValueRef}; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::{ColumnId, RegionId}; @@ -224,9 +432,17 @@ mod tests { } fn create_filter(column_name: &str, value: &str) -> SimpleFilterEvaluator { + create_filter_with_op(column_name, Operator::Eq, value) + } + + fn create_filter_with_op( + column_name: &str, + op: Operator, + value: T, + ) -> SimpleFilterEvaluator { let expr = Expr::BinaryExpr(BinaryExpr { left: Box::new(Expr::Column(Column::from_name(column_name))), - op: Operator::Eq, + op, right: Box::new(value.lit()), }); SimpleFilterEvaluator::try_new(&expr).unwrap() @@ -331,4 +547,77 @@ mod tests { let mut filter = DensePrimaryKeyFilter::new(metadata, filters, codec); assert!(filter.matches(&pk)); } + + #[test] + fn test_dense_primary_key_filter_order_ops() { + let metadata = setup_metadata(); + let pk = encode_dense_pk(&metadata, create_test_row()); + let codec = DensePrimaryKeyCodec::new(&metadata); + + let cases = [ + (Operator::Gt, "greptime-frontend-6989d9899-22221", true), + (Operator::GtEq, "greptime-frontend-6989d9899-22222", true), + (Operator::Lt, "greptime-frontend-6989d9899-22223", true), + (Operator::LtEq, "greptime-frontend-6989d9899-22221", false), + ]; + + for (op, value, expected) in cases { + let filters = Arc::new(vec![create_filter_with_op("pod", op, value)]); + let mut filter = DensePrimaryKeyFilter::new(metadata.clone(), filters, codec.clone()); + assert_eq!(expected, filter.matches(&pk)); + } + } + + #[test] + fn test_sparse_primary_key_filter_order_ops() { + let metadata = setup_metadata(); + let pk = encode_sparse_pk(&metadata, create_test_row()); + let codec = SparsePrimaryKeyCodec::new(&metadata); + + let cases = [ + (Operator::Gt, "greptime-frontend-6989d9899-22221", true), + (Operator::GtEq, "greptime-frontend-6989d9899-22222", true), + (Operator::Lt, "greptime-frontend-6989d9899-22223", true), + (Operator::LtEq, "greptime-frontend-6989d9899-22221", false), + ]; + + for (op, value, expected) in cases { + let filters = Arc::new(vec![create_filter_with_op("pod", op, value)]); + let mut filter = SparsePrimaryKeyFilter::new(metadata.clone(), filters, codec.clone()); + assert_eq!(expected, filter.matches(&pk)); + } + } + + #[test] + fn test_dense_primary_key_filter_float_eq_fallback() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("f", ConcreteDataType::float64_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + greptime_timestamp(), + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 2, + }) + .primary_key(vec![1]); + let metadata = Arc::new(builder.build().unwrap()); + + let codec = DensePrimaryKeyCodec::new(&metadata); + let mut pk = Vec::new(); + codec + .encode_to_vec([ValueRef::Float64(OrderedFloat(-0.0))].into_iter(), &mut pk) + .unwrap(); + + let filters = Arc::new(vec![create_filter_with_op("f", Operator::Eq, 0.0_f64)]); + let mut filter = DensePrimaryKeyFilter::new(metadata, filters, codec); + + assert!(filter.matches(&pk)); + } } diff --git a/src/mito-codec/src/row_converter/dense.rs b/src/mito-codec/src/row_converter/dense.rs index 961e36115aa..6cc70feaeaf 100644 --- a/src/mito-codec/src/row_converter/dense.rs +++ b/src/mito-codec/src/row_converter/dense.rs @@ -430,6 +430,52 @@ impl DensePrimaryKeyCodec { &self.ordered_primary_key_columns[pos].1 } + /// Advances `deserializer` to the start of value at `pos`. + /// + /// Returns the offset of the value at `pos` in `bytes`. + fn advance_to_value_at( + &self, + bytes: &[u8], + pos: usize, + offsets_buf: &mut Vec, + deserializer: &mut Deserializer<&[u8]>, + ) -> Result { + if pos < offsets_buf.len() { + // We computed the offset before. + let offset = offsets_buf[pos]; + deserializer.advance(offset); + return Ok(offset); + } + + if offsets_buf.is_empty() { + let mut offset = 0; + // Skip values before `pos`. + for i in 0..pos { + // Offset to skip before reading value i. + offsets_buf.push(offset); + let skip = self.field_at(i).skip_deserialize(bytes, deserializer)?; + offset += skip; + } + // Offset to skip before reading this value. + offsets_buf.push(offset); + Ok(offset) + } else { + // Offsets are not enough. + let value_start = offsets_buf.len() - 1; + // Advances to decode value at `value_start`. + let mut offset = offsets_buf[value_start]; + deserializer.advance(offset); + for i in value_start..pos { + // Skip value i. + let skip = self.field_at(i).skip_deserialize(bytes, deserializer)?; + // Offset for the value at i + 1. + offset += skip; + offsets_buf.push(offset); + } + Ok(offset) + } + } + /// Decode value at `pos` in `bytes`. /// /// The i-th element in offsets buffer is how many bytes to skip in order to read value at `pos`. @@ -440,46 +486,30 @@ impl DensePrimaryKeyCodec { offsets_buf: &mut Vec, ) -> Result { let mut deserializer = Deserializer::new(bytes); - if pos < offsets_buf.len() { - // We computed the offset before. - let to_skip = offsets_buf[pos]; - deserializer.advance(to_skip); - return self.field_at(pos).deserialize(&mut deserializer); - } - - if offsets_buf.is_empty() { - let mut offset = 0; - // Skip values before `pos`. - for i in 0..pos { - // Offset to skip before reading value i. - offsets_buf.push(offset); - let skip = self - .field_at(i) - .skip_deserialize(bytes, &mut deserializer)?; - offset += skip; - } - // Offset to skip before reading this value. - offsets_buf.push(offset); - } else { - // Offsets are not enough. - let value_start = offsets_buf.len() - 1; - // Advances to decode value at `value_start`. - let mut offset = offsets_buf[value_start]; - deserializer.advance(offset); - for i in value_start..pos { - // Skip value i. - let skip = self - .field_at(i) - .skip_deserialize(bytes, &mut deserializer)?; - // Offset for the value at i + 1. - offset += skip; - offsets_buf.push(offset); - } - } + self.advance_to_value_at(bytes, pos, offsets_buf, &mut deserializer)?; self.field_at(pos).deserialize(&mut deserializer) } + /// Returns the encoded bytes at `pos` in `bytes`. + /// + /// The i-th element in offsets buffer is how many bytes to skip in order to read value at + /// `pos`. + pub fn encoded_value_at<'a>( + &self, + bytes: &'a [u8], + pos: usize, + offsets_buf: &mut Vec, + ) -> Result<&'a [u8]> { + let mut deserializer = Deserializer::new(bytes); + let offset = self.advance_to_value_at(bytes, pos, offsets_buf, &mut deserializer)?; + + let len = self + .field_at(pos) + .skip_deserialize(bytes, &mut deserializer)?; + Ok(&bytes[offset..offset + len]) + } + pub fn estimated_size(&self) -> usize { self.ordered_primary_key_columns .iter() @@ -666,6 +696,59 @@ mod tests { ) } + #[test] + fn test_encoded_value_at() { + let data_types = [ + ConcreteDataType::string_datatype(), + ConcreteDataType::int32_datatype(), + ConcreteDataType::string_datatype(), + ]; + let encoder = DensePrimaryKeyCodec::with_fields( + data_types + .iter() + .enumerate() + .map(|(idx, t)| (idx as ColumnId, SortField::new(t.clone()))) + .collect::>(), + ); + let row = [Value::String("hello".into()), Value::Int32(42), Value::Null]; + let value_ref = row.iter().map(|v| v.as_value_ref()).collect::>(); + let encoded_pk = encoder.encode(value_ref.iter().cloned()).unwrap(); + + let mut offsets = Vec::new(); + let mut combined = Vec::new(); + for (pos, value) in row.iter().enumerate() { + let encoded_value = encoder + .encoded_value_at(&encoded_pk, pos, &mut offsets) + .unwrap(); + combined.extend_from_slice(encoded_value); + + let field = SortField::new(data_types[pos].clone()); + let mut expected = Vec::new(); + let mut serializer = Serializer::new(&mut expected); + field + .serialize(&mut serializer, &value.as_value_ref()) + .unwrap(); + assert_eq!(encoded_value, expected.as_slice()); + } + assert_eq!(combined, encoded_pk); + assert_eq!(offsets.len(), row.len()); + + // Verify the offsets buffer can be reused for random access. + for (pos, value) in row.iter().enumerate().rev() { + let encoded_value = encoder + .encoded_value_at(&encoded_pk, pos, &mut offsets) + .unwrap(); + + let field = SortField::new(data_types[pos].clone()); + let mut expected = Vec::new(); + let mut serializer = Serializer::new(&mut expected); + field + .serialize(&mut serializer, &value.as_value_ref()) + .unwrap(); + assert_eq!(encoded_value, expected.as_slice()); + } + } + #[test] fn test_memcmp_dictionary() { // Test Dictionary diff --git a/src/mito-codec/src/row_converter/sparse.rs b/src/mito-codec/src/row_converter/sparse.rs index 731de5c0b08..4638ddcefb6 100644 --- a/src/mito-codec/src/row_converter/sparse.rs +++ b/src/mito-codec/src/row_converter/sparse.rs @@ -296,6 +296,29 @@ impl SparsePrimaryKeyCodec { let field = self.get_field(column_id).unwrap(); field.deserialize(&mut deserializer) } + + /// Returns the encoded bytes of the given `column_id` in `pk`. + /// + /// Returns `Ok(None)` if the `column_id` is missing in `pk`. + pub fn encoded_value_for_column<'a>( + &self, + pk: &'a [u8], + offsets_map: &mut HashMap, + column_id: ColumnId, + ) -> Result> { + let Some(offset) = self.has_column(pk, offsets_map, column_id) else { + return Ok(None); + }; + + let Some(field) = self.get_field(column_id) else { + return Ok(None); + }; + + let mut deserializer = Deserializer::new(pk); + deserializer.advance(offset); + let len = field.skip_deserialize(pk, &mut deserializer)?; + Ok(Some(&pk[offset..offset + len])) + } } impl PrimaryKeyCodec for SparsePrimaryKeyCodec { @@ -693,4 +716,55 @@ mod tests { assert_eq!(value.as_value_ref(), expected_value); } } + + #[test] + fn test_encoded_value_for_column() { + let region_metadata = test_region_metadata(); + let codec = SparsePrimaryKeyCodec::new(®ion_metadata); + let mut buffer = Vec::new(); + let row = test_row(); + codec + .encode_to_vec(row.clone().into_iter(), &mut buffer) + .unwrap(); + assert!(!buffer.is_empty()); + + let mut offsets_map = HashMap::new(); + for column_id in [ + RESERVED_COLUMN_ID_TABLE_ID, + RESERVED_COLUMN_ID_TSID, + 1, + 2, + 3, + 4, + 5, + ] { + let encoded_value = codec + .encoded_value_for_column(&buffer, &mut offsets_map, column_id) + .unwrap() + .unwrap(); + let expected_value = row + .iter() + .find(|(id, _)| *id == column_id) + .unwrap() + .1 + .clone(); + let data_type = match column_id { + RESERVED_COLUMN_ID_TABLE_ID => ConcreteDataType::uint32_datatype(), + RESERVED_COLUMN_ID_TSID => ConcreteDataType::uint64_datatype(), + _ => ConcreteDataType::string_datatype(), + }; + let field = SortField::new(data_type); + let mut expected_encoded = Vec::new(); + let mut serializer = Serializer::new(&mut expected_encoded); + field.serialize(&mut serializer, &expected_value).unwrap(); + assert_eq!(encoded_value, expected_encoded.as_slice()); + } + + for column_id in [6_u32, 7_u32, 999_u32] { + let encoded_value = codec + .encoded_value_for_column(&buffer, &mut offsets_map, column_id) + .unwrap(); + assert!(encoded_value.is_none()); + } + } }