rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-23 04:12:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(partition): coerce partition exprs against the batch schema before routing (#9752)

Browse Source 
* fix(partition): coerce partition exprs against the batch schema before routing

Batch partition routing (MultiDimPartitionRule::split_record_batch) builds
physical expressions via PartitionExpr::try_as_physical_expr, which did no
type coercion. After widening the time index unit (e.g. ms -> ns via
ALTER TABLE ... MODIFY COLUMN ts TIMESTAMP_NS), the persisted partition
bounds keep their original unit, so evaluating the bounds against the batch
failed with:

  Invalid comparison operation: Timestamp(ns) < Timestamp(ms)

Row inserts were unaffected because find_region compares Values and
Timestamp::cmp converts units; scans were unaffected because
PredicateGroup::new coerces the region partition expr against the region
schema first.

Coerce the logical partition expression against the evaluation schema in
try_as_physical_expr before creating the physical expression, the same way
PredicateGroup does. This also covers the region partition filter built for
SST pruning (src/mito2/src/sst/parquet/reader.rs), which shares this
conversion.

Fixes #9747

Signed-off-by: Ning Sun <sunning@greptime.com>

* test: regenerate change_col_type result after comment update

The ts_widen_part comment in change_col_type.sql was reworded after the
result file was last generated; sqlness echoes SQL comments into the
result, so the stale comment made the case fail. Regenerate it.

Signed-off-by: Ning Sun <sunning@greptime.com>

---------

Signed-off-by: Ning Sun <sunning@greptime.com>
This commit is contained in:
Ning Sun
2026-10-20 08:17:52 +00:00
committed by GitHub
parent b9aa627180
commit a7230f7a92
8 changed files with 225 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+2
View File
@@ -10489,9 +10489,11 @@ dependencies = [
"common-error",
"common-macro",
"common-meta",
"common-time",
"criterion 0.7.0",
"datafusion-common",
"datafusion-expr",
"datafusion-optimizer",
"datafusion-physical-expr",
"datatypes",
"moka",
src/partition/Cargo.toml
+2
View File
@@ -15,6 +15,7 @@ common-macro.workspace = true
common-meta.workspace = true
datafusion-common.workspace = true
datafusion-expr.workspace = true
datafusion-optimizer.workspace = true
datafusion-physical-expr.workspace = true
datatypes.workspace = true
moka = { workspace = true, features = ["future"] }
@@ -27,6 +28,7 @@ store-api.workspace = true
table.workspace = true
[dev-dependencies]
common-time.workspace = true
criterion.workspace = true
rand = "0.8"
src/partition/src/error.rs
+9
View File
@@ -206,6 +206,14 @@ pub enum Error {
location: Location,
},
#[snafu(display("Failed to coerce partition expr against the schema"))]
CoercePartitionExpr {
#[snafu(source)]
error: datafusion_common::error::DataFusionError,
#[snafu(implicit)]
location: Location,
},
#[snafu(display("Partition expr value is not supported: {:?}", value))]
UnsupportedPartitionExprValue {
value: Value,
@@ -275,6 +283,7 @@ impl ErrorExt for Error {
Error::UnexpectedColumnType { .. } => StatusCode::Internal,
Error::ToDFSchema { .. } => StatusCode::Internal,
Error::CreatePhysicalExpr { .. } => StatusCode::Internal,
Error::CoercePartitionExpr { .. } => StatusCode::Internal,
Error::UnsupportedPartitionExprValue { .. } => StatusCode::InvalidArguments,
Error::ConvertPartitionExprValue { .. } => StatusCode::InvalidArguments,
}
src/partition/src/expr.rs
+16 -1
View File
@@ -21,6 +21,8 @@ use datafusion_common::{ScalarValue, ToDFSchema};
use datafusion_expr::Expr;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::simplify::SimplifyContext;
use datafusion_optimizer::simplify_expressions::ExprSimplifier;
use datafusion_physical_expr::{PhysicalExpr, create_physical_expr};
use datatypes::arrow;
use datatypes::value::{
@@ -357,8 +359,21 @@ impl PartitionExpr {
.clone()
.to_dfschema_ref()
.context(error::ToDFSchemaSnafu)?;
let execution_props = &ExecutionProps::default();
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,
src/partition/src/multi_dim.rs
+42 -2
View File
@@ -798,9 +798,10 @@ mod tests {
mod test_split_record_batch {
use std::sync::Arc;
use common_time::timestamp::Timestamp;
use datafusion_common::ScalarValue;
use datatypes::arrow::array::{Int64Array, StringArray};
use datatypes::arrow::datatypes::{DataType, Field, Schema};
use datatypes::arrow::array::{Int64Array, StringArray, TimestampNanosecondArray};
use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use datatypes::arrow::record_batch::RecordBatch;
use rand::Rng;
@@ -857,6 +858,45 @@ mod test_split_record_batch {
}
}
#[test]
fn test_split_record_batch_widened_time_index_unit() {
// https://github.com/GreptimeTeam/greptimedb/issues/9747
// The partition bounds keep the unit they were created with (millisecond)
// after the time index column is widened to nanosecond, so the batch
// schema and the bound values mix timestamp units.
let bound = Value::from(Timestamp::new_millisecond(1_000));
let rule = MultiDimPartitionRule::try_new(
vec!["ts".to_string()],
vec![0, 1],
vec![col("ts").lt(bound.clone()), col("ts").gt_eq(bound)],
true,
)
.unwrap();
let schema = Arc::new(Schema::new(vec![Field::new(
"ts",
DataType::Timestamp(TimeUnit::Nanosecond, None),
false,
)]));
let batch = RecordBatch::try_new(
schema,
vec![Arc::new(TimestampNanosecondArray::from(vec![
500_000_000,
2_000_000_000,
]))],
)
.unwrap();
// 500_000_000ns < 1000ms and 2_000_000_000ns >= 1000ms: one row per region.
let expected = rule.split_record_batch_naive(&batch).unwrap();
assert_eq!(expected.len(), 2);
let result = rule.split_record_batch(&batch).unwrap();
assert_eq!(result.len(), expected.len());
for (region, value) in &result {
assert_eq!(value.array(), expected.get(region).unwrap());
}
}
#[test]
fn test_split_record_batch_empty() {
// Create a simple MultiDimPartitionRule
tests-integration/src/grpc/flight.rs
+109 -3
View File
@@ -52,7 +52,9 @@ mod test {
};
use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef};
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector};
use datatypes::vectors::{
Int32Vector, StringVector, TimestampMillisecondVector, TimestampNanosecondVector,
};
use futures_util::{Stream, StreamExt};
use hyper_util::rt::TokioIo;
use itertools::Itertools;
@@ -837,6 +839,82 @@ mod test {
);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_flight_bulk_widen_partitioned_time_index() {
// https://github.com/GreptimeTeam/greptimedb/issues/9747
// Bulk writes route through `MultiDimPartitionRule::split_record_batch`,
// which evaluates the persisted partition bounds against the batch. The
// bounds keep the unit they were created with, so widening the time
// index unit must not mix timestamp units during that evaluation.
common_telemetry::init_default_ut_logging();
let (db, server) = setup_grpc_server(
StorageType::File,
"test_flight_bulk_widen_partitioned_time_index",
)
.await;
let addr = server.bind_addr().unwrap().to_string();
let client = Client::with_urls(vec![addr]);
let client = Database::new_with_dbname("greptime-public", client);
client
.sql(
"CREATE TABLE widen_ts (host STRING, ts TIMESTAMP TIME INDEX) \
PARTITION ON COLUMNS (ts) (\
ts < '2024-01-02 00:00:00', \
ts >= '2024-01-02 00:00:00'\
)",
)
.await
.unwrap();
client
.sql("ALTER TABLE widen_ts MODIFY COLUMN ts TIMESTAMP_NS")
.await
.unwrap();
// 2024-01-02 00:00:00 UTC == 1704153600000 ms; rows below and above the
// bound, expressed in the widened nanosecond unit ("c" only exists in ns).
let schema = Arc::new(Schema::new(vec![
ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
ColumnSchema::new(
"ts",
ConcreteDataType::timestamp_nanosecond_datatype(),
false,
)
.with_time_index(true),
]));
let batch = RecordBatch::new(
schema,
vec![
Arc::new(StringVector::from_vec(vec![
"a".to_string(),
"b".to_string(),
"c".to_string(),
])) as VectorRef,
Arc::new(TimestampNanosecondVector::from_vec(vec![
1_704_153_599_500_000_000,
1_704_153_600_000_000_000,
1_704_153_600_000_000_001,
])) as VectorRef,
],
)
.unwrap();
test_put_record_batches_to_table(&client, "widen_ts", vec![batch]).await;
let sql = "select host, ts from widen_ts order by ts";
let expected = "\
+------+-------------------------------+
| host | ts |
+------+-------------------------------+
| a | 2024-01-01T23:59:59.500 |
| b | 2024-01-02T00:00:00 |
| c | 2024-01-02T00:00:00.000000001 |
+------+-------------------------------+";
query_and_expect(db.frontend().as_ref(), sql, expected).await;
}
#[tokio::test(flavor = "multi_thread")]
async fn test_standalone_flight_do_put_missing_nullable_columns() {
common_telemetry::init_default_ut_logging();
@@ -1124,15 +1202,36 @@ mod test {
test_put_record_batches_with_hints(client, record_batches, &[]).await;
}
async fn test_put_record_batches_to_table(
client: &Database,
table: &str,
record_batches: Vec<RecordBatch>,
) {
test_put_record_batches_with_hints_and_table(client, table, record_batches, &[]).await;
}
async fn test_put_record_batches_with_hints(
client: &Database,
record_batches: Vec<RecordBatch>,
hints: &[(&str, &str)],
) {
test_put_record_batches_with_hints_and_table(client, "foo", record_batches, hints).await
}
async fn test_put_record_batches_with_hints_and_table(
client: &Database,
table: &str,
record_batches: Vec<RecordBatch>,
hints: &[(&str, &str)],
) {
let table = table.to_string();
let requests_count = record_batches.len();
let schema = record_batches[0].schema.arrow_schema().clone();
let stream = futures::stream::once(async move { encode_put_schema(schema.as_ref()) })
let stream =
futures::stream::once(
async move { encode_put_schema_for_table(schema.as_ref(), &table) },
)
.chain(
tokio_stream::iter(record_batches)
.enumerate()
@@ -1169,11 +1268,18 @@ mod test {
}
fn encode_put_schema(schema: &datatypes::arrow::datatypes::Schema) -> FlightData {
encode_put_schema_for_table(schema, "foo")
}
fn encode_put_schema_for_table(
schema: &datatypes::arrow::datatypes::Schema,
table: &str,
) -> FlightData {
let mut schema_data = FlightEncoder::default().encode_schema(schema);
schema_data.app_metadata = serde_json::to_vec(&DoPutMetadata::new(0)).unwrap().into();
schema_data.flight_descriptor = Some(FlightDescriptor {
r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32,
path: vec!["foo".to_string()],
path: vec![table.to_string()],
..Default::default()
});
schema_data
tests/cases/standalone/common/alter/change_col_type.result
+35
View File
@@ -341,6 +341,41 @@ SELECT * FROM ts_widen_part ORDER BY ts;
| b | 2.0 | 2024-01-02T00:00:00.000250 |
+------+-----+----------------------------+
-- force the rows into SSTs so scans read old-unit data through the widened
-- schema of a time-index-partitioned table (https://github.com/GreptimeTeam/greptimedb/issues/9747)
ADMIN FLUSH_TABLE('ts_widen_part');
+------------------------------------+
| ADMIN FLUSH_TABLE('ts_widen_part') |
+------------------------------------+
| 0 |
+------------------------------------+
SELECT * FROM ts_widen_part ORDER BY ts;
+------+-----+----------------------------+
| host | val | ts |
+------+-----+----------------------------+
| a | 1.0 | 2024-01-01T23:59:59.500 |
| b | 2.0 | 2024-01-02T00:00:00.000250 |
+------+-----+----------------------------+
SELECT * FROM ts_widen_part WHERE ts >= '2024-01-02 00:00:00' ORDER BY ts;
+------+-----+----------------------------+
| host | val | ts |
+------+-----+----------------------------+
| b | 2.0 | 2024-01-02T00:00:00.000250 |
+------+-----+----------------------------+
SELECT * FROM ts_widen_part WHERE ts < '2024-01-02 00:00:00' ORDER BY ts;
+------+-----+-------------------------+
| host | val | ts |
+------+-----+-------------------------+
| a | 1.0 | 2024-01-01T23:59:59.500 |
+------+-----+-------------------------+
DROP TABLE ts_widen_part;
Affected Rows: 0
tests/cases/standalone/common/alter/change_col_type.sql
+10
View File
@@ -119,6 +119,16 @@ SELECT * FROM ts_widen_part ORDER BY ts;
INSERT INTO ts_widen_part VALUES ("b", 2, "2024-01-02 00:00:00.000250");
SELECT * FROM ts_widen_part ORDER BY ts;
-- force the rows into SSTs so scans read old-unit data through the widened
-- schema of a time-index-partitioned table (https://github.com/GreptimeTeam/greptimedb/issues/9747)
ADMIN FLUSH_TABLE('ts_widen_part');
SELECT * FROM ts_widen_part ORDER BY ts;
SELECT * FROM ts_widen_part WHERE ts >= '2024-01-02 00:00:00' ORDER BY ts;
SELECT * FROM ts_widen_part WHERE ts < '2024-01-02 00:00:00' ORDER BY ts;
DROP TABLE ts_widen_part;
CREATE TABLE ts_overflow (host STRING, ts TIMESTAMP TIME INDEX);