diff --git a/Cargo.lock b/Cargo.lock index 946472d1df0..73bd260e9de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", diff --git a/src/partition/Cargo.toml b/src/partition/Cargo.toml index 6773cc5878b..43f6c8e792a 100644 --- a/src/partition/Cargo.toml +++ b/src/partition/Cargo.toml @@ -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" diff --git a/src/partition/src/error.rs b/src/partition/src/error.rs index 065ff9d5250..83b9b767346 100644 --- a/src/partition/src/error.rs +++ b/src/partition/src/error.rs @@ -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, } diff --git a/src/partition/src/expr.rs b/src/partition/src/expr.rs index 560dc517060..03e7249797b 100644 --- a/src/partition/src/expr.rs +++ b/src/partition/src/expr.rs @@ -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, diff --git a/src/partition/src/multi_dim.rs b/src/partition/src/multi_dim.rs index 29e45126343..25dcc0fb89d 100644 --- a/src/partition/src/multi_dim.rs +++ b/src/partition/src/multi_dim.rs @@ -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 diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index e74d4d5bf70..80030174424 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -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, + ) { + 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, 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, + 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 diff --git a/tests/cases/standalone/common/alter/change_col_type.result b/tests/cases/standalone/common/alter/change_col_type.result index d362215792b..92a986cd73b 100644 --- a/tests/cases/standalone/common/alter/change_col_type.result +++ b/tests/cases/standalone/common/alter/change_col_type.result @@ -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 diff --git a/tests/cases/standalone/common/alter/change_col_type.sql b/tests/cases/standalone/common/alter/change_col_type.sql index c7d206eaf7c..2da3094a6f1 100644 --- a/tests/cases/standalone/common/alter/change_col_type.sql +++ b/tests/cases/standalone/common/alter/change_col_type.sql @@ -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);