rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-29 15:06:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(partition): support iceberg_hash in partition rules (#9693)

Browse Source 
* feat(partition): support iceberg_hash in partition rules

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(partition): reject partition column type changes and quote rule identifiers

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(table): keep time index widening for partition columns

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* refactor(partition): replace PartitionFunction with IcebergHash

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* docs(partition): describe iceberg_hash in the function module

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(meta): validate iceberg_hash partition bounds on ALTER and trim comments

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* test: quote partition rule identifiers in SHOW CREATE TABLE expectations

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* test(partition): check hash rules in information_schema and repartition target columns

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(partition): coerce hash partition rules against the evaluation schema

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(partition): keep quotes in partition column names when building filters

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
dennis zhuang
2026-10-27 09:23:00 +00:00
committed by GitHub
parent 2117f8b059
commit 2e84833ac8
29 changed files with 1947 additions and 351 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/meta/src/ddl/alter_table.rs
+3 -7
View File
@@ -150,8 +150,6 @@ impl AlterTableProcedure {
)?;
self.new_table_info = Some(new_table_info);
// Reject schema changes that the persisted partition expressions of some
// region cannot survive, before any region is altered.
self.validate_region_partition_exprs().await?;
// Safety: Checked in `AlterTableProcedure::new`.
@@ -248,9 +246,8 @@ impl AlterTableProcedure {
Ok(Status::executing(true))
}
/// Validates the persisted partition expressions of every region against the
/// schema proposed by this ALTER TABLE, so an ALTER that would break a region's
/// partition predicate is rejected before any region is modified.
/// Rejects type changes the persisted partition expressions can't take, before
/// any region is altered.
async fn validate_region_partition_exprs(&self) -> Result<()> {
// Safety: filled in `fill_table_info`.
let old_schema = &self
@@ -275,8 +272,7 @@ impl AlterTableProcedure {
.get_physical_table_route(self.data.table_id())
.await?;
if physical_table_id != self.data.table_id() {
// Partition expressions are persisted on the regions of the physical
// table; the altered schema here belongs to a logical table.
// Partition expressions live on the physical table's regions.
return Ok(());
}
src/common/meta/src/ddl/alter_table/partition_compat.rs
+55 -118
View File
@@ -12,39 +12,17 @@
// See the License for the specific language governing permissions and
// limitations under the License.
//! Preflight validation of persisted partition expressions against an ALTER TABLE schema change.
//! Rejects an ALTER TABLE whose column type changes would break the persisted
//! partition expressions of a region, before any region is altered.
//!
//! Persisted partition bounds keep the types they were written with. A column whose
//! type changes while a partition expression references it breaks the region on both
//! paths that consume the bounds:
//! Bounds keep the type they were written with. Row routing compares them with
//! row values as `Value`s, which only stays valid across a timestamp unit
//! widening; scans coerce the bound to the common comparison type, which can
//! overflow (a millisecond bound in year 3000 cast to nanoseconds).
//!
//! * On scan, the region coerces the bound literal and the new column type to a
//! common comparison type while building the partition predicate. If no common type
//! exists, or the bound cannot be cast to it (a classic example is a millisecond
//! timestamp bound that overflows when the time index column is widened to
//! nanoseconds), every scan of that region fails at that point.
//! * On write, the region routes rows through `MultiDimPartitionRule::find_region`,
//! which compares the persisted bound against row values directly (`Value::cmp`): a
//! bound whose variant no longer matches the column's panics on ordered comparisons
//! and misroutes silently on equality. Only a timestamp unit widening keeps the two
//! comparable — `Timestamp` comparisons convert units before comparing — so every
//! other type change of a referenced column is rejected outright, and the widening
//! keeps the scan check below.
//!
//! Such an ALTER must be rejected before it is submitted to any region, otherwise
//! already-altered regions keep failing while the rest of the table is untouched.
//!
//! The `partition` crate depends on `common-meta`, so this crate cannot depend on
//! `partition` back and cannot reuse [`partition::expr::PartitionExpr`]. The `Persisted*`
//! types below mirror the JSON representation of `partition::partition::PartitionBound` and
//! `partition::expr::PartitionExpr` and must be kept in sync with `partition/src/partition.rs`
//! and `partition/src/expr.rs`; round-trip tests in `partition` run real serialized
//! expressions through this validator so drift fails in CI. The compatibility check itself
//! deliberately repeats the same conversion (value -> [`ScalarValue`]) and the same
//! comparison coercion the region performs when it builds the partition predicate
//! (`ExprSimplifier::coerce` / [`comparison_coercion`]): both sides are cast to the common
//! comparison type, and the persisted bound must be castable to that type with
//! [`ScalarValue::cast_to`]. A bound that would fail on scan is rejected here as well.
//! `partition` depends on this crate, so the `Persisted*` types mirror the JSON of
//! `partition::partition::PartitionBound` and `partition::expr::PartitionExpr`.
//! A test in `partition` runs real serialized expressions through this validator.
use std::collections::HashMap;
@@ -62,10 +40,9 @@ use store_api::storage::RegionId;
use crate::error::{PartitionExprIncompatibleSnafu, Result};
use crate::rpc::router::RegionRoute;
/// Mirrors `partition::partition::PartitionBound`.
#[derive(Debug, Deserialize)]
enum PersistedPartitionBound {
/// Deprecated since 0.9.0. Only deserialized to keep unknown old formats parseable.
/// Deprecated since 0.9.0.
#[allow(dead_code)]
Value(Value),
/// Deprecated since 0.15.0.
@@ -73,7 +50,6 @@ enum PersistedPartitionBound {
Expr(PersistedPartitionExpr),
}
/// Mirrors `partition::expr::PartitionExpr`.
#[derive(Debug, Deserialize)]
struct PersistedPartitionExpr {
lhs: Box<PersistedOperand>,
@@ -81,15 +57,15 @@ struct PersistedPartitionExpr {
rhs: Box<PersistedOperand>,
}
/// Mirrors `partition::expr::Operand`.
#[derive(Debug, Deserialize)]
enum PersistedOperand {
Column(String),
Value(Value),
Expr(PersistedPartitionExpr),
/// `iceberg_hash(column)`.
IcebergHash(String),
}
/// Mirrors `partition::expr::RestrictedOp`.
#[derive(Debug, Deserialize)]
enum PersistedOp {
Eq,
@@ -102,10 +78,7 @@ enum PersistedOp {
Or,
}
/// Returns the columns whose data type differs between `old_schema` and `new_schema`,
/// mapped to (old type, new type). Columns only present in one of the schemas (added or
/// dropped) are not included; dropping a column referenced by a partition expression is
/// already rejected by the table metadata validation.
/// Returns the columns present in both schemas whose type differs, as (old, new).
pub(crate) fn changed_column_types(
old_schema: &Schema,
new_schema: &Schema,
@@ -126,12 +99,10 @@ pub(crate) fn changed_column_types(
changed
}
/// Validates that the persisted partition bounds of every region are still representable
/// with the types of `changed_columns`. Returns an error naming the first offending
/// region without modifying anything.
/// Returns an error naming the first region whose partition expression can't take
/// the type changes in `changed_columns`.
///
/// Public so that round-trip tests in the `partition` crate can run real serialized
/// partition expressions through the same mirror used here.
/// Public so that tests in `partition` can run real serialized expressions through it.
pub fn validate_region_partition_bounds(
changed_columns: &HashMap<String, (ConcreteDataType, ConcreteDataType)>,
region_routes: &[RegionRoute],
@@ -149,8 +120,7 @@ pub fn validate_region_partition_bounds(
let bound: PersistedPartitionBound = match serde_json::from_str(&expr_json) {
Ok(bound) => bound,
Err(err) => {
// A malformed persisted expression already fails every scan of this
// region. It is unrelated to the schema change, so don't block the ALTER.
// Such a region already fails every scan, regardless of this ALTER.
warn!(
"Skipping partition expression validation for region {}, \
the persisted expression is not parseable: {}, expr: {}",
@@ -159,9 +129,8 @@ pub fn validate_region_partition_bounds(
continue;
}
};
// Deprecated bound formats are not used to build partition predicates.
let PersistedPartitionBound::Expr(expr) = bound else {
// Deprecated bound formats don't carry expressions and are not used
// when the region builds its partition predicate.
continue;
};
@@ -187,7 +156,6 @@ fn validate_expr(
| PersistedOp::LtEq
| PersistedOp::Gt
| PersistedOp::GtEq => {
// Comparison between a column and a bound literal in either order.
validate_comparison(&expr.lhs, &expr.rhs, changed_columns, region_id)?;
validate_comparison(&expr.rhs, &expr.lhs, changed_columns, region_id)?;
}
@@ -203,7 +171,9 @@ fn validate_operand(
) -> Result<()> {
match operand {
PersistedOperand::Expr(expr) => validate_expr(expr, changed_columns, region_id),
PersistedOperand::Column(_) | PersistedOperand::Value(_) => Ok(()),
PersistedOperand::Column(_)
| PersistedOperand::Value(_)
| PersistedOperand::IcebergHash(_) => Ok(()),
}
}
@@ -213,83 +183,56 @@ fn validate_comparison(
changed_columns: &HashMap<String, (ConcreteDataType, ConcreteDataType)>,
region_id: RegionId,
) -> Result<()> {
let (PersistedOperand::Column(column), PersistedOperand::Value(value)) = (operand, other)
else {
return Ok(());
let (column, value) = match (operand, other) {
(PersistedOperand::Column(column), PersistedOperand::Value(value)) => (column, value),
// The hash depends on the column type, so any change reroutes rows. Table
// metadata rejects such changes first; this covers stale partition columns.
(PersistedOperand::IcebergHash(column), PersistedOperand::Value(_)) => {
return match changed_columns.get(column) {
Some((_, target_type)) => incompatible(region_id, column, target_type),
None => Ok(()),
};
}
_ => return Ok(()),
};
let Some((old_type, target_type)) = changed_columns.get(column) else {
return Ok(());
};
// Values not representable as a partition bound (lists, structs, json) are
// rejected by the partition parser and cannot appear in a persisted expression.
let Some(scalar) = value_to_scalar_value(value) else {
return Ok(());
};
// The write path (`MultiDimPartitionRule::find_region`) compares the persisted
// bound against row values directly: `Value::cmp` panics on mismatched variants
// and equality misroutes silently. Only a timestamp unit widening keeps the
// bound and the rows comparable (`Timestamp` comparisons convert units first),
// so every other type change of a referenced column is rejected here.
if !old_type.is_timestamp_unit_widening_to(target_type) {
warn!(
"Persisted partition bound would be compared against values of a different \
type after the ALTER, region_id: {}, column: '{}', target_type: {}",
region_id, column, target_type
);
return PartitionExprIncompatibleSnafu {
region_id,
column: column.clone(),
target_type: target_type.to_string(),
}
.fail();
return incompatible(region_id, column, target_type);
}
// For the widening, mirror `ExprSimplifier::coerce` in `PredicateGroup::new`:
// the region casts both the column and the persisted bound to the common
// comparison type. The bound is not cast to the column type itself, so the
// widening is accepted as long as the bound survives the common type; a bound
// that cannot be cast to it — or a pair of types without a common comparison
// type — fails every scan.
let target_arrow_type = target_type.as_arrow_type();
let Some(common_type) = comparison_coercion(&target_arrow_type, &scalar.data_type()) else {
warn!(
"Persisted partition bound has no common comparison type with the altered \
column type, region_id: {}, column: '{}', target_type: {}",
region_id, column, target_type
);
return PartitionExprIncompatibleSnafu {
region_id,
column: column.clone(),
target_type: target_type.to_string(),
}
.fail();
// Mirrors the coercion `PredicateGroup::new` applies on scan: the bound must
// cast to the common comparison type.
let Some(common_type) = comparison_coercion(&target_type.as_arrow_type(), &scalar.data_type())
else {
return incompatible(region_id, column, target_type);
};
if scalar.data_type() != common_type
&& let Err(error) = scalar.cast_to(&common_type)
{
warn!(
error;
"Persisted partition bound is not representable with the altered column type, \
region_id: {}, column: '{}', target_type: {}",
region_id,
column,
target_type
);
return PartitionExprIncompatibleSnafu {
region_id,
column: column.clone(),
target_type: target_type.to_string(),
}
.fail();
warn!(error; "Partition bound of region {} can't cast to {}", region_id, common_type);
return incompatible(region_id, column, target_type);
}
Ok(())
}
/// Converts a persisted bound value to a [`ScalarValue`], mirroring
/// `partition::expr::Operand::try_as_logical_expr`. Returns `None` for values that
/// cannot appear in a persisted partition expression.
fn incompatible(region_id: RegionId, column: &str, target_type: &ConcreteDataType) -> Result<()> {
PartitionExprIncompatibleSnafu {
region_id,
column,
target_type: target_type.to_string(),
}
.fail()
}
/// Converts a bound the same way as `partition::expr::Operand::try_as_logical_expr`.
/// Returns `None` for values that can't be partition bounds.
fn value_to_scalar_value(value: &Value) -> Option<ScalarValue> {
Some(match value {
Value::Boolean(v) => ScalarValue::Boolean(Some(*v)),
@@ -508,10 +451,8 @@ mod tests {
#[test]
fn test_field_narrowing_is_rejected_for_persisted_bounds() {
// BIGINT bound outside the INT range with the column narrowed to INT: the
// scan would coerce both sides to Int64, but the region's write path
// compares the persisted Int64 bound against Int32 rows directly and
// panics, so the ALTER must be rejected.
// Scans would coerce both sides, but row routing compares Int64 bounds with
// Int32 rows.
let bound =
r#"{"Expr":{"lhs":{"Column":"v"},"op":"Lt","rhs":{"Value":{"Int64":3000000000}}}}"#;
let err = validate_region_partition_bounds(
@@ -528,9 +469,7 @@ mod tests {
#[test]
fn test_non_timestamp_widening_is_rejected_for_persisted_bounds() {
// Widening INT to BIGINT still changes the variant rows carry: the
// region's write path compares the persisted Int32 bound against Int64
// rows directly and panics, so the ALTER must be rejected as well.
// Row routing would compare Int32 bounds with Int64 rows.
let bound = r#"{"Expr":{"lhs":{"Column":"v"},"op":"Lt","rhs":{"Value":{"Int32":10}}}}"#;
let err = validate_region_partition_bounds(
&changed(
@@ -546,9 +485,7 @@ mod tests {
#[test]
fn test_bound_without_common_comparison_type_is_rejected() {
// A boolean bound has no comparison coercion with a timestamp column, so
// even a widening whose bounds are normally representable fails every
// predicate build; the ALTER must be rejected.
// No common comparison type between a boolean and a timestamp.
let bound =
r#"{"Expr":{"lhs":{"Column":"ts"},"op":"Lt","rhs":{"Value":{"Boolean":true}}}}"#;
let err = validate_region_partition_bounds(
src/mito2/src/engine/partition_filter_test.rs
+115 -1
View File
@@ -16,8 +16,10 @@
use api::v1::Rows;
use common_recordbatch::RecordBatches;
use datatypes::arrow::array::AsArray;
use datatypes::value::Value;
use partition::expr::col;
use partition::expr::{Operand, col};
use partition::function::IcebergHash;
use store_api::region_engine::RegionEngine;
use store_api::region_request::{
EnterStagingRequest, RegionFlushRequest, RegionRequest, StagingPartitionDirective,
@@ -168,3 +170,115 @@ async fn test_partition_filter_basic_with_format(flat_format: bool) {
flat_format
);
}
#[tokio::test]
async fn test_partition_filter_iceberg_hash() {
test_partition_filter_iceberg_hash_with_format(false).await;
test_partition_filter_iceberg_hash_with_format(true).await;
}
/// Narrows a region to `iceberg_hash(tag_0) < 0` and checks rows flushed under
/// the old partition expr are filtered by the hash.
async fn test_partition_filter_iceberg_hash_with_format(flat_format: bool) {
let mut env = TestEnv::new().await;
let engine = env
.create_engine(MitoConfig {
default_flat_format: flat_format,
..Default::default()
})
.await;
let region_id = RegionId::new(1024, 0);
let request = CreateRequestBuilder::new()
.partition_expr_json(Some(range_expr_string("field_0", 0., 99.)))
.build();
let column_schemas = rows_schema(&request);
engine
.handle_request(region_id, RegionRequest::Create(request))
.await
.unwrap();
let flush = || async {
engine
.handle_request(
region_id,
RegionRequest::Flush(RegionFlushRequest::default()),
)
.await
.unwrap();
};
put_rows(
&engine,
region_id,
Rows {
schema: column_schemas.clone(),
rows: build_rows(0, 20),
},
)
.await;
flush().await;
let in_partition = |i: &usize| {
IcebergHash::evaluate_value(&Value::String(i.to_string().as_str().into())).unwrap()
< Value::Int32(0)
};
let hash_expr = Operand::IcebergHash(IcebergHash::new("tag_0")).lt(Value::Int32(0));
engine
.handle_request(
region_id,
RegionRequest::EnterStaging(EnterStagingRequest {
partition_directive: StagingPartitionDirective::UpdatePartitionExpr(
hash_expr.as_json_str().unwrap(),
),
}),
)
.await
.unwrap();
put_rows(
&engine,
region_id,
Rows {
schema: column_schemas,
rows: (20..40)
.filter(in_partition)
.flat_map(|i| build_rows(i, i + 1))
.collect(),
},
)
.await;
flush().await;
use store_api::region_engine::SettableRegionRoleState;
engine
.set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader)
.await
.unwrap();
let request = ScanRequest {
projection: Some(vec![0]),
..Default::default()
};
let scanner = engine.scanner(region_id, request).await.unwrap();
let batches = RecordBatches::try_collect(scanner.scan().await.unwrap())
.await
.unwrap();
let mut actual = vec![];
for batch in batches.iter() {
let tags = datatypes::arrow::compute::cast(
batch.column(0),
&datatypes::arrow::datatypes::DataType::Utf8,
)
.unwrap();
actual.extend(
tags.as_string::<i32>()
.iter()
.map(|v| v.unwrap().to_string()),
);
}
actual.sort();
let mut expected = (0..40)
.filter(in_partition)
.map(|i| i.to_string())
.collect::<Vec<_>>();
expected.sort();
// Rows flushed before staging fall on both sides of the hash bound.
assert!((0..20).any(|i| in_partition(&i)) && !(0..20).all(|i| in_partition(&i)));
assert_eq!(actual, expected, "flat_format: {flat_format}");
}
src/operator/src/statement/ddl.rs
+173 -17
View File
@@ -72,6 +72,7 @@ use datatypes::value::Value;
use datatypes::vectors::{StringVector, VectorRef};
use humantime::parse_duration;
use partition::expr::{Operand, PartitionExpr, RestrictedOp};
use partition::function::IcebergHash;
use partition::multi_dim::MultiDimPartitionRule;
use query::parser::QueryStatement;
use query::plan::extract_and_rewrite_full_table_names;
@@ -92,7 +93,10 @@ use sql::statements::create::{
CreateExternalTable, CreateFlow, CreateTable, CreateTableLike, CreateView, Partitions,
};
use sql::statements::statement::Statement;
use sqlparser::ast::{Expr, Ident, UnaryOperator, Value as ParserValue};
use sqlparser::ast::{
Expr, Function, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, UnaryOperator,
Value as ParserValue,
};
use store_api::metric_engine_consts::{
LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
};
@@ -2919,6 +2923,15 @@ fn ensure_partition_operand_columns_in_target(
ensure_partition_operand_columns_in_target(&expr.lhs, target_partition_columns)?;
ensure_partition_operand_columns_in_target(&expr.rhs, target_partition_columns)?;
}
Operand::IcebergHash(hash) => ensure!(
target_partition_columns.contains(&hash.column().to_string()),
InvalidPartitionRuleSnafu {
reason: format!(
"partition expression references column '{}' that is not in target partition columns",
hash.column()
)
}
),
Operand::Value(_) => {}
}
@@ -2970,6 +2983,22 @@ fn convert_one_expr(
let value = convert_value(&v.value, data_type, timezone, Some(*unary_op))?;
(Operand::Value(value), op, Operand::Column(column_name))
}
// func, val
(Expr::Function(function), literal)
if !matches!(op, RestrictedOp::And | RestrictedOp::Or) =>
{
let hash = convert_iceberg_hash(function, column_name_and_type)?;
let value = convert_literal(literal, IcebergHash::output_type(), timezone)?;
(Operand::IcebergHash(hash), op, Operand::Value(value))
}
// val, func
(literal, Expr::Function(function))
if !matches!(op, RestrictedOp::And | RestrictedOp::Or) =>
{
let hash = convert_iceberg_hash(function, column_name_and_type)?;
let value = convert_literal(literal, IcebergHash::output_type(), timezone)?;
(Operand::Value(value), op, Operand::IcebergHash(hash))
}
(Expr::BinaryOp { .. }, Expr::BinaryOp { .. }) => {
// sub-expr must against another sub-expr
let lhs = convert_one_expr(left, column_name_and_type, timezone)?;
@@ -2987,6 +3016,70 @@ fn convert_one_expr(
Ok(PartitionExpr::new(lhs, op, rhs))
}
/// Converts an `iceberg_hash` call, which takes exactly one partition column.
fn convert_iceberg_hash(
function: &Function,
column_name_and_type: &HashMap<&String, ConcreteDataType>,
) -> Result<IcebergHash> {
let invalid = || {
InvalidPartitionRuleSnafu {
reason: format!("unsupported partition function call {function}"),
}
.build()
};
let FunctionArguments::List(args) = &function.args else {
return Err(invalid());
};
let [FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Identifier(ident)))] =
args.args.as_slice()
else {
return Err(invalid());
};
if function.filter.is_some()
|| function.over.is_some()
|| function.null_treatment.is_some()
|| !function.within_group.is_empty()
|| !matches!(function.parameters, FunctionArguments::None)
|| args.duplicate_treatment.is_some()
|| !args.clauses.is_empty()
{
return Err(invalid());
}
ensure!(
function
.name
.to_string()
.eq_ignore_ascii_case(IcebergHash::NAME),
InvalidPartitionRuleSnafu {
reason: format!("unsupported partition function {}", function.name),
}
);
let (column_name, data_type) = convert_identifier(ident, column_name_and_type)?;
ensure!(
IcebergHash::accepts(&data_type),
InvalidPartitionRuleSnafu {
reason: format!(
"partition function {} does not support column type {data_type}",
IcebergHash::NAME
),
}
);
Ok(IcebergHash::new(column_name))
}
fn convert_literal(expr: &Expr, data_type: ConcreteDataType, timezone: &Timezone) -> Result<Value> {
match expr {
Expr::Value(value) => convert_value(&value.value, data_type, timezone, None),
Expr::UnaryOp { op, expr } if let Expr::Value(value) = &**expr => {
convert_value(&value.value, data_type, timezone, Some(*op))
}
_ => InvalidPartitionRuleSnafu {
reason: format!("expect a literal in partition expr, got {expr}"),
}
.fail(),
}
}
fn convert_identifier(
ident: &Ident,
column_name_and_type: &HashMap<&String, ConcreteDataType>,
@@ -3748,6 +3841,66 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;";
assert!(!NAME_PATTERN_REG.is_match("#"));
}
#[test]
fn test_convert_iceberg_hash() {
let names = ["host", "count", "device", "measure", "host-name"].map(String::from);
let column_name_and_type = HashMap::from([
(&names[0], ConcreteDataType::string_datatype()),
(&names[1], ConcreteDataType::int16_datatype()),
(&names[2], ConcreteDataType::uint64_datatype()),
(&names[3], ConcreteDataType::float64_datatype()),
(&names[4], ConcreteDataType::string_datatype()),
]);
let timezone = Timezone::from_tz_string("UTC").unwrap();
let dialect = GreptimeDbDialect {};
let parse = |sql: &str| {
Parser::new(&dialect)
.try_with_sql(sql)
.unwrap()
.parse_expr()
.unwrap()
};
for sql in [
"iceberg_hash(host) < 0",
"ICEBERG_HASH(device) >= -1073741824",
"0 > iceberg_hash(count)",
"iceberg_hash(host) < 100 AND host < 'm'",
r#"iceberg_hash("host-name") < 0 AND "host-name" < 'm'"#,
] {
let expr = convert_one_expr(&parse(sql), &column_name_and_type, &timezone).unwrap();
// SHOW CREATE TABLE prints the rule back as SQL, quoting identifiers.
let mut partitions = Partitions {
column_list: vec![],
exprs: vec![expr.to_parser_expr()],
};
partitions.set_quote('"');
let restored = parse(&partitions.exprs[0].to_string());
assert_eq!(
convert_one_expr(&restored, &column_name_and_type, &timezone).unwrap(),
expr,
"{sql}"
);
}
for sql in [
"iceberg_hash(measure) < 0",
"iceberg_hash(host, count) < 0",
"iceberg_hash('a') < 0",
"iceberg_hash(missing) < 0",
"iceberg_hash(DISTINCT host) < 0",
"md5(host) < 0",
"iceberg_hash(host) < host",
"iceberg_hash(host) < 'abc'",
"iceberg_hash(host) < 2147483648",
] {
assert!(
convert_one_expr(&parse(sql), &column_name_and_type, &timezone).is_err(),
"{sql}"
);
}
}
#[test]
fn test_partition_expr_equivalence_with_swapped_operands() {
let column_name = "device_id".to_string();
@@ -3818,25 +3971,28 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;";
(&area, ConcreteDataType::string_datatype()),
]);
let dialect = GreptimeDbDialect {};
let mut parser = Parser::new(&dialect)
.try_with_sql("device_id >= 100")
.unwrap();
let remaining_old_expr = convert_one_expr(
&parser.parse_expr().unwrap(),
&column_name_and_type,
&timezone,
)
.unwrap();
let target_partition_columns = HashSet::from([&area]);
for sql in ["device_id >= 100", "iceberg_hash(device_id) >= 0"] {
let mut parser = Parser::new(&dialect).try_with_sql(sql).unwrap();
let remaining_old_expr = convert_one_expr(
&parser.parse_expr().unwrap(),
&column_name_and_type,
&timezone,
)
.unwrap();
let err = ensure_partition_expr_columns_in_target(
&[remaining_old_expr],
&target_partition_columns,
)
.unwrap_err();
let err = ensure_partition_expr_columns_in_target(
&[remaining_old_expr],
&target_partition_columns,
)
.unwrap_err();
assert!(err.to_string().contains("device_id"));
assert!(err.to_string().contains("target partition columns"));
assert!(err.to_string().contains("device_id"), "{sql}");
assert!(
err.to_string().contains("target partition columns"),
"{sql}"
);
}
}
#[test]
src/partition/benches/bench_split_record_batch.rs
+60 -2
View File
@@ -22,7 +22,8 @@ use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::value::Value;
use partition::PartitionRule;
use partition::expr::{Operand, col};
use partition::expr::{Operand, PartitionExpr, col};
use partition::function::IcebergHash;
use partition::multi_dim::MultiDimPartitionRule;
use rand::Rng;
use store_api::storage::RegionNumber;
@@ -219,9 +220,66 @@ fn bench_split_record_batch_vs_row(c: &mut Criterion) {
group.finish();
}
/// Splits `dim` into `num_partitions` ranges, `bounds` being the inner boundaries.
fn range_rule(
column: &str,
dim: Operand,
bounds: impl IntoIterator<Item = i32>,
) -> MultiDimPartitionRule {
let bounds = bounds.into_iter().map(Value::Int32).collect::<Vec<_>>();
let mut exprs: Vec<PartitionExpr> = vec![dim.clone().lt(bounds[0].clone())];
for pair in bounds.windows(2) {
exprs.push(
dim.clone()
.gt_eq(pair[0].clone())
.and(dim.clone().lt(pair[1].clone())),
);
}
exprs.push(dim.gt_eq(bounds[bounds.len() - 1].clone()));
let regions = (0..exprs.len()).map(|v| v as u32).collect();
MultiDimPartitionRule::try_new(vec![column.to_string()], regions, exprs, true).unwrap()
}
fn hash_rule(column: &str, num_partitions: i64) -> MultiDimPartitionRule {
let width = (1i64 << 32) / num_partitions;
range_rule(
column,
Operand::IcebergHash(IcebergHash::new(column)),
(1..num_partitions).map(|k| (i32::MIN as i64 + k * width) as i32),
)
}
fn bench_split_record_batch_hash(c: &mut Criterion) {
let mut group = c.benchmark_group("split_record_batch_hash");
for num_partitions in [4, 16, 64] {
for num_rows in [1000, 100000] {
let batch = create_test_batch(num_rows);
// Baseline: the same number of plain column ranges over a0 in [0, 100).
let baseline = range_rule(
"a0",
col("a0"),
(1..num_partitions).map(|k| (k * 100 / num_partitions) as i32),
);
group.bench_function(format!("column_range_{num_partitions}_{num_rows}"), |b| {
b.iter(|| black_box(baseline.split_record_batch(black_box(&batch))).unwrap())
});
for (key, column) in [("int", "a0"), ("string", "a1")] {
let rule = hash_rule(column, num_partitions);
group.bench_function(format!("hash_{key}_{num_partitions}_{num_rows}"), |b| {
b.iter(|| black_box(rule.split_record_batch(black_box(&batch))).unwrap())
});
}
}
}
group.finish();
}
criterion_group!(
benches,
bench_split_record_batch_naive_vs_optimized,
bench_split_record_batch_vs_row
bench_split_record_batch_vs_row,
bench_split_record_batch_hash
);
criterion_main!(benches);
src/partition/src/checker.rs
+58 -13
View File
@@ -19,7 +19,7 @@ use datatypes::arrow::array::{BooleanArray, Float64Array, Float64Builder, Record
use datatypes::arrow::datatypes::{DataType, Field, Schema};
use datatypes::value::OrderedF64;
use crate::collider::{CHECK_STEP, Collider, NORMALIZE_STEP};
use crate::collider::{CHECK_STEP, Collider, NORMALIZE_STEP, PartitionDim};
use crate::error::{
CheckpointNotCoveredSnafu, CheckpointOverlappedSnafu, DuplicateExprSnafu, Result,
};
@@ -29,12 +29,26 @@ use crate::multi_dim::MultiDimPartitionRule;
pub struct PartitionChecker<'a> {
rule: &'a MultiDimPartitionRule,
collider: Collider<'a>,
/// Dimensions in the matrix. A dimension's field name in the matrix is its
/// index here, since a column may be named like a function call.
dims: Vec<PartitionDim>,
}
impl<'a> PartitionChecker<'a> {
pub fn try_new(rule: &'a MultiDimPartitionRule) -> Result<Self> {
let collider = Collider::new(rule.exprs())?;
Ok(Self { rule, collider })
let mut dims = collider
.normalized_values
.iter()
.filter(|(_, values)| !values.is_empty())
.map(|(dim, _)| dim.clone())
.collect::<Vec<_>>();
dims.sort_unstable();
Ok(Self {
rule,
collider,
dims,
})
}
pub fn check(&self) -> Result<()> {
@@ -76,19 +90,19 @@ impl<'a> PartitionChecker<'a> {
// TODO(ruihang): merge atomic exprs to improve checker's performance
// matrix test
let field_names = (0..self.dims.len())
.map(|index| index.to_string())
.collect::<Vec<_>>();
let mut matrix_foundation = HashMap::new();
for (col, values) in self.collider.normalized_values.iter() {
if values.is_empty() {
continue;
}
for (dim, field_name) in self.dims.iter().zip(&field_names) {
let values = &self.collider.normalized_values[dim];
let mut cornerstones = Vec::with_capacity(values.len() * 2 + 1);
cornerstones.push(values[0].1 - CHECK_STEP);
for value in values {
cornerstones.push(value.1);
cornerstones.push(value.1 + CHECK_STEP);
}
matrix_foundation.insert(col.as_str(), cornerstones);
matrix_foundation.insert(field_name.as_str(), cornerstones);
}
// If there are no values, the rule is empty and valid.
@@ -104,7 +118,12 @@ impl<'a> PartitionChecker<'a> {
.collider
.atomic_exprs
.iter()
.map(|expr| expr.to_physical_expr(matrix_generator.schema()))
.map(|expr| {
expr.to_physical_expr(matrix_generator.schema(), |dim| {
// Every nucleon's dimension has a value, so it is in `dims`.
self.dims.binary_search(dim).unwrap().to_string()
})
})
.collect::<Vec<_>>();
for batch in matrix_generator {
results.clear();
@@ -158,14 +177,20 @@ impl<'a> PartitionChecker<'a> {
let mut check_point = String::new();
let schema = batch.schema();
for (col_index, normalized_value) in normalized_row.iter().enumerate() {
let col_name = schema.field(col_index).name();
let dim = schema
.field(col_index)
.name()
.parse::<usize>()
.ok()
.and_then(|index| self.dims.get(index));
let col_name = dim.map(ToString::to_string).unwrap_or_default();
if col_index > 0 {
check_point.push_str(", ");
}
// Check if point is on NORMALIZE_STEP or between steps
if let Some(values) = self.collider.normalized_values.get(col_name) {
if let Some(values) = dim.and_then(|dim| self.collider.normalized_values.get(dim)) {
let normalize_step = NORMALIZE_STEP.0;
// Check if the normalized value is on a NORMALIZE_STEP boundary
@@ -350,6 +375,26 @@ mod tests {
use crate::expr::col;
use crate::multi_dim::MultiDimPartitionRule;
#[test]
fn test_hash_and_same_named_column_are_distinct() {
use crate::expr::Operand;
use crate::function::IcebergHash;
let hash = IcebergHash::new("host");
// A column can be named like the hash call, but it is another dimension, so these
// two rules don't cover each other's complement.
let result = MultiDimPartitionRule::try_new(
vec!["host".to_string(), hash.to_string()],
vec![1, 2],
vec![
Operand::IcebergHash(hash.clone()).lt(Value::Int32(0)),
col(hash.to_string()).gt_eq(Value::Int32(0)),
],
true,
);
assert!(result.is_err());
}
#[test]
fn test_matrix_generator_single_column() {
let mut matrix_foundation = HashMap::new();
@@ -582,8 +627,8 @@ mod tests {
let checker = PartitionChecker::try_new(&rule).unwrap();
let schema = Arc::new(Schema::new(vec![
Field::new("host", DataType::Float64, false),
Field::new("value", DataType::Float64, false),
Field::new("0", DataType::Float64, false),
Field::new("1", DataType::Float64, false),
]));
let host_array = Float64Array::from(vec![-0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]);
let value_array = Float64Array::from(vec![-0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]);
src/partition/src/collider.rs
+93 -40
View File
@@ -26,7 +26,7 @@
//! force inside nucleons.
use std::collections::HashMap;
use std::fmt::Debug;
use std::fmt::{Debug, Display, Formatter};
use std::sync::Arc;
use datafusion_expr::Operator;
@@ -38,11 +38,45 @@ use datatypes::value::{OrderedF64, OrderedFloat, Value};
use crate::error;
use crate::error::Result;
use crate::expr::{Operand, PartitionExpr, RestrictedOp};
use crate::function::IcebergHash;
const ZERO: OrderedF64 = OrderedFloat(0.0f64);
pub(crate) const NORMALIZE_STEP: OrderedF64 = OrderedFloat(1.0f64);
pub(crate) const CHECK_STEP: OrderedF64 = OrderedFloat(0.5f64);
/// A dimension of the partition space: a column, or the hash of a column.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum PartitionDim {
Column(String),
IcebergHash(IcebergHash),
}
impl PartitionDim {
fn from_operand(operand: &Operand) -> Option<Self> {
match operand {
Operand::Column(column) => Some(Self::Column(column.clone())),
Operand::IcebergHash(hash) => Some(Self::IcebergHash(hash.clone())),
Operand::Value(_) | Operand::Expr(_) => None,
}
}
pub fn to_operand(&self) -> Operand {
match self {
Self::Column(column) => Operand::Column(column.clone()),
Self::IcebergHash(hash) => Operand::IcebergHash(hash.clone()),
}
}
}
impl Display for PartitionDim {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Column(column) => write!(f, "{column}"),
Self::IcebergHash(hash) => write!(f, "{hash}"),
}
}
}
/// Represents an "atomic" Expression, which isn't composed (OR-ed) of other expressions.
#[allow(unused)]
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -55,10 +89,15 @@ pub struct AtomicExpr {
}
impl AtomicExpr {
pub fn to_physical_expr(&self, schema: &Schema) -> Arc<dyn PhysicalExpr> {
/// Builds the conjunction over `schema`, where `field_name` maps each dimension to its field.
pub fn to_physical_expr(
&self,
schema: &Schema,
field_name: impl Fn(&PartitionDim) -> String,
) -> Arc<dyn PhysicalExpr> {
let mut exprs = Vec::with_capacity(self.nucleons.len());
for nucleon in &self.nucleons {
exprs.push(nucleon.to_physical_expr(schema));
exprs.push(nucleon.to_physical_expr(schema, &field_name(&nucleon.column)));
}
let result: Arc<dyn PhysicalExpr> = exprs
.into_iter()
@@ -79,23 +118,23 @@ impl PartialOrd for AtomicExpr {
/// This struct is used to compose [`AtomicExpr`], hence "nucleon".
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct NucleonExpr {
column: String,
column: PartitionDim,
op: GluonOp,
/// Normalized [`Value`].
value: OrderedF64,
}
impl NucleonExpr {
pub fn to_physical_expr(&self, schema: &Schema) -> Arc<dyn PhysicalExpr> {
pub fn to_physical_expr(&self, schema: &Schema, field_name: &str) -> Arc<dyn PhysicalExpr> {
Arc::new(BinaryExpr::new(
col(&self.column, schema).unwrap(),
col(field_name, schema).unwrap(),
self.op.to_operator(),
lit(*self.value.as_ref()),
))
}
/// Get the column name
pub fn column(&self) -> &str {
/// Get the dimension
pub fn column(&self) -> &PartitionDim {
&self.column
}
@@ -109,12 +148,8 @@ impl NucleonExpr {
&self.op
}
pub fn new(column: impl Into<String>, op: GluonOp, value: OrderedF64) -> Self {
Self {
column: column.into(),
op,
value,
}
pub fn new(column: PartitionDim, op: GluonOp, value: OrderedF64) -> Self {
Self { column, op, value }
}
}
@@ -153,22 +188,22 @@ pub struct Collider<'a> {
source_exprs: &'a [PartitionExpr],
pub atomic_exprs: Vec<AtomicExpr>,
/// A map of column name to a list of `(value, normalized value)` pairs.
/// A map of dimension to a list of `(value, normalized value)` pairs.
///
/// The normalized value is used for comparison. The normalization process keeps the order of the values.
pub normalized_values: HashMap<String, Vec<(Value, OrderedF64)>>,
pub normalized_values: HashMap<PartitionDim, Vec<(Value, OrderedF64)>>,
}
impl<'a> Collider<'a> {
pub fn new(source_exprs: &'a [PartitionExpr]) -> Result<Self> {
// first walk to collect all values
let mut values: HashMap<String, Vec<Value>> = HashMap::new();
let mut values: HashMap<PartitionDim, Vec<Value>> = HashMap::new();
for expr in source_exprs {
Self::collect_column_values_from_expr(expr, &mut values)?;
}
// normalize values, assumes all values on a column are the same type
let mut normalized_values: HashMap<String, HashMap<Value, OrderedF64>> =
let mut normalized_values: HashMap<PartitionDim, HashMap<Value, OrderedF64>> =
HashMap::with_capacity(values.len());
for (column, mut column_values) in values {
column_values.sort_unstable();
@@ -215,13 +250,14 @@ impl<'a> Collider<'a> {
/// Helper to collect values with their associated columns from an expression
fn collect_column_values_from_expr(
expr: &PartitionExpr,
values: &mut HashMap<String, Vec<Value>>,
values: &mut HashMap<PartitionDim, Vec<Value>>,
) -> Result<()> {
// Handle binary operations between column and value
match (&*expr.lhs, &*expr.rhs) {
(Operand::Column(col), Operand::Value(val))
| (Operand::Value(val), Operand::Column(col)) => {
values.entry(col.clone()).or_default().push(val.clone());
(dim, Operand::Value(val)) | (Operand::Value(val), dim)
if let Some(dim) = PartitionDim::from_operand(dim) =>
{
values.entry(dim).or_default().push(val.clone());
Ok(())
}
(Operand::Expr(left_expr), Operand::Expr(right_expr)) => {
@@ -241,7 +277,7 @@ impl<'a> Collider<'a> {
fn collide_expr(
expr: &PartitionExpr,
index: usize,
normalized_values: &HashMap<String, HashMap<Value, OrderedF64>>,
normalized_values: &HashMap<PartitionDim, HashMap<Value, OrderedF64>>,
result: &mut Vec<AtomicExpr>,
) -> Result<()> {
match expr.op {
@@ -300,7 +336,7 @@ impl<'a> Collider<'a> {
fn collect_nucleons_from_expr(
expr: &PartitionExpr,
nucleons: &mut Vec<NucleonExpr>,
normalized_values: &HashMap<String, HashMap<Value, OrderedF64>>,
normalized_values: &HashMap<PartitionDim, HashMap<Value, OrderedF64>>,
) -> Result<()> {
match expr.op {
RestrictedOp::And => {
@@ -325,7 +361,7 @@ impl<'a> Collider<'a> {
fn collect_nucleons_from_operand(
operand: &Operand,
nucleons: &mut Vec<NucleonExpr>,
normalized_values: &HashMap<String, HashMap<Value, OrderedF64>>,
normalized_values: &HashMap<PartitionDim, HashMap<Value, OrderedF64>>,
) -> Result<()> {
match operand {
Operand::Expr(expr) => {
@@ -346,7 +382,7 @@ impl<'a> Collider<'a> {
lhs: &Operand,
op: &RestrictedOp,
rhs: &Operand,
normalized_values: &HashMap<String, HashMap<Value, OrderedF64>>,
normalized_values: &HashMap<PartitionDim, HashMap<Value, OrderedF64>>,
) -> Result<NucleonExpr> {
let gluon_op = match op {
RestrictedOp::Eq => GluonOp::Eq,
@@ -365,19 +401,19 @@ impl<'a> Collider<'a> {
};
match (lhs, rhs) {
(Operand::Column(col), Operand::Value(val)) => {
if let Some(column_values) = normalized_values.get(col)
(dim, Operand::Value(val)) if let Some(dim) = PartitionDim::from_operand(dim) => {
if let Some(column_values) = normalized_values.get(&dim)
&& let Some(&normalized_val) = column_values.get(val)
{
return Ok(NucleonExpr {
column: col.clone(),
column: dim,
op: gluon_op,
value: normalized_val,
});
}
}
(Operand::Value(val), Operand::Column(col)) => {
if let Some(column_values) = normalized_values.get(col)
(Operand::Value(val), dim) if let Some(dim) = PartitionDim::from_operand(dim) => {
if let Some(column_values) = normalized_values.get(&dim)
&& let Some(&normalized_val) = column_values.get(val)
{
// Flip the operation for value op column
@@ -389,7 +425,7 @@ impl<'a> Collider<'a> {
op => op, // Eq and NotEq remain the same
};
return Ok(NucleonExpr {
column: col.clone(),
column: dim,
op: flipped_op,
value: normalized_val,
});
@@ -436,7 +472,7 @@ mod test {
assert_eq!(collider.normalized_values.len(), 4);
// Check age column - should have 2 unique values (25, 30)
let age_values = &collider.normalized_values["age"];
let age_values = &collider.normalized_values[&PartitionDim::Column("age".to_string())];
assert_eq!(age_values.len(), 2);
assert_eq!(
age_values,
@@ -447,7 +483,7 @@ mod test {
);
// Check name column - should have 2 values
let name_values = &collider.normalized_values["name"];
let name_values = &collider.normalized_values[&PartitionDim::Column("name".to_string())];
assert_eq!(name_values.len(), 2);
assert_eq!(
name_values,
@@ -458,7 +494,8 @@ mod test {
);
// Check active column - should have 2 values
let active_values = &collider.normalized_values["active"];
let active_values =
&collider.normalized_values[&PartitionDim::Column("active".to_string())];
assert_eq!(active_values.len(), 2);
assert_eq!(
active_values,
@@ -469,7 +506,7 @@ mod test {
);
// Check score column - should have 2 values
let score_values = &collider.normalized_values["score"];
let score_values = &collider.normalized_values[&PartitionDim::Column("score".to_string())];
assert_eq!(score_values.len(), 2);
assert_eq!(
score_values,
@@ -665,10 +702,26 @@ mod test {
// Should have normalized values for all 4 columns
assert_eq!(collider.normalized_values.len(), 4);
assert!(collider.normalized_values.contains_key("name"));
assert!(collider.normalized_values.contains_key("age"));
assert!(collider.normalized_values.contains_key("status"));
assert!(collider.normalized_values.contains_key("score"));
assert!(
collider
.normalized_values
.contains_key(&PartitionDim::Column("name".to_string()))
);
assert!(
collider
.normalized_values
.contains_key(&PartitionDim::Column("age".to_string()))
);
assert!(
collider
.normalized_values
.contains_key(&PartitionDim::Column("status".to_string()))
);
assert!(
collider
.normalized_values
.contains_key(&PartitionDim::Column("score".to_string()))
);
}
#[test]
src/partition/src/expr.rs
+143 -62
View File
@@ -31,9 +31,13 @@ use datatypes::value::{
use serde::{Deserialize, Serialize};
use snafu::ResultExt;
use sql::statements::value_to_sql_value;
use sqlparser::ast::{BinaryOperator as ParserBinaryOperator, Expr as ParserExpr, Ident};
use sqlparser::ast::{
BinaryOperator as ParserBinaryOperator, Expr as ParserExpr, Function, FunctionArg,
FunctionArgExpr, FunctionArgumentList, FunctionArguments, Ident, ObjectName,
};
use crate::error;
use crate::function::IcebergHash;
use crate::partition::PartitionBound;
/// Struct for partition expression. This can be converted back to sqlparser's [Expr].
@@ -52,6 +56,7 @@ pub enum Operand {
Column(String),
Value(Value),
Expr(PartitionExpr),
IcebergHash(IcebergHash),
}
pub fn col(column_name: impl Into<String>) -> Operand {
@@ -67,7 +72,7 @@ impl From<Value> for Operand {
impl Operand {
pub fn try_as_logical_expr(&self) -> error::Result<Expr> {
match self {
Self::Column(c) => Ok(datafusion_expr::col(format!(r#""{}""#, c))),
Self::Column(c) => Ok(Expr::Column(datafusion_common::Column::from_name(c))),
Self::Value(v) => {
let scalar_value = match v {
Value::Boolean(v) => ScalarValue::Boolean(Some(*v)),
@@ -108,6 +113,38 @@ impl Operand {
Ok(datafusion_expr::lit(scalar_value))
}
Self::Expr(e) => e.try_as_logical_expr(),
Self::IcebergHash(f) => Ok(f.to_logical_expr()),
}
}
/// Whether this operand is a partition dimension that values are compared against.
fn is_dimension(&self) -> bool {
matches!(self, Self::Column(_) | Self::IcebergHash(_))
}
fn to_parser_expr(&self) -> ParserExpr {
// Safety: Partition rule won't contains unsupported value type.
// Otherwise it will be rejected by the parser.
match self {
Self::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())),
Self::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()),
Self::Expr(e) => e.to_parser_expr(),
Self::IcebergHash(f) => ParserExpr::Function(Function {
name: ObjectName::from(vec![Ident::new(IcebergHash::NAME)]),
uses_odbc_syntax: false,
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
duplicate_treatment: None,
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
ParserExpr::Identifier(Ident::new(f.column())),
))],
clauses: vec![],
}),
filter: None,
null_treatment: None,
over: None,
within_group: vec![],
}),
}
}
@@ -142,6 +179,7 @@ impl Display for Operand {
Self::Column(v) => write!(f, "{v}"),
Self::Value(v) => write!(f, "{v}"),
Self::Expr(v) => write!(f, "{v}"),
Self::IcebergHash(v) => write!(f, "{v}"),
}
}
}
@@ -239,7 +277,7 @@ impl PartitionExpr {
rhs: Box::new(rhs),
};
if matches!(&*expr.lhs, Operand::Value(_)) && matches!(&*expr.rhs, Operand::Column(_)) {
if matches!(&*expr.lhs, Operand::Value(_)) && expr.rhs.is_dimension() {
std::mem::swap(&mut expr.lhs, &mut expr.rhs);
expr.op = expr.op.invert_for_swap();
}
@@ -258,24 +296,10 @@ impl PartitionExpr {
///
/// [Expr]: ParserExpr
pub fn to_parser_expr(&self) -> ParserExpr {
// Safety: Partition rule won't contains unsupported value type.
// Otherwise it will be rejected by the parser.
let lhs = match &*self.lhs {
Operand::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())),
Operand::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()),
Operand::Expr(e) => e.to_parser_expr(),
};
let rhs = match &*self.rhs {
Operand::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())),
Operand::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()),
Operand::Expr(e) => e.to_parser_expr(),
};
ParserExpr::BinaryOp {
left: Box::new(lhs),
left: Box::new(self.lhs.to_parser_expr()),
op: self.op.to_parser_op(),
right: Box::new(rhs),
right: Box::new(self.rhs.to_parser_expr()),
}
}
@@ -284,10 +308,10 @@ impl PartitionExpr {
// Batch routing and region filters treat a NULL result as false, so only the
// comparisons that hold for NULL need an explicit `IS NULL`.
let comparison = match (self.lhs.as_ref(), self.rhs.as_ref()) {
(column @ Operand::Column(_), Operand::Value(value)) => {
(column, Operand::Value(value)) if column.is_dimension() => {
Some((column, value, self.op.clone()))
}
(Operand::Value(value), column @ Operand::Column(_)) => {
(Operand::Value(value), column) if column.is_dimension() => {
Some((column, value, self.op.invert_for_swap()))
}
_ => None,
@@ -355,32 +379,7 @@ impl PartitionExpr {
&self,
schema: &arrow::datatypes::SchemaRef,
) -> error::Result<Arc<dyn PhysicalExpr>> {
let df_schema = schema
.clone()
.to_dfschema_ref()
.context(error::ToDFSchemaSnafu)?;
let expr = self.try_as_logical_expr()?;
// Persisted partition bounds keep the unit (and type) they were created
// with, so after a schema change like widening the time index unit the
// expression may mix types with the columns it is evaluated against.
// Coerce it against the schema so comparisons share a common type, the
// same way `PredicateGroup` does for region partition filters.
let simplifier = ExprSimplifier::new(
SimplifyContext::builder()
.with_schema(df_schema.clone())
.build(),
);
let expr = simplifier
.coerce(expr, &df_schema)
.context(error::CoercePartitionExprSnafu)?;
let execution_props = &ExecutionProps::default();
create_physical_expr(
&expr,
&df_schema,
execution_props,
&PhysicalPlanningContext::default(),
)
.context(error::CreatePhysicalExprSnafu)
plan_physical_expr(self.try_as_logical_expr()?, schema)
}
pub fn and(self, rhs: PartitionExpr) -> PartitionExpr {
@@ -432,11 +431,45 @@ impl PartitionExpr {
Operand::Expr(e) => {
e.collect_column_names(columns);
}
Operand::IcebergHash(f) => {
columns.insert(f.column().to_string());
}
Operand::Value(_) => {}
}
}
}
/// Plans a logical partition expression for evaluation against `schema`.
pub(crate) fn plan_physical_expr(
expr: Expr,
schema: &arrow::datatypes::SchemaRef,
) -> error::Result<Arc<dyn PhysicalExpr>> {
let df_schema = schema
.clone()
.to_dfschema_ref()
.context(error::ToDFSchemaSnafu)?;
// Persisted partition bounds keep the unit (and type) they were created
// with, so after a schema change like widening the time index unit the
// expression may mix types with the columns it is evaluated against.
// Coerce it against the schema so comparisons share a common type, the
// same way `PredicateGroup` does for region partition filters.
let simplifier = ExprSimplifier::new(
SimplifyContext::builder()
.with_schema(df_schema.clone())
.build(),
);
let expr = simplifier
.coerce(expr, &df_schema)
.context(error::CoercePartitionExprSnafu)?;
create_physical_expr(
&expr,
&df_schema,
&ExecutionProps::default(),
&PhysicalPlanningContext::default(),
)
.context(error::CreatePhysicalExprSnafu)
}
impl Display for PartitionExpr {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{} {} {}", self.lhs, self.op, self.rhs)
@@ -813,6 +846,42 @@ mod tests {
assert!(expr5.is_none());
}
#[test]
fn test_serde_iceberg_hash_partition_expr() {
let hash = Operand::IcebergHash(IcebergHash::new("host"));
let expr = hash.clone().lt(Value::Int32(0));
let json = expr.as_json_str().unwrap();
assert_eq!(
json,
r#"{"Expr":{"lhs":{"IcebergHash":"host"},"op":"Lt","rhs":{"Value":{"Int32":0}}}}"#
);
assert_eq!(PartitionExpr::from_json_str(&json).unwrap().unwrap(), expr);
let swapped =
r#"{"Expr":{"lhs":{"Value":{"Int32":0}},"op":"Gt","rhs":{"IcebergHash":"host"}}}"#;
assert_eq!(
PartitionExpr::from_json_str(swapped).unwrap().unwrap(),
expr
);
assert_eq!(expr.to_string(), "iceberg_hash(host) < 0");
}
#[test]
fn test_logical_expr_keeps_quotes_in_column_names() {
use datatypes::arrow::datatypes::{DataType, Field, Schema};
// Declared as `"a""b"`; the stored name contains the quote.
let name = r#"a"b"#;
let schema = Arc::new(Schema::new(vec![Field::new(name, DataType::Utf8, true)]));
for expr in [
col(name).lt(Value::from("m")),
Operand::IcebergHash(IcebergHash::new(name)).lt(Value::Int32(0)),
] {
expr.try_as_physical_expr(&schema)
.unwrap_or_else(|e| panic!("{expr}: {e}"));
}
}
#[test]
fn test_collect_column_names() {
// Simple expression: col_a = 1 should give {col_a}
@@ -854,10 +923,8 @@ mod tests {
assert!(columns.contains("c"));
}
/// Real serialized expressions must survive the round trip through
/// `common-meta`'s persisted-expression mirror: whenever that mirror drifts
/// from these types, deserialization fails, validation is skipped and the
/// overflowing cases below stop producing `PartitionExprIncompatible`.
/// `common-meta` mirrors these types to validate ALTERs. If the mirror drifts,
/// it can't parse the expressions and skips validation, which fails these cases.
#[test]
fn test_serialized_expressions_validate_against_type_changes() {
let overflow_ms = serde_json::from_str::<Value>(
@@ -884,10 +951,7 @@ mod tests {
HashMap::from([(column.to_string(), (old_type, new_type))])
};
// Compile-time tripwire: exhaustively matching the operand and operator
// variants makes a new variant break this test at compile time, so the
// `partition_compat::Persisted*` mirror and these cases must be extended
// together instead of silently skipping validation.
// Fails to compile on a new variant, which the mirror must learn too.
fn walk(expr: &PartitionExpr) {
walk_operand(&expr.lhs);
walk_operand(&expr.rhs);
@@ -904,13 +968,11 @@ mod tests {
}
fn walk_operand(operand: &Operand) {
match operand {
Operand::Column(_) | Operand::Value(_) => {}
Operand::Column(_) | Operand::Value(_) | Operand::IcebergHash(_) => {}
Operand::Expr(expr) => walk(expr),
}
}
// Every expression shape must be rejected while the bound overflows
// nanoseconds, with the error naming the offending column.
let simple = col("ts").lt(overflow_ms.clone());
let nested_and = col("ts")
.gt_eq(start_ms.clone())
@@ -920,7 +982,10 @@ mod tests {
RestrictedOp::Or,
Operand::Expr(col("ts").lt(overflow_ms.clone())),
);
for expr in [&simple, &nested_and, &nested_or] {
let with_hash = col("ts")
.lt(overflow_ms.clone())
.and(Operand::IcebergHash(IcebergHash::new("host")).lt(Value::Int32(0)));
for expr in [&simple, &nested_and, &nested_or, &with_hash] {
walk(expr);
let err = validate_region_partition_bounds(
&changed("ts", ts_ms.clone(), ts_ns.clone()),
@@ -941,9 +1006,6 @@ mod tests {
validate_region_partition_bounds(&changed("ts", ts_ms, ts_ns), &routes(&compatible))
.unwrap();
// Narrowing a field column must be rejected as well: the region's write
// path compares the persisted Int64 bound against Int32 rows directly
// and panics, even though the scan would coerce both sides.
let narrowing = col("v").lt(Value::Int64(3_000_000_000));
let err = validate_region_partition_bounds(
&changed(
@@ -961,5 +1023,24 @@ mod tests {
),
"unexpected error: {err}"
);
// The hash of a column depends on its type.
let hashed = Operand::IcebergHash(IcebergHash::new("v")).lt(Value::Int32(0));
let err = validate_region_partition_bounds(
&changed(
"v",
ConcreteDataType::int32_datatype(),
ConcreteDataType::int64_datatype(),
),
&routes(&hashed),
)
.unwrap_err();
assert!(
matches!(
err,
Error::PartitionExprIncompatible { ref column, .. } if column == "v"
),
"unexpected error: {err}"
);
}
}
src/partition/src/function.rs
+345
View File
@@ -0,0 +1,345 @@
// 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.
//! `iceberg_hash`, the hash allowed on the column side of a partition rule.
//!
//! The hash decides which region a row is written to and is persisted in
//! partition metadata, so its result for a given input must never change across
//! versions. A different algorithm needs a new type.
//!
//! Rule simplification and coverage checks treat the hash as an independent
//! dimension, which holds because it is defined for every value of the accepted
//! column types.
use std::fmt::{Display, Formatter};
use std::sync::Arc;
use datafusion_common::arrow::array::{Array, ArrayRef, AsArray, Int32Array};
use datafusion_common::arrow::compute::take;
use datafusion_common::arrow::datatypes::{
DataType, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
};
use datafusion_common::{DataFusionError, ScalarValue};
use datafusion_expr::{
ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};
use datatypes::prelude::ConcreteDataType;
use datatypes::value::Value;
use serde::{Deserialize, Serialize};
use crate::error::{self, Result};
/// `iceberg_hash(column)`: the 32-bit Murmur3 hash defined by the Iceberg table
/// spec (appendix B) over a string or integer column, as an `INT`.
///
/// Integers of every width hash as Iceberg longs. `UINT64` values hash their bits
/// as a long, so values up to `i64::MAX` match Iceberg. `NULL` hashes to `NULL`.
/// The field is the hashed column.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct IcebergHash(String);
impl IcebergHash {
pub const NAME: &str = "iceberg_hash";
pub fn new(column: impl Into<String>) -> Self {
Self(column.into())
}
pub fn column(&self) -> &str {
&self.0
}
pub fn accepts(data_type: &ConcreteDataType) -> bool {
data_type.is_string() || data_type.is_unsigned() || is_signed_integer(data_type)
}
/// The type of the result, which is also the type of the literals it is compared with.
pub fn output_type() -> ConcreteDataType {
ConcreteDataType::int32_datatype()
}
pub fn evaluate_value(value: &Value) -> Result<Value> {
let hash = match value {
Value::Null => return Ok(Value::Null),
Value::String(v) => murmur3_32(v.as_utf8().as_bytes()),
Value::Int8(v) => hash_long(*v as i64),
Value::Int16(v) => hash_long(*v as i64),
Value::Int32(v) => hash_long(*v as i64),
Value::Int64(v) => hash_long(*v),
Value::UInt8(v) => hash_long(*v as i64),
Value::UInt16(v) => hash_long(*v as i64),
Value::UInt32(v) => hash_long(*v as i64),
Value::UInt64(v) => hash_long(*v as i64),
other => {
return error::UnexpectedSnafu {
err_msg: format!("Unsupported value for {}: {other:?}", Self::NAME),
}
.fail();
}
};
Ok(Value::Int32(hash))
}
pub fn evaluate_array(array: &dyn Array) -> datafusion_common::Result<ArrayRef> {
iceberg_hash_array(array)
}
pub fn to_logical_expr(&self) -> Expr {
let column = Expr::Column(datafusion_common::Column::from_name(self.column()));
ScalarUDF::new_from_impl(IcebergHashUdf::default()).call(vec![column])
}
}
impl Display for IcebergHash {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{}({})", Self::NAME, self.column())
}
}
fn is_signed_integer(data_type: &ConcreteDataType) -> bool {
matches!(
data_type,
ConcreteDataType::Int8(_)
| ConcreteDataType::Int16(_)
| ConcreteDataType::Int32(_)
| ConcreteDataType::Int64(_)
)
}
fn iceberg_hash_array(array: &dyn Array) -> datafusion_common::Result<ArrayRef> {
let hashes: Int32Array = match array.data_type() {
DataType::Utf8 => array
.as_string::<i32>()
.iter()
.map(|v| v.map(|v| murmur3_32(v.as_bytes())))
.collect(),
DataType::LargeUtf8 => array
.as_string::<i64>()
.iter()
.map(|v| v.map(|v| murmur3_32(v.as_bytes())))
.collect(),
DataType::Utf8View => array
.as_string_view()
.iter()
.map(|v| v.map(|v| murmur3_32(v.as_bytes())))
.collect(),
DataType::Int8 => array
.as_primitive::<Int8Type>()
.unary(|v| hash_long(v as i64)),
DataType::Int16 => array
.as_primitive::<Int16Type>()
.unary(|v| hash_long(v as i64)),
DataType::Int32 => array
.as_primitive::<Int32Type>()
.unary(|v| hash_long(v as i64)),
DataType::Int64 => array.as_primitive::<Int64Type>().unary(hash_long),
DataType::UInt8 => array
.as_primitive::<UInt8Type>()
.unary(|v| hash_long(v as i64)),
DataType::UInt16 => array
.as_primitive::<UInt16Type>()
.unary(|v| hash_long(v as i64)),
DataType::UInt32 => array
.as_primitive::<UInt32Type>()
.unary(|v| hash_long(v as i64)),
DataType::UInt64 => array
.as_primitive::<UInt64Type>()
.unary(|v| hash_long(v as i64)),
DataType::Dictionary(_, _) => {
let dictionary = array.as_any_dictionary();
let values = iceberg_hash_array(dictionary.values().as_ref())?;
return Ok(take(values.as_ref(), dictionary.keys(), None)?);
}
other => {
return Err(DataFusionError::Execution(format!(
"{} does not support type {other}",
IcebergHash::NAME
)));
}
};
Ok(Arc::new(hashes))
}
fn hash_long(value: i64) -> i32 {
murmur3_32(&value.to_le_bytes())
}
/// MurmurHash3 x86 32-bit with seed 0, returned as a signed integer like Iceberg.
fn murmur3_32(data: &[u8]) -> i32 {
const C1: u32 = 0xcc9e_2d51;
const C2: u32 = 0x1b87_3593;
fn mix_k(k: u32) -> u32 {
k.wrapping_mul(C1).rotate_left(15).wrapping_mul(C2)
}
let mut h = 0u32;
let (blocks, tail) = data.as_chunks::<4>();
for block in blocks {
h = (h ^ mix_k(u32::from_le_bytes(*block)))
.rotate_left(13)
.wrapping_mul(5)
.wrapping_add(0xe654_6b64);
}
if !tail.is_empty() {
let k = tail
.iter()
.rev()
.fold(0u32, |k, byte| (k << 8) | *byte as u32);
h ^= mix_k(k);
}
// Murmur3 mixes in the length modulo 2^32.
h ^= data.len() as u32;
h ^= h >> 16;
h = h.wrapping_mul(0x85eb_ca6b);
h ^= h >> 13;
h = h.wrapping_mul(0xc2b2_ae35);
h ^= h >> 16;
h as i32
}
/// DataFusion form of [`IcebergHash`], so storage engines
/// can evaluate a region's partition expression as a filter.
#[derive(Debug, PartialEq, Eq, Hash)]
struct IcebergHashUdf {
signature: Signature,
}
impl Default for IcebergHashUdf {
fn default() -> Self {
Self {
signature: Signature::any(1, Volatility::Immutable),
}
}
}
impl ScalarUDFImpl for IcebergHashUdf {
fn name(&self) -> &str {
IcebergHash::NAME
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
Ok(DataType::Int32)
}
fn invoke_with_args(
&self,
args: ScalarFunctionArgs,
) -> datafusion_common::Result<ColumnarValue> {
let [arg] = args.args.as_slice() else {
return Err(DataFusionError::Internal(format!(
"{} expects one argument, got {}",
IcebergHash::NAME,
args.args.len()
)));
};
match arg {
ColumnarValue::Array(array) => {
Ok(ColumnarValue::Array(iceberg_hash_array(array.as_ref())?))
}
ColumnarValue::Scalar(scalar) => {
let hashes = iceberg_hash_array(scalar.to_array()?.as_ref())?;
Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
&hashes, 0,
)?))
}
}
}
}
#[cfg(test)]
mod tests {
use datafusion_common::arrow::array::{
Int64Array, StringArray, StringDictionaryBuilder, StringViewArray, UInt64Array,
};
use datafusion_common::arrow::datatypes::UInt16Type;
use super::*;
#[test]
fn test_iceberg_spec_vectors() {
// Iceberg table spec, appendix B.
assert_eq!(hash_long(34), 2017239379);
// date 2017-11-16 hashes as int 17486.
assert_eq!(hash_long(17486), -653330422);
// time 22:31:08 hashes as long 81068000000 micros.
assert_eq!(hash_long(81068000000), -662762989);
assert_eq!(murmur3_32("iceberg".as_bytes()), 1210000089);
assert_eq!(murmur3_32(&[0, 1, 2, 3]), -188683207);
}
#[test]
fn test_unsigned_policy() {
assert_eq!(
IcebergHash::evaluate_value(&Value::UInt64(34)).unwrap(),
Value::Int32(2017239379)
);
assert_eq!(
IcebergHash::evaluate_value(&Value::UInt64(u64::MAX)).unwrap(),
Value::Int32(hash_long(-1))
);
}
#[test]
fn test_array_matches_value() {
let strings = [Some("iceberg"), None, Some(""), Some("中🙂")];
let mut dictionary = StringDictionaryBuilder::<UInt16Type>::new();
for v in strings {
dictionary.append_option(v);
}
let string_arrays: Vec<ArrayRef> = vec![
Arc::new(StringArray::from(strings.to_vec())),
Arc::new(StringViewArray::from(strings.to_vec())),
Arc::new(dictionary.finish()),
];
let string_values = strings.map(|v| v.map_or(Value::Null, Value::from));
let ints = [Some(34), None, Some(i64::MIN), Some(-1)];
let unsigned = [Some(34), None, Some(u64::MAX), Some(0)];
let cases = string_arrays
.into_iter()
.map(|array| (array, string_values.to_vec()))
.chain([
(
Arc::new(Int64Array::from(ints.to_vec())) as ArrayRef,
ints.map(|v| v.map_or(Value::Null, Value::Int64)).to_vec(),
),
(
Arc::new(UInt64Array::from(unsigned.to_vec())) as ArrayRef,
unsigned
.map(|v| v.map_or(Value::Null, Value::UInt64))
.to_vec(),
),
]);
for (array, values) in cases {
let hashes = IcebergHash::evaluate_array(array.as_ref()).unwrap();
let hashes = hashes.as_primitive::<Int32Type>();
for (row, value) in values.iter().enumerate() {
let expected = IcebergHash::evaluate_value(value).unwrap();
let actual = if hashes.is_null(row) {
Value::Null
} else {
Value::Int32(hashes.value(row))
};
assert_eq!(actual, expected, "{:?} row {row}", array.data_type());
}
}
}
}
src/partition/src/lib.rs
+1
View File
@@ -19,6 +19,7 @@ pub mod checker;
pub mod collider;
pub mod error;
pub mod expr;
pub mod function;
pub mod manager;
pub mod multi_dim;
pub mod overlap;
src/partition/src/multi_dim.rs
+293 -16
View File
@@ -14,17 +14,19 @@
use std::any::Any;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::collections::{BTreeSet, HashMap};
use std::sync::{Arc, RwLock};
use datafusion_expr::ColumnarValue;
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_expr::{ColumnarValue, Expr};
use datafusion_physical_expr::PhysicalExpr;
use datatypes::arrow;
use datatypes::arrow::array::{Array, BooleanArray, BooleanBufferBuilder, RecordBatch};
use datatypes::arrow::buffer::BooleanBuffer;
use datatypes::arrow::compute::prep_null_mask_filter;
use datatypes::arrow::datatypes::Schema;
use datatypes::arrow::datatypes::{Field, Schema};
use datatypes::data_type::DataType as _;
use datatypes::prelude::Value;
use datatypes::vectors::{Helper, VectorRef};
use serde::{Deserialize, Serialize};
@@ -34,13 +36,26 @@ use store_api::storage::RegionNumber;
use crate::PartitionRule;
use crate::checker::PartitionChecker;
use crate::error::{self, Result, UndefinedColumnSnafu};
use crate::expr::{Operand, PartitionExpr, RestrictedOp};
use crate::expr::{Operand, PartitionExpr, RestrictedOp, plan_physical_expr};
use crate::function::IcebergHash;
use crate::partition::RegionMask;
/// The default region number when no partition exprs are matched.
const DEFAULT_REGION: RegionNumber = 0;
type PhysicalExprCache = Option<(Vec<Arc<dyn PhysicalExpr>>, Arc<Schema>)>;
type PhysicalExprCache = Option<Arc<PhysicalPartitionExprs>>;
/// Partition predicates planned for one input schema.
#[derive(Debug)]
struct PhysicalPartitionExprs {
input_schema: Arc<Schema>,
exprs: Vec<Arc<dyn PhysicalExpr>>,
/// Input indices of the columns under `iceberg_hash`. Each hash is evaluated
/// once per batch and appended to the input as `eval_schema` describes, and
/// `exprs` read the results instead of hashing per predicate.
hash_inputs: Vec<usize>,
eval_schema: Arc<Schema>,
}
/// Multi-Dimiension partition rule. RFC [here](https://github.com/GreptimeTeam/greptimedb/blob/main/docs/rfcs/2024-02-21-multi-dimension-partition-rule/rfc.md)
///
@@ -137,6 +152,14 @@ impl MultiDimPartitionRule {
let r = &values[*index];
Self::perform_op(l, &expr.op, r)
}
(Operand::IcebergHash(hash), Operand::Value(r)) => {
let l = self.evaluate_hash(hash, values)?;
Self::perform_op(&l, &expr.op, r)
}
(Operand::Value(l), Operand::IcebergHash(hash)) => {
let r = self.evaluate_hash(hash, values)?;
Self::perform_op(l, &expr.op, &r)
}
(Operand::Expr(lhs), Operand::Expr(rhs)) => {
let lhs = self.evaluate_expr(lhs, values)?;
let rhs = self.evaluate_expr(rhs, values)?;
@@ -150,6 +173,15 @@ impl MultiDimPartitionRule {
}
}
fn evaluate_hash(&self, hash: &IcebergHash, values: &[Value]) -> Result<Value> {
let column = hash.column();
let index = self
.name_to_index
.get(column)
.context(UndefinedColumnSnafu { column })?;
IcebergHash::evaluate_value(&values[*index])
}
fn perform_op(lhs: &Value, op: &RestrictedOp, rhs: &Value) -> Result<bool> {
let result = match op {
RestrictedOp::Eq => lhs.eq(rhs),
@@ -239,27 +271,38 @@ impl MultiDimPartitionRule {
}
let physical_exprs = {
let cache_read_guard = self.physical_expr_cache.read().unwrap();
if let Some((cached_exprs, schema)) = cache_read_guard.as_ref()
&& schema == record_batch.schema_ref()
if let Some(cached) = cache_read_guard.as_ref()
&& &cached.input_schema == record_batch.schema_ref()
{
cached_exprs.clone()
cached.clone()
} else {
drop(cache_read_guard); // Release the read lock before acquiring write lock
let schema = record_batch.schema();
let new_cache = self
.exprs
.iter()
.map(|e| e.try_as_physical_expr(&schema))
.collect::<Result<Vec<_>>>()?;
let new_cache = Arc::new(self.plan_physical_exprs(record_batch.schema())?);
let mut cache_write_guard = self.physical_expr_cache.write().unwrap();
cache_write_guard.replace((new_cache.clone(), schema));
cache_write_guard.replace(new_cache.clone());
new_cache
}
};
let eval_batch;
let record_batch = if physical_exprs.hash_inputs.is_empty() {
record_batch
} else {
let mut columns = record_batch.columns().to_vec();
for index in &physical_exprs.hash_inputs {
columns.push(
IcebergHash::evaluate_array(record_batch.column(*index).as_ref())
.context(error::EvaluateRecordBatchSnafu)?,
);
}
eval_batch = RecordBatch::try_new(physical_exprs.eval_schema.clone(), columns)
.context(error::ComputeArrowKernelSnafu)?;
&eval_batch
};
let mut result: HashMap<u32, RegionMask> = physical_exprs
.exprs
.iter()
.zip(self.regions.iter())
.filter_map(|(expr, region_num)| {
@@ -325,6 +368,97 @@ impl MultiDimPartitionRule {
}
Ok(result)
}
fn plan_physical_exprs(&self, input_schema: Arc<Schema>) -> Result<PhysicalPartitionExprs> {
let mut hashes = BTreeSet::new();
for expr in &self.exprs {
collect_hashes(expr, &mut hashes);
}
if hashes.is_empty() {
let exprs = self
.exprs
.iter()
.map(|e| e.try_as_physical_expr(&input_schema))
.collect::<Result<Vec<_>>>()?;
return Ok(PhysicalPartitionExprs {
eval_schema: input_schema.clone(),
input_schema,
exprs,
hash_inputs: vec![],
});
}
let mut fields = input_schema.fields().to_vec();
let mut hash_columns = HashMap::with_capacity(hashes.len());
let mut hash_inputs = Vec::with_capacity(hashes.len());
for hash in hashes {
let index =
input_schema
.index_of(hash.column())
.ok()
.context(UndefinedColumnSnafu {
column: hash.column(),
})?;
let mut name = format!("__partition_hash_{}", hash_inputs.len());
while input_schema.field_with_name(&name).is_ok() {
name.push('_');
}
fields.push(Arc::new(Field::new(
&name,
IcebergHash::output_type().as_arrow_type(),
true,
)));
hash_columns.insert(
hash.to_logical_expr(),
Expr::Column(datafusion_common::Column::from_name(name)),
);
hash_inputs.push(index);
}
let eval_schema = Arc::new(Schema::new(fields));
let exprs = self
.exprs
.iter()
.map(|expr| {
let expr = expr
.try_as_logical_expr()?
.transform(|expr| {
Ok(match hash_columns.get(&expr) {
Some(column) => Transformed::yes(column.clone()),
None => Transformed::no(expr),
})
})
.context(error::CreatePhysicalExprSnafu)?
.data;
debug_assert!(
!expr
.exists(|expr| Ok(matches!(expr, Expr::ScalarFunction(_))))
.unwrap(),
"iceberg_hash should read the precomputed columns: {expr}"
);
plan_physical_expr(expr, &eval_schema)
})
.collect::<Result<Vec<_>>>()?;
Ok(PhysicalPartitionExprs {
input_schema,
exprs,
hash_inputs,
eval_schema,
})
}
}
fn collect_hashes(expr: &PartitionExpr, hashes: &mut BTreeSet<IcebergHash>) {
for operand in [expr.lhs(), expr.rhs()] {
match operand {
Operand::IcebergHash(hash) => {
hashes.insert(hash.clone());
}
Operand::Expr(expr) => collect_hashes(expr, hashes),
Operand::Column(_) | Operand::Value(_) => {}
}
}
}
fn columnar_value_to_boolean_array(
@@ -858,6 +992,71 @@ mod test_split_record_batch {
}
}
#[test]
fn test_split_record_batch_widened_time_index_unit_with_hash() {
// The precomputed-hash path must coerce millisecond bounds against a
// nanosecond time index too.
let bound = Value::from(Timestamp::new_millisecond(1_000));
let hash = Operand::IcebergHash(IcebergHash::new("host"));
let exprs = [
(
col("ts").lt(bound.clone()),
hash.clone().lt(Value::Int32(0)),
),
(
col("ts").lt(bound.clone()),
hash.clone().gt_eq(Value::Int32(0)),
),
(
col("ts").gt_eq(bound.clone()),
hash.clone().lt(Value::Int32(0)),
),
(col("ts").gt_eq(bound), hash.gt_eq(Value::Int32(0))),
]
.into_iter()
.map(|(ts, hash)| ts.and(hash))
.collect();
let rule = MultiDimPartitionRule::try_new(
vec!["ts".to_string(), "host".to_string()],
vec![1, 2, 3, 4],
exprs,
true,
)
.unwrap();
let hosts = ["a", "b", "c", "server42"];
let mut ts = vec![];
let mut host = vec![];
for nanos in [500_000_000, 2_000_000_000] {
for h in hosts {
ts.push(nanos);
host.push(h);
}
}
let schema = Arc::new(Schema::new(vec![
Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), false),
Field::new("host", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(TimestampNanosecondArray::from(ts)),
Arc::new(StringArray::from(host)),
],
)
.unwrap();
let expected = rule.split_record_batch_naive(&batch).unwrap();
let result = rule.split_record_batch(&batch).unwrap();
for (region, expected) in expected {
let actual = result
.get(&region)
.map(|mask| mask.array().clone())
.unwrap_or_else(|| BooleanArray::from(vec![false; batch.num_rows()]));
assert_eq!(actual, expected, "region {region}");
}
}
#[test]
fn test_split_record_batch_widened_time_index_unit() {
// https://github.com/GreptimeTeam/greptimedb/issues/9747
@@ -1021,6 +1220,84 @@ mod test_split_record_batch {
assert_eq!(total_rows, batch.num_rows());
}
#[test]
fn test_hash_partition_batch_matches_row() {
use datatypes::arrow::array::{Array, StringDictionaryBuilder, UInt64Array};
use datatypes::arrow::datatypes::{DataType, Field, UInt32Type};
let host_hash = Operand::IcebergHash(IcebergHash::new("host"));
let tenant_hash = Operand::IcebergHash(IcebergHash::new("tenant"));
let exprs = [
host_hash
.clone()
.lt(Value::Int32(0))
.and(col("tenant").lt(Value::UInt64(100))),
host_hash
.clone()
.lt(Value::Int32(0))
.and(col("tenant").gt_eq(Value::UInt64(100))),
host_hash
.clone()
.gt_eq(Value::Int32(0))
.and(tenant_hash.clone().lt(Value::Int32(0))),
host_hash
.gt_eq(Value::Int32(0))
.and(tenant_hash.gt_eq(Value::Int32(0))),
]
.iter()
.map(|expr| {
PartitionExpr::from_json_str(&expr.as_json_str().unwrap())
.unwrap()
.unwrap()
})
.collect();
let rule = MultiDimPartitionRule::try_new(
vec!["host".to_string(), "tenant".to_string()],
vec![1, 2, 3, 4],
exprs,
true,
)
.unwrap();
let hosts = [
None,
Some(""),
Some("a"),
Some("b"),
Some("中🙂"),
Some("server42"),
];
let tenants = [None, Some(0), Some(99), Some(100), Some(u64::MAX)];
let mut host_builder = StringDictionaryBuilder::<UInt32Type>::new();
let mut tenant_values = vec![];
for host in hosts {
for tenant in tenants {
host_builder.append_option(host);
tenant_values.push(tenant);
}
}
let host_array = Arc::new(host_builder.finish());
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("host", host_array.data_type().clone(), true),
Field::new("tenant", DataType::UInt64, true),
])),
vec![host_array, Arc::new(UInt64Array::from(tenant_values))],
)
.unwrap();
let by_row = rule.split_record_batch_naive(&batch).unwrap();
let by_batch = rule.split_record_batch(&batch).unwrap();
assert!(!by_batch.contains_key(&DEFAULT_REGION));
for (region, expected) in by_row {
let actual = by_batch
.get(&region)
.map(|mask| mask.array().clone())
.unwrap_or_else(|| BooleanArray::from(vec![false; batch.num_rows()]));
assert_eq!(actual, expected, "region {region}");
}
}
#[test]
fn test_comparison_batch_matches_row() {
use datatypes::arrow::array::Array;
src/partition/src/simplify.rs
+30 -16
View File
@@ -23,8 +23,8 @@ use std::ops::Bound;
use datatypes::value::{OrderedF64, Value};
use crate::collider::{AtomicExpr, Collider, GluonOp, NucleonExpr};
use crate::expr::{Operand, PartitionExpr, RestrictedOp, col};
use crate::collider::{AtomicExpr, Collider, GluonOp, NucleonExpr, PartitionDim};
use crate::expr::{Operand, PartitionExpr, RestrictedOp};
/// Attempts to simplify a merged partition expression (typically an `OR` of multiple partitions)
/// into an equivalent but shorter expression.
@@ -36,7 +36,7 @@ pub fn simplify_merged_partition_expr(expr: PartitionExpr) -> PartitionExpr {
try_simplify_merged_partition_expr(&expr).unwrap_or(expr)
}
type DenormValues = BTreeMap<String, BTreeMap<OrderedF64, Value>>;
type DenormValues = BTreeMap<PartitionDim, BTreeMap<OrderedF64, Value>>;
fn try_simplify_merged_partition_expr(expr: &PartitionExpr) -> Option<PartitionExpr> {
let collider = Collider::new(std::slice::from_ref(expr)).ok()?;
@@ -87,7 +87,7 @@ fn term_from_atomic(atomic: &AtomicExpr, denorm_values: &DenormValues) -> Option
let interval = interval_from_nucleons(&atomic.nucleons[start..i])?;
if !interval.is_unbounded() {
constraints.insert(column.to_string(), interval);
constraints.insert(column.clone(), interval);
}
}
@@ -105,7 +105,7 @@ fn interval_from_nucleons(nucleons: &[NucleonExpr]) -> Option<Interval> {
#[derive(Debug, Clone, PartialEq, Eq)]
struct Term {
// Only stores constrained columns. Missing column means unbounded.
constraints: BTreeMap<String, Interval>,
constraints: BTreeMap<PartitionDim, Interval>,
}
impl Term {
@@ -358,11 +358,11 @@ fn simplify_terms(mut terms: Vec<Term>) -> Option<Vec<Term>> {
fn try_merge_terms(a: &Term, b: &Term) -> Option<Term> {
// Find the only differing column (treat missing as unbounded).
let mut diff_col: Option<&str> = None;
let mut diff_col: Option<&PartitionDim> = None;
let mut cols = BTreeSet::new();
cols.extend(a.constraints.keys().map(|s| s.as_str()));
cols.extend(b.constraints.keys().map(|s| s.as_str()));
cols.extend(a.constraints.keys());
cols.extend(b.constraints.keys());
for col in cols {
let a_interval = a.constraints.get(col);
@@ -387,7 +387,7 @@ fn try_merge_terms(a: &Term, b: &Term) -> Option<Term> {
constraints.remove(diff_col);
}
UnionInterval::Interval(interval) => {
constraints.insert(diff_col.to_string(), interval);
constraints.insert(diff_col.clone(), interval);
}
}
@@ -440,7 +440,7 @@ fn term_to_expr(term: &Term, denorm_values: &DenormValues) -> Option<PartitionEx
}
fn interval_to_exprs(
column: &str,
column: &PartitionDim,
interval: &Interval,
denorm_values: &DenormValues,
) -> Option<Vec<PartitionExpr>> {
@@ -457,7 +457,7 @@ fn interval_to_exprs(
match (lower, upper) {
(Included(lv), Included(uv)) if lv == uv => {
return Some(vec![col(column).eq(col_values.get(lv)?.clone())]);
return Some(vec![column.to_operand().eq(col_values.get(lv)?.clone())]);
}
(Excluded(lv), Excluded(uv)) if lv == uv => return None,
(Included(lv), Excluded(uv)) if lv == uv => return None,
@@ -468,13 +468,13 @@ fn interval_to_exprs(
let mut exprs = Vec::new();
match lower {
Unbounded => {}
Included(v) => exprs.push(col(column).gt_eq(col_values.get(v)?.clone())),
Excluded(v) => exprs.push(col(column).gt(col_values.get(v)?.clone())),
Included(v) => exprs.push(column.to_operand().gt_eq(col_values.get(v)?.clone())),
Excluded(v) => exprs.push(column.to_operand().gt(col_values.get(v)?.clone())),
}
match upper {
Unbounded => {}
Included(v) => exprs.push(col(column).lt_eq(col_values.get(v)?.clone())),
Excluded(v) => exprs.push(col(column).lt(col_values.get(v)?.clone())),
Included(v) => exprs.push(column.to_operand().lt_eq(col_values.get(v)?.clone())),
Excluded(v) => exprs.push(column.to_operand().lt(col_values.get(v)?.clone())),
}
Some(exprs)
@@ -487,7 +487,7 @@ mod tests {
use datatypes::value::{OrderedFloat, Value};
use super::*;
use crate::expr::Operand;
use crate::expr::col;
fn or(lhs: PartitionExpr, rhs: PartitionExpr) -> PartitionExpr {
PartitionExpr::new(Operand::Expr(lhs), RestrictedOp::Or, Operand::Expr(rhs))
@@ -508,6 +508,20 @@ mod tests {
assert_eq!(simplified.to_string(), "device_id < 100");
}
#[test]
fn simplify_adjacent_hash_ranges() {
use crate::function::IcebergHash;
let hash = Operand::IcebergHash(IcebergHash::new("host"));
let left = hash.clone().lt(Value::Int32(0));
let right = hash
.clone()
.gt_eq(Value::Int32(0))
.and(hash.clone().lt(Value::Int32(100)));
let simplified = simplify_merged_partition_expr(or(left, right));
assert_eq!(simplified, hash.lt(Value::Int32(100)));
}
#[test]
fn simplify_adjacent_ranges() {
// host < 'h0' OR (host >= 'h0' AND host < 'h1') -> host < 'h1'
src/query/src/dist_plan/predicate_extractor.rs
+2 -1
View File
@@ -227,7 +227,7 @@ impl ExpressionChecker {
result: &mut Vec<PartitionExpr>,
) {
match operand {
Operand::Column(_) | Operand::Value(_) => {
Operand::Column(_) | Operand::Value(_) | Operand::IcebergHash(_) => {
// This shouldn't happen in well-formed expressions
}
Operand::Expr(expr) => {
@@ -252,6 +252,7 @@ impl ExpressionChecker {
) -> bool {
match operand {
Operand::Column(col) => partition_columns.contains(col),
Operand::IcebergHash(hash) => partition_columns.contains(hash.column()),
Operand::Value(_) => true, // Values are always safe
Operand::Expr(expr) => {
Self::expr_only_involves_partition_columns(expr, partition_columns)
src/query/src/dist_plan/region_pruner.rs
+55
View File
@@ -206,6 +206,10 @@ impl ConstraintPruner {
true
}
// Query predicates never constrain a function dimension, so its values all
// come from partition rules, already typed as the function output.
(Operand::IcebergHash(_), Operand::Value(_))
| (Operand::Value(_), Operand::IcebergHash(_)) => true,
_ => false,
}
}
@@ -270,6 +274,57 @@ mod tests {
assert!(pruned.contains(&RegionId::new(1, 2)));
}
#[test]
fn test_constraint_pruning_with_hash_dimension() {
use partition::function::IcebergHash;
let hash = Operand::IcebergHash(IcebergHash::new("host"));
let partitions = [
col("ts")
.lt(Value::Int64(100))
.and(hash.clone().lt(Value::Int32(0))),
col("ts")
.lt(Value::Int64(100))
.and(hash.clone().gt_eq(Value::Int32(0))),
col("ts")
.gt_eq(Value::Int64(100))
.and(hash.clone().lt(Value::Int32(0))),
col("ts")
.gt_eq(Value::Int64(100))
.and(hash.gt_eq(Value::Int32(0))),
]
.into_iter()
.enumerate()
.map(|(i, expr)| create_test_partition_info(i as u64 + 1, Some(expr)))
.collect::<Vec<_>>();
let column_datatypes = HashMap::from_iter([
("ts".to_string(), ConcreteDataType::int64_datatype()),
("host".to_string(), ConcreteDataType::string_datatype()),
]);
let prune = |query: PartitionExpr| {
let mut regions =
ConstraintPruner::prune_regions(&[query], &partitions, column_datatypes.clone())
.unwrap();
regions.sort();
regions
};
// A filter on the hashed column can't exclude any hash range.
assert_eq!(
prune(col("host").eq(Value::String("a".into()))),
(1..=4).map(|i| RegionId::new(1, i)).collect::<Vec<_>>()
);
// Plain columns in the same rule still prune.
assert_eq!(
prune(
col("host")
.eq(Value::String("a".into()))
.and(col("ts").gt_eq(Value::Int64(100)))
),
vec![RegionId::new(1, 3), RegionId::new(1, 4)]
);
}
#[test]
fn test_constraint_pruning_in_list() {
let partitions = vec![
src/sql/src/parsers/create_parser.rs
+13 -2
View File
@@ -28,8 +28,8 @@ pub(crate) use json::parse_json2_type_and_options;
pub use json::parse_json2_type_hint_path;
use snafu::{OptionExt, ResultExt, ensure};
use sqlparser::ast::{
ColumnOption, ColumnOptionDef, DataType, Expr, KeyOrIndexDisplay, NullsDistinctOption,
PrimaryKeyConstraint, UniqueConstraint,
ColumnOption, ColumnOptionDef, DataType, Expr, FunctionArg, FunctionArgExpr, FunctionArguments,
KeyOrIndexDisplay, NullsDistinctOption, PrimaryKeyConstraint, UniqueConstraint,
};
use sqlparser::dialect::keywords::Keyword;
use sqlparser::keywords::ALL_KEYWORDS;
@@ -1300,6 +1300,17 @@ fn ensure_one_expr(expr: &Expr, columns: &[&Column]) -> Result<()> {
);
Ok(())
}
// The function and its argument shape are validated when the rule is converted.
Expr::Function(function) => {
if let FunctionArguments::List(args) = &function.args {
for arg in &args.args {
if let FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) = arg {
ensure_one_expr(expr, columns)?;
}
}
}
Ok(())
}
Expr::Value(_) => Ok(()),
Expr::UnaryOp { expr, .. } => {
ensure_one_expr(expr, columns)?;
src/sql/src/statements/create.rs
+9 -2
View File
@@ -14,6 +14,7 @@
use std::collections::HashMap;
use std::fmt::{Display, Formatter, Write};
use std::ops::ControlFlow;
use common_catalog::consts::FILE_ENGINE;
use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings};
@@ -22,7 +23,7 @@ use datatypes::schema::{CompressionOptions, FulltextOptions, SkippingIndexOption
use itertools::Itertools;
use serde::Serialize;
use snafu::ResultExt;
use sqlparser::ast::{ColumnOptionDef, DataType, Expr};
use sqlparser::ast::{ColumnOptionDef, DataType, Expr, visit_expressions_mut};
use sqlparser_derive::{Visit, VisitMut};
use crate::ast::{ColumnDef, Ident, ObjectName};
@@ -395,11 +396,17 @@ pub struct Partitions {
}
impl Partitions {
/// set quotes to all [Ident]s from column list
/// Sets quotes to all column [Ident]s, in the column list and in the expressions.
pub fn set_quote(&mut self, quote_style: char) {
self.column_list
.iter_mut()
.for_each(|c| c.quote_style = Some(quote_style));
let _ = visit_expressions_mut(&mut self.exprs, |expr| {
if let Expr::Identifier(ident) = expr {
ident.quote_style = Some(quote_style);
}
ControlFlow::<()>::Continue(())
});
}
}
src/table/src/metadata.rs
+54
View File
@@ -1093,6 +1093,16 @@ impl TableMeta {
}
);
} else {
// Partition rules keep bounds typed as the column, and functions such as
// `iceberg_hash` route by the typed value, so a new type would reroute or
// break routing. Time index widening above keeps timestamps comparable.
ensure!(
!self.partition_key_indices.contains(&index),
error::InvalidAlterRequestSnafu {
table: table_name,
err: format!("Not allowed to change partition column '{}'", column.name)
}
);
ensure!(
column
.data_type
@@ -1911,6 +1921,50 @@ mod tests {
}
}
#[test]
fn test_modify_partition_column_type() {
// `col2` is nullable and not in the primary key, so only its partition role
// decides whether the change is allowed.
let alter_kind = AlterKind::ModifyColumnTypes {
columns: vec![ModifyColumnTypeRequest {
column_name: "col2".to_string(),
target_type: ConcreteDataType::float64_datatype(),
}],
};
for (partition_key_indices, allowed) in [(vec![], true), (vec![2], false)] {
let meta = TableMetaBuilder::empty()
.schema(Arc::new(new_test_schema()))
.primary_key_indices(vec![0])
.partition_key_indices(partition_key_indices)
.engine("engine")
.next_column_id(3)
.build()
.unwrap();
let result = meta.builder_with_alter_kind("my_table", &alter_kind);
assert_eq!(result.is_ok(), allowed);
}
// Widening the unit of a time index that is also a partition column is allowed.
let meta = TableMetaBuilder::empty()
.schema(Arc::new(new_test_schema()))
.primary_key_indices(vec![0])
.partition_key_indices(vec![1])
.engine("engine")
.next_column_id(3)
.build()
.unwrap();
let alter_kind = AlterKind::ModifyColumnTypes {
columns: vec![ModifyColumnTypeRequest {
column_name: "ts".to_string(),
target_type: ConcreteDataType::timestamp_microsecond_datatype(),
}],
};
assert!(
meta.builder_with_alter_kind("my_table", &alter_kind)
.is_ok()
);
}
#[test]
fn test_add_columns() {
let schema = Arc::new(new_test_schema());