diff --git a/Cargo.lock b/Cargo.lock index 765638526b464..c8f9fcb3400ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4298,6 +4298,7 @@ dependencies = [ "libm", "match-template", "md-5", + "murmur3", "naive-cityhash", "num-traits", "proptest", @@ -5035,6 +5036,7 @@ dependencies = [ "backoff", "bytes", "chrono", + "databend-common-ast", "databend-common-base", "databend-common-catalog", "databend-common-column", diff --git a/Cargo.toml b/Cargo.toml index d87fa441dbad9..cfec6de7cbc64 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -355,6 +355,7 @@ memchr = { version = "2", default-features = false } micromarshal = "0.7.0" mlua = { version = "0.11", features = ["lua54", "vendored", "async", "serialize"] } mockall = "0.11.2" +murmur3 = "0.5.2" mysql_async = { version = "0.34", default-features = false, features = ["native-tls-tls"] } naive-cityhash = "0.2.0" ndarray = "0.15.6" diff --git a/src/query/ast/src/ast/statements/table.rs b/src/query/ast/src/ast/statements/table.rs index cd959c720fa91..eeb79013be9c9 100644 --- a/src/query/ast/src/ast/statements/table.rs +++ b/src/query/ast/src/ast/statements/table.rs @@ -162,7 +162,7 @@ pub struct CreateTableStmt { pub uri_location: Option, pub cluster_by: Option, pub table_options: BTreeMap, - pub iceberg_table_partition: Option>, + pub partition_by: Option>, pub table_properties: Option>, pub as_query: Option>, pub table_type: TableType, @@ -211,6 +211,12 @@ impl Display for CreateTableStmt { write!(f, " {uri_location}")?; } + if let Some(partition_by) = &self.partition_by { + write!(f, " PARTITION BY(")?; + write_comma_separated_list(f, partition_by)?; + write!(f, ")")?; + } + if let Some(cluster_by) = &self.cluster_by { write!(f, " {cluster_by}")?; } @@ -221,12 +227,6 @@ impl Display for CreateTableStmt { write_space_separated_string_map(f, &self.table_options)?; } - if let Some(iceberg_table_partition) = &self.iceberg_table_partition { - write!(f, " PARTITION BY(")?; - write_comma_separated_list(f, iceberg_table_partition)?; - write!(f, ")")?; - } - if let Some(table_properties) = &self.table_properties { write!(f, " PROPERTIES(")?; write_space_separated_string_map(f, table_properties)?; diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs index b1c08851e6ff7..ec023299b324b 100644 --- a/src/query/ast/src/parser/statement.rs +++ b/src/query/ast/src/parser/statement.rs @@ -1131,6 +1131,12 @@ pub fn statement_body(i: Input) -> IResult { }) }, ); + let create_table_partition_by = map( + rule! { + #table_option ~ PARTITION ~ ^BY ~ ^"(" ~ ^#comma_separated_list1(expr) ~ ^")" + }, + |(table_options, _, _, _, exprs, _)| (table_options, exprs), + ); let create_table = map_res( rule! { CREATE ~ ( OR ~ ^REPLACE )? ~ (TEMP| TEMPORARY|TRANSIENT)? ~ TABLE ~ ( IF ~ ^NOT ~ ^EXISTS )? @@ -1138,9 +1144,9 @@ pub fn statement_body(i: Input) -> IResult { ~ #create_table_source? ~ ( #engine )? ~ ( #uri_location )? + ~ #create_table_partition_by? ~ ( CLUSTER ~ ^BY ~ LINEAR? ~ ^"(" ~ ^#comma_separated_list1(expr) ~ ^")" )? ~ ( #table_option )? - ~ ( PARTITION ~ ^BY ~ ^"(" ~ ^#comma_separated_list1(ident) ~ ^")" )? ~ ( PROPERTIES ~ #connection_options )? ~ ( AS ~ ^#query )? }, @@ -1154,9 +1160,9 @@ pub fn statement_body(i: Input) -> IResult { source, engine, uri_location, + opt_partition_by, opt_cluster_by, opt_table_options, - opt_iceberg_table_partition_by, opt_table_properties, opt_as_query, )| { @@ -1168,6 +1174,10 @@ pub fn statement_body(i: Input) -> IResult { Some(TEMP) | Some(TEMPORARY) => TableType::Temporary, _ => unreachable!(), }; + let (mut table_options, partition_by) = opt_partition_by + .map(|(options, exprs)| (options, Some(exprs))) + .unwrap_or_default(); + table_options.extend(opt_table_options.unwrap_or_default()); Ok(Statement::CreateTable(CreateTableStmt { create_option, catalog, @@ -1179,9 +1189,8 @@ pub fn statement_body(i: Input) -> IResult { cluster_by: opt_cluster_by.map(|(_, _, _, _, exprs, _)| ClusterOption { cluster_exprs: exprs, }), - table_options: opt_table_options.unwrap_or_default(), - iceberg_table_partition: opt_iceberg_table_partition_by - .map(|(_, _, _, cols, _)| cols), + table_options, + partition_by, table_properties: opt_table_properties.map(|(_, properties)| properties), as_query: opt_as_query.map(|(_, query)| Box::new(query)), table_type, diff --git a/src/query/ast/src/visit/statement_table.rs b/src/query/ast/src/visit/statement_table.rs index b1ed8167cb774..ed104ab73552d 100644 --- a/src/query/ast/src/visit/statement_table.rs +++ b/src/query/ast/src/visit/statement_table.rs @@ -149,9 +149,9 @@ impl Walk for CreateTableStmt { if let Some(cluster_by) = &self.cluster_by { try_walk!(cluster_by.walk(visitor)); } - if let Some(partitions) = &self.iceberg_table_partition { - for ident in partitions { - try_walk!(ident.walk(visitor)); + if let Some(partitions) = &self.partition_by { + for expr in partitions { + try_walk!(expr.walk(visitor)); } } if let Some(as_query) = &self.as_query { @@ -173,9 +173,9 @@ impl WalkMut for CreateTableStmt { if let Some(cluster_by) = &mut self.cluster_by { try_walk!(cluster_by.walk_mut(visitor)); } - if let Some(partitions) = &mut self.iceberg_table_partition { - for ident in partitions { - try_walk!(ident.walk_mut(visitor)); + if let Some(partitions) = &mut self.partition_by { + for expr in partitions { + try_walk!(expr.walk_mut(visitor)); } } if let Some(as_query) = &mut self.as_query { diff --git a/src/query/ast/tests/it/parser.rs b/src/query/ast/tests/it/parser.rs index 4e40ba8bf5e68..09cf7dbc0fab0 100644 --- a/src/query/ast/tests/it/parser.rs +++ b/src/query/ast/tests/it/parser.rs @@ -1393,6 +1393,32 @@ fn test_file_format_trim_space_option() { assert!(displayed.contains("TRIM_SPACE = true".to_uppercase().as_str())); } +#[test] +fn test_create_table_options_before_partition_by() { + let cases = [ + "CREATE TABLE t(c INT) ENGINE=ICEBERG LOCATION='s3://bucket/path' CONNECTION_NAME='conn' PARTITION BY (c)", + "CREATE TABLE iceberg.db.t(c INT) LOCATION='s3://bucket/path' PARTITION BY (c)", + "CREATE TABLE t(a INT) ENGINE=FUSE ROW_PER_BLOCK=1 PARTITION BY (a)", + "CREATE TABLE t(a INT) ROW_PER_BLOCK=1 PARTITION BY (a)", + ]; + for sql in cases { + let tokens = tokenize_sql(sql).unwrap(); + let (stmt, _) = parse_sql(&tokens, Dialect::PostgreSQL).unwrap(); + + let displayed = stmt.to_string(); + let displayed_uppercase = displayed.to_uppercase(); + let partition_pos = displayed_uppercase.find("PARTITION BY").unwrap(); + for option in ["LOCATION", "CONNECTION_NAME", "ROW_PER_BLOCK"] { + if let Some(option_pos) = displayed_uppercase.find(option) { + assert!(partition_pos < option_pos); + } + } + + let tokens = tokenize_sql(&displayed).unwrap(); + parse_sql(&tokens, Dialect::PostgreSQL).unwrap(); + } +} + #[test] fn test_stage_local_filesystem_uri_errors() { let cases = [ diff --git a/src/query/ast/tests/it/testdata/stmt-error.txt b/src/query/ast/tests/it/testdata/stmt-error.txt index 1b43b10fccce4..3f07b88b0b55b 100644 --- a/src/query/ast/tests/it/testdata/stmt-error.txt +++ b/src/query/ast/tests/it/testdata/stmt-error.txt @@ -94,8 +94,10 @@ error: --> SQL:1:55 | 1 | create table a (c1 decimal(38), c2 int) partition by (); - | ------ ^ unexpected `)`, expecting , , or `IDENTIFIER` - | | + | ------ ^ + | | | + | | expecting ``, '', '', '', 'TRUE', 'FALSE', or more ... + | | while parsing expression | while parsing `CREATE [OR REPLACE] TABLE [IF NOT EXISTS] [.] [] []` diff --git a/src/query/ast/tests/it/testdata/stmt.txt b/src/query/ast/tests/it/testdata/stmt.txt index 8165f7d9e87be..9600ad2879c21 100644 --- a/src/query/ast/tests/it/testdata/stmt.txt +++ b/src/query/ast/tests/it/testdata/stmt.txt @@ -2420,7 +2420,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -2477,7 +2477,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -2549,23 +2549,45 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: Some( + partition_by: Some( [ - Identifier { + ColumnRef { span: Some( 54..56, ), - name: "c1", - quote: None, - ident_type: None, + column: ColumnRef { + database: None, + table: None, + column: Name( + Identifier { + span: Some( + 54..56, + ), + name: "c1", + quote: None, + ident_type: None, + }, + ), + }, }, - Identifier { + ColumnRef { span: Some( 58..60, ), - name: "c2", - quote: None, - ident_type: None, + column: ColumnRef { + database: None, + table: None, + column: Name( + Identifier { + span: Some( + 58..60, + ), + name: "c2", + quote: None, + ident_type: None, + }, + ), + }, }, ], ), @@ -2630,7 +2652,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -2684,7 +2706,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -2753,7 +2775,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -2865,7 +2887,7 @@ CreateTable( }, ), table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -2956,7 +2978,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -3047,7 +3069,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: Some( Query { @@ -3305,7 +3327,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -3414,7 +3436,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -3497,7 +3519,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -3662,7 +3684,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -3802,7 +3824,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -3916,7 +3938,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -4030,7 +4052,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -4092,7 +4114,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -4138,7 +4160,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -4215,7 +4237,7 @@ CreateTable( ), cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -5586,7 +5608,7 @@ CreateTable( "compression": "zstd", "storage_format": "native", }, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -5817,7 +5839,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -7511,7 +7533,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -7601,7 +7623,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -7670,7 +7692,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -7735,7 +7757,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -7766,7 +7788,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: Some( Query { @@ -7908,7 +7930,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -7977,7 +7999,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -20412,7 +20434,7 @@ CreateTable( table_options: { "comment": "Comment types type speedily ' \\\\ '' Fun!", }, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Normal, @@ -20609,7 +20631,7 @@ CreateTable( uri_location: None, cluster_by: None, table_options: {}, - iceberg_table_partition: None, + partition_by: None, table_properties: None, as_query: None, table_type: Temporary, diff --git a/src/query/expression/src/constant_folder.rs b/src/query/expression/src/constant_folder.rs index 2322640f365f5..0a4b909ba6329 100644 --- a/src/query/expression/src/constant_folder.rs +++ b/src/query/expression/src/constant_folder.rs @@ -871,7 +871,7 @@ impl<'a, Index: ColumnIndex> ConstantFolder<'a, Index> { // Extract constraints from each expression for arg in args { - if let Some(constraint) = self.extract_range_constraint(arg) { + if let Some(constraint) = RangeConstraint::try_from_expr(arg) { column_constraints .entry(constraint.column_id.clone()) .or_default() @@ -902,54 +902,6 @@ impl<'a, Index: ColumnIndex> ConstantFolder<'a, Index> { None // No conclusive mutual exclusion found } - /// Extract range constraint from a comparison expression - fn extract_range_constraint(&self, expr: &Expr) -> Option> { - if let Expr::FunctionCall(FunctionCall { function, args, .. }) = expr { - if args.len() != 2 { - return None; - } - - let op = function.signature.name.as_str(); - if !matches!(op, "gt" | "gte" | "lt" | "lte" | "eq" | "noteq") { - return None; - } - - // Try both orders: column op constant and constant op column - if let (Some(column_ref), Some(constant)) = - (args[0].as_column_ref(), args[1].as_constant()) - { - return Some(RangeConstraint { - column_id: column_ref.id.clone(), - data_type: column_ref.data_type.clone(), - operator: op.to_string(), - constant: constant.scalar.clone(), - is_flipped: false, - }); - } else if let (Some(constant), Some(column_ref)) = - (args[0].as_constant(), args[1].as_column_ref()) - { - // Flip the operator for constant op column - let flipped_op = match op { - "gt" => "lt", - "gte" => "lte", - "lt" => "gt", - "lte" => "gte", - "eq" => "eq", - "noteq" => "noteq", - _ => return None, - }; - return Some(RangeConstraint { - column_id: column_ref.id.clone(), - data_type: column_ref.data_type.clone(), - operator: flipped_op.to_string(), - constant: constant.scalar.clone(), - is_flipped: true, - }); - } - } - None - } - /// Check if two range constraints are mutually exclusive pub fn are_constraints_mutually_exclusive( &self, @@ -1100,3 +1052,58 @@ pub struct RangeConstraint { pub constant: Scalar, pub is_flipped: bool, // true if original was constant op column } + +impl RangeConstraint { + /// Extracts a normalized column-to-constant comparison. Comparisons with + /// the constant on the left are flipped so the column is always the lhs. + pub fn try_from_expr(expr: &Expr) -> Option { + let Expr::FunctionCall(call) = expr else { + return None; + }; + Self::try_from_function_call(call) + } + + pub fn try_from_function_call(call: &FunctionCall) -> Option { + let FunctionCall { function, args, .. } = call; + if args.len() != 2 { + return None; + } + + let op = function.signature.name.as_str(); + if !matches!(op, "gt" | "gte" | "lt" | "lte" | "eq" | "noteq") { + return None; + } + + if let (Some(column_ref), Some(constant)) = (args[0].as_column_ref(), args[1].as_constant()) + { + return Some(Self { + column_id: column_ref.id.clone(), + data_type: column_ref.data_type.clone(), + operator: op.to_string(), + constant: constant.scalar.clone(), + is_flipped: false, + }); + } + + let (Some(constant), Some(column_ref)) = (args[0].as_constant(), args[1].as_column_ref()) + else { + return None; + }; + let operator = match op { + "gt" => "lt", + "gte" => "lte", + "lt" => "gt", + "lte" => "gte", + "eq" => "eq", + "noteq" => "noteq", + _ => unreachable!(), + }; + Some(Self { + column_id: column_ref.id.clone(), + data_type: column_ref.data_type.clone(), + operator: operator.to_string(), + constant: constant.scalar.clone(), + is_flipped: true, + }) + } +} diff --git a/src/query/functions/Cargo.toml b/src/query/functions/Cargo.toml index 1b3dd1fcd622b..f6e2137a6aa29 100644 --- a/src/query/functions/Cargo.toml +++ b/src/query/functions/Cargo.toml @@ -41,6 +41,7 @@ jsonb = { workspace = true } libm = { workspace = true } match-template = { workspace = true } md-5 = { workspace = true } +murmur3 = { workspace = true } naive-cityhash = { workspace = true } num-traits = { workspace = true } proptest = { workspace = true } diff --git a/src/query/functions/src/scalars/hash.rs b/src/query/functions/src/scalars/hash.rs index c93ed7c5a497a..b6570d2a9a5dd 100644 --- a/src/query/functions/src/scalars/hash.rs +++ b/src/query/functions/src/scalars/hash.rs @@ -102,6 +102,8 @@ pub fn register(registry: &mut FunctionRegistry) { }); } + register_bucket(registry); + // Could be used in exchange registry .scalar_builder("siphash64") @@ -222,6 +224,99 @@ pub fn register(registry: &mut FunctionRegistry) { ); } +fn register_bucket(registry: &mut FunctionRegistry) { + registry + .register_passthrough_nullable_2_arg::, StringType, NumberType, _, _>( + "bucket", + |_, _, _| FunctionDomain::MayThrow, + vectorize_with_builder_2_arg::, StringType, NumberType>( + |buckets, value, output, ctx| { + output.push(bucket(buckets, value.as_bytes(), output.len(), ctx)); + }, + ), + ); + registry + .register_passthrough_nullable_2_arg::, DateType, NumberType, _, _>( + "bucket", + |_, _, _| FunctionDomain::MayThrow, + vectorize_with_builder_2_arg::, DateType, NumberType>( + |buckets, value, output, ctx| { + output.push(bucket(buckets, &value.to_le_bytes(), output.len(), ctx)); + }, + ), + ); + registry.register_passthrough_nullable_2_arg::< + NumberType, + TimestampType, + NumberType, + _, + _, + >( + "bucket", + |_, _, _| FunctionDomain::MayThrow, + vectorize_with_builder_2_arg::, TimestampType, NumberType>( + |buckets, value, output, ctx| { + output.push(bucket( + buckets, + &value.to_le_bytes(), + output.len(), + ctx, + )); + }, + ), + ); + for num_type in ALL_INTEGER_TYPES { + with_integer_mapped_type!(|NUM_TYPE| match num_type { + NumberDataType::NUM_TYPE => { + registry.register_passthrough_nullable_2_arg::< + NumberType, + NumberType, + NumberType, + _, + _, + >( + "bucket", + |_, _, _| FunctionDomain::MayThrow, + vectorize_with_builder_2_arg::< + NumberType, + NumberType, + NumberType, + >(|buckets, value, output, ctx| { + let value: i64 = value.as_(); + output.push(bucket( + buckets, + &value.to_le_bytes(), + output.len(), + ctx, + )); + }), + ); + } + _ => unreachable!(), + }); + } +} + +fn bucket( + buckets: u64, + bytes: &[u8], + row: usize, + ctx: &mut databend_common_expression::EvalContext, +) -> i32 { + let Ok(buckets) = i32::try_from(buckets) else { + ctx.set_error(row, "bucket count must be between 1 and 2147483647"); + return 0; + }; + if buckets == 0 { + ctx.set_error(row, "bucket count must be greater than zero"); + return 0; + } + + let mut bytes = bytes; + let hash = murmur3::murmur3_32(&mut bytes, 0).unwrap() as i32; + (hash & i32::MAX) % buckets +} + fn register_simple_domain_type_hash(registry: &mut FunctionRegistry) where for<'a> T::ScalarRef<'a>: DFHash { registry diff --git a/src/query/functions/src/scalars/timestamp/src/datetime.rs b/src/query/functions/src/scalars/timestamp/src/datetime.rs index f4fa722e5acbf..16613c760852b 100644 --- a/src/query/functions/src/scalars/timestamp/src/datetime.rs +++ b/src/query/functions/src/scalars/timestamp/src/datetime.rs @@ -2674,7 +2674,12 @@ fn register_rounder_functions(registry: &mut FunctionRegistry) { ); registry.register_1_arg::( "to_start_of_day", - |_, _| FunctionDomain::Full, + |ctx, domain| { + FunctionDomain::Domain(SimpleDomain { + min: round_timestamp(domain.min, &ctx.tz, Round::Day), + max: round_timestamp(domain.max, &ctx.tz, Round::Day), + }) + }, |val, ctx| round_timestamp(val, &ctx.func_ctx.tz, Round::Day), ); registry.register_1_arg::( diff --git a/src/query/functions/tests/it/scalars/datetime.rs b/src/query/functions/tests/it/scalars/datetime.rs index f5e50327738a9..0bf72746045f9 100644 --- a/src/query/functions/tests/it/scalars/datetime.rs +++ b/src/query/functions/tests/it/scalars/datetime.rs @@ -861,6 +861,10 @@ fn test_rounder_functions(file: &mut impl Write) { ); run_ast(file, "to_start_of_hour(to_timestamp(1630812366))", &[]); run_ast(file, "to_start_of_day(to_timestamp(1630812366))", &[]); + run_ast(file, "to_start_of_day(order_time)", &[( + "order_time", + TimestampType::from_data(vec![0, 86_400_000_001]), + )]); run_ast(file, "time_slot(to_timestamp(1630812366))", &[]); run_ast(file, "to_monday(to_timestamp(1630812366))", &[]); run_ast(file, "to_start_of_week(to_timestamp(1630812366))", &[]); diff --git a/src/query/functions/tests/it/scalars/hash.rs b/src/query/functions/tests/it/scalars/hash.rs index 0a8439e7a39f4..a12a000b87bf3 100644 --- a/src/query/functions/tests/it/scalars/hash.rs +++ b/src/query/functions/tests/it/scalars/hash.rs @@ -33,6 +33,19 @@ fn test_hash() { test_siphash64(file); test_xxhash64(file); test_xxhash32(file); + test_bucket(file); +} + +fn test_bucket(file: &mut impl Write) { + run_ast(file, "bucket(16, 34)", &[]); + run_ast(file, "bucket(16, 'iceberg')", &[]); + run_ast(file, "bucket(16, to_date(10000))", &[]); + run_ast(file, "bucket(16, to_datetime(100000))", &[]); + run_ast(file, "bucket(0, 1)", &[]); + run_ast(file, "bucket(8, a)", &[( + "a", + Int64Type::from_data(vec![1, 2, 3, 4]), + )]); } fn test_md5(file: &mut impl Write) { diff --git a/src/query/functions/tests/it/scalars/testdata/datetime.txt b/src/query/functions/tests/it/scalars/testdata/datetime.txt index 1474a23c5eaf9..a8939cd110dfa 100644 --- a/src/query/functions/tests/it/scalars/testdata/datetime.txt +++ b/src/query/functions/tests/it/scalars/testdata/datetime.txt @@ -3951,6 +3951,27 @@ output domain : {1630800000000000..=1630800000000000} output : '2021-09-05 00:00:00.000000' +ast : to_start_of_day(order_time) +raw expr : to_start_of_day(order_time::Timestamp) +checked expr : to_start_of_day(order_time) +evaluation: ++--------+------------------------------+------------------------------+ +| | order_time | Output | ++--------+------------------------------+------------------------------+ +| Type | Timestamp | Timestamp | +| Domain | {0..=86400000001} | {0..=86400000000} | +| Row 0 | '1970-01-01 00:00:00.000000' | '1970-01-01 00:00:00.000000' | +| Row 1 | '1970-01-02 00:00:00.000001' | '1970-01-02 00:00:00.000000' | ++--------+------------------------------+------------------------------+ +evaluation (internal): ++------------+-------------------------------------+ +| Column | Data | ++------------+-------------------------------------+ +| order_time | Column(Timestamp([0, 86400000001])) | +| Output | Timestamp([0, 86400000000]) | ++------------+-------------------------------------+ + + ast : time_slot(to_timestamp(1630812366)) raw expr : time_slot(to_timestamp(1630812366)) checked expr : time_slot(CAST(CAST(1630812366_u32 AS Int64) AS Timestamp)) diff --git a/src/query/functions/tests/it/scalars/testdata/function_list.txt b/src/query/functions/tests/it/scalars/testdata/function_list.txt index 5879be79072ec..66a18bfe6c8f9 100644 --- a/src/query/functions/tests/it/scalars/testdata/function_list.txt +++ b/src/query/functions/tests/it/scalars/testdata/function_list.txt @@ -934,6 +934,28 @@ Functions overloads: 1 bitmap_xor(Bitmap NULL, Bitmap NULL) :: Bitmap NULL 0 blake3(String) :: String 1 blake3(String NULL) :: String NULL +0 bucket(UInt64, String) :: Int32 +1 bucket(UInt64 NULL, String NULL) :: Int32 NULL +2 bucket(UInt64, Date) :: Int32 +3 bucket(UInt64 NULL, Date NULL) :: Int32 NULL +4 bucket(UInt64, Timestamp) :: Int32 +5 bucket(UInt64 NULL, Timestamp NULL) :: Int32 NULL +6 bucket(UInt64, UInt8) :: Int32 +7 bucket(UInt64 NULL, UInt8 NULL) :: Int32 NULL +8 bucket(UInt64, UInt16) :: Int32 +9 bucket(UInt64 NULL, UInt16 NULL) :: Int32 NULL +10 bucket(UInt64, UInt32) :: Int32 +11 bucket(UInt64 NULL, UInt32 NULL) :: Int32 NULL +12 bucket(UInt64, UInt64) :: Int32 +13 bucket(UInt64 NULL, UInt64 NULL) :: Int32 NULL +14 bucket(UInt64, Int8) :: Int32 +15 bucket(UInt64 NULL, Int8 NULL) :: Int32 NULL +16 bucket(UInt64, Int16) :: Int32 +17 bucket(UInt64 NULL, Int16 NULL) :: Int32 NULL +18 bucket(UInt64, Int32) :: Int32 +19 bucket(UInt64 NULL, Int32 NULL) :: Int32 NULL +20 bucket(UInt64, Int64) :: Int32 +21 bucket(UInt64 NULL, Int64 NULL) :: Int32 NULL 0 build_bitmap(Array(UInt8 NULL)) :: Bitmap 1 build_bitmap(Array(UInt8 NULL) NULL) :: Bitmap NULL 2 build_bitmap(Array(UInt16 NULL)) :: Bitmap diff --git a/src/query/functions/tests/it/scalars/testdata/hash.txt b/src/query/functions/tests/it/scalars/testdata/hash.txt index f6a46b45593a5..8e4fa33423e8a 100644 --- a/src/query/functions/tests/it/scalars/testdata/hash.txt +++ b/src/query/functions/tests/it/scalars/testdata/hash.txt @@ -851,3 +851,71 @@ evaluation (internal): +--------+----------------------------------+ +ast : bucket(16, 34) +raw expr : bucket(16, 34) +checked expr : bucket(CAST(16_u8 AS UInt64), 34_u8) +optimized expr : 3_i32 +output type : Int32 +output domain : {3..=3} +output : 3 + + +ast : bucket(16, 'iceberg') +raw expr : bucket(16, 'iceberg') +checked expr : bucket(CAST(16_u8 AS UInt64), "iceberg") +optimized expr : 9_i32 +output type : Int32 +output domain : {9..=9} +output : 9 + + +ast : bucket(16, to_date(10000)) +raw expr : bucket(16, to_date(10000)) +checked expr : bucket(CAST(16_u8 AS UInt64), CAST(CAST(10000_u16 AS Int64) AS Date)) +optimized expr : 12_i32 +output type : Int32 +output domain : {12..=12} +output : 12 + + +ast : bucket(16, to_datetime(100000)) +raw expr : bucket(16, to_datetime(100000)) +checked expr : bucket(CAST(16_u8 AS UInt64), CAST(CAST(100000_u32 AS Int64) AS Timestamp)) +optimized expr : 15_i32 +output type : Int32 +output domain : {15..=15} +output : 15 + + +error: + --> SQL:1:1 + | +1 | bucket(0, 1) + | ^^^^^^^^^^^^ bucket count must be greater than zero while evaluating function `bucket(0, 1)` in expr `bucket(CAST(0 AS UInt64), 1)` + + + +ast : bucket(8, a) +raw expr : bucket(8, a::Int64) +checked expr : bucket(CAST(8_u8 AS UInt64), a) +optimized expr : bucket(8_u64, a) +evaluation: ++--------+---------+---------+ +| | a | Output | ++--------+---------+---------+ +| Type | Int64 | Int32 | +| Domain | {1..=4} | Unknown | +| Row 0 | 1 | 4 | +| Row 1 | 2 | 4 | +| Row 2 | 3 | 3 | +| Row 3 | 4 | 6 | ++--------+---------+---------+ +evaluation (internal): ++--------+-----------------------------+ +| Column | Data | ++--------+-----------------------------+ +| a | Column(Int64([1, 2, 3, 4])) | +| Output | Int32([4, 4, 3, 6]) | ++--------+-----------------------------+ + + diff --git a/src/query/service/src/interpreters/common/table_option_validation.rs b/src/query/service/src/interpreters/common/table_option_validation.rs index d94700283b0a0..74e6602602404 100644 --- a/src/query/service/src/interpreters/common/table_option_validation.rs +++ b/src/query/service/src/interpreters/common/table_option_validation.rs @@ -72,6 +72,7 @@ use databend_storages_common_table_meta::table::OPT_KEY_SEGMENT_FORMAT; use databend_storages_common_table_meta::table::OPT_KEY_STORAGE_FORMAT; use databend_storages_common_table_meta::table::OPT_KEY_TABLE_COMPRESSION; use databend_storages_common_table_meta::table::OPT_KEY_TEMP_PREFIX; +use databend_storages_common_table_meta::table::OPT_KEY_WRITE_DISTRIBUTION_MODE; pub use databend_storages_common_table_meta::table::analyze_count_min_sketch_error_rate_from_options; pub use databend_storages_common_table_meta::table::analyze_top_n_size_from_options; use log::error; @@ -101,6 +102,7 @@ pub static CREATE_FUSE_OPTIONS: LazyLock> = LazyLock::new( r.insert(OPT_KEY_DATABASE_ID); r.insert(OPT_KEY_COMMENT); r.insert(OPT_KEY_CHANGE_TRACKING); + r.insert(OPT_KEY_WRITE_DISTRIBUTION_MODE); r.insert(OPT_KEY_ENGINE); diff --git a/src/query/service/src/interpreters/interpreter_insert.rs b/src/query/service/src/interpreters/interpreter_insert.rs index 8ee499146a54c..0ad41dcc6115e 100644 --- a/src/query/service/src/interpreters/interpreter_insert.rs +++ b/src/query/service/src/interpreters/interpreter_insert.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeSet; use std::sync::Arc; use chrono::Duration; @@ -22,21 +23,29 @@ use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_expression::DataBlock; use databend_common_expression::DataSchema; +use databend_common_expression::DataSchemaRef; use databend_common_expression::FromData; +use databend_common_expression::RemoteExpr; use databend_common_expression::SendableDataBlockStream; use databend_common_expression::types::UInt64Type; +use databend_common_functions::BUILTIN_FUNCTIONS; use databend_common_meta_app::schema::Constraint; use databend_common_pipeline::sources::AsyncSourcer; use databend_common_pipeline_transforms::TransformPipelineHelper; #[cfg(feature = "storage-stage")] use databend_common_pipeline_transforms::columns::TransformAddConstColumns; +use databend_common_sql::ColumnBinding; +use databend_common_sql::MetadataRef; use databend_common_sql::NameResolutionContext; use databend_common_sql::binder::ConstraintExprBinder; +use databend_common_sql::executor::physical_plans::FragmentKind; use databend_common_sql::executor::physical_plans::MutationKind; +use databend_common_sql::executor::physical_plans::SortDesc; use databend_common_sql::plans::Insert; use databend_common_sql::plans::InsertInputSource; use databend_common_sql::plans::InsertValue; use databend_common_sql::plans::Plan; +use databend_common_storages_fuse::FuseTable; use databend_common_storages_fuse::operations::TransformConstraintVerify; #[cfg(feature = "storage-stage")] use databend_query_storage_stage_support::build_streaming_load_pipeline; @@ -49,12 +58,15 @@ use crate::interpreters::InterpreterPtr; use crate::interpreters::common::check_deduplicate_label; use crate::interpreters::common::dml_build_update_stream_req; use crate::physical_plans::DistributedInsertSelect; +use crate::physical_plans::EvalScalar; use crate::physical_plans::Exchange; use crate::physical_plans::IPhysicalPlan; use crate::physical_plans::PhysicalPlan; use crate::physical_plans::PhysicalPlanBuilder; use crate::physical_plans::PhysicalPlanCast; use crate::physical_plans::PhysicalPlanMeta; +use crate::physical_plans::Sort; +use crate::physical_plans::SortStep; use crate::pipelines::PipelineBuildResult; use crate::pipelines::PipelineBuilder; use crate::pipelines::RawValueSource; @@ -95,6 +107,93 @@ impl InsertInterpreter { let cast_needed = select_schema.as_ref() != &DataSchema::from(output_schema.as_ref()); Ok(cast_needed) } + + fn build_hash_distributed_input( + &self, + input: PhysicalPlan, + table: &Arc, + insert_schema: &DataSchemaRef, + select_column_bindings: &[ColumnBinding], + metadata: &MetadataRef, + cast_needed: bool, + ) -> Result { + let Ok(fuse_table) = FuseTable::try_from_table(table.as_ref()) else { + return Ok(input); + }; + if cast_needed || !fuse_table.use_hash_write_distribution() { + return Ok(input); + } + + let Some(partition_info) = fuse_table.partition_pruning_info(self.ctx.clone()) else { + return Ok(input); + }; + let input_schema = input.output_schema()?; + let mut partition_exprs = Vec::with_capacity(partition_info.partition_keys.len()); + for partition_key in partition_info.partition_keys { + let expr = partition_key.as_expr(&BUILTIN_FUNCTIONS); + if expr + .column_refs() + .keys() + .any(|name| insert_schema.index_of(name).is_err()) + { + return Ok(input); + } + partition_exprs.push(expr.project_column_ref(|name| { + let position = insert_schema.index_of(name)?; + let binding = select_column_bindings.get(position).ok_or_else(|| { + ErrorCode::Internal("INSERT SELECT column binding is missing") + })?; + input_schema.index_of(&binding.index.to_string()) + })?); + } + + let input_columns = input_schema.num_fields(); + let mut eval_exprs = Vec::with_capacity(partition_exprs.len()); + let mut hash_keys = Vec::with_capacity(partition_exprs.len()); + let mut order_by = Vec::with_capacity(partition_exprs.len()); + for (index, expr) in partition_exprs.into_iter().enumerate() { + let data_type = expr.data_type().clone(); + let display_name = format!("partition_key_{index}"); + let symbol = metadata + .write() + .add_derived_column(display_name.clone(), data_type.clone()); + eval_exprs.push((expr.as_remote_expr(), symbol)); + hash_keys.push(RemoteExpr::ColumnRef { + span: None, + id: input_columns + index, + data_type, + display_name: display_name.clone(), + }); + order_by.push(SortDesc { + asc: true, + nulls_first: false, + order_by: symbol, + display_name, + }); + } + + let projections: BTreeSet<_> = (0..input_columns + eval_exprs.len()).collect(); + let eval = PhysicalPlan::new(EvalScalar::create(input, eval_exprs, projections, None)); + let exchange = PhysicalPlan::new(Exchange { + meta: PhysicalPlanMeta::new("Exchange"), + input: eval, + kind: FragmentKind::GlobalShuffle, + keys: hash_keys, + ignore_exchange: false, + allow_adjust_parallelism: true, + }); + Ok(PhysicalPlan::new(Sort { + meta: PhysicalPlanMeta::new("Sort"), + input: exchange, + order_by, + limit: None, + step: SortStep::Single, + pre_projection: None, + broadcast_id: None, + enable_fixed_rows: self.ctx.get_settings().get_enable_fixed_rows_sort()?, + stat_info: None, + })) + } } #[async_trait::async_trait] @@ -230,28 +329,47 @@ impl Interpreter for InsertInterpreter { && let Some(exchange) = Exchange::from_physical_plan(&select_plan) { // insert can be dispatched to different nodes if table support_distributed_insert - let input = exchange.input.clone(); + let insert_schema = self.plan.dest_schema(); + let cast_needed = self.check_schema_cast(plan)?; + let input = self.build_hash_distributed_input( + exchange.input.clone(), + &table1, + &insert_schema, + &select_column_bindings, + metadata, + cast_needed, + )?; exchange.derive(vec![PhysicalPlan::new(DistributedInsertSelect { input, table_info, select_schema: plan.schema(), select_column_bindings, - insert_schema: self.plan.dest_schema(), - cast_needed: self.check_schema_cast(plan)?, + insert_schema, + cast_needed, table_meta_timestamps, meta: PhysicalPlanMeta::new("DistributedInsertSelect"), })]) } else { // insert should wait until all nodes finished + let insert_schema = self.plan.dest_schema(); + let cast_needed = self.check_schema_cast(plan)?; + let input = self.build_hash_distributed_input( + select_plan, + &table1, + &insert_schema, + &select_column_bindings, + metadata, + cast_needed, + )?; PhysicalPlan::new(DistributedInsertSelect { // TODO: we reuse the id of other plan here, // which is not correct. We should generate a new id for insert. - input: select_plan, + input, table_info, select_schema: plan.schema(), select_column_bindings, - insert_schema: self.plan.dest_schema(), - cast_needed: self.check_schema_cast(plan)?, + insert_schema, + cast_needed, table_meta_timestamps, meta: PhysicalPlanMeta::new("DistributedInsertSelect"), }) diff --git a/src/query/service/src/interpreters/interpreter_table_create.rs b/src/query/service/src/interpreters/interpreter_table_create.rs index bfe7c753e8f3e..c12b94a301904 100644 --- a/src/query/service/src/interpreters/interpreter_table_create.rs +++ b/src/query/service/src/interpreters/interpreter_table_create.rs @@ -59,6 +59,7 @@ use databend_storages_common_table_meta::meta::TableSnapshot; use databend_storages_common_table_meta::meta::Versioned; use databend_storages_common_table_meta::table::OPT_KEY_COMMENT; use databend_storages_common_table_meta::table::OPT_KEY_ENABLE_COPY_DEDUP_FULL_PATH; +use databend_storages_common_table_meta::table::OPT_KEY_PARTITION_BY; use databend_storages_common_table_meta::table::OPT_KEY_SEGMENT_FORMAT; use databend_storages_common_table_meta::table::OPT_KEY_SNAPSHOT_LOCATION; use databend_storages_common_table_meta::table::OPT_KEY_STORAGE_FORMAT; @@ -507,7 +508,8 @@ impl CreateTableInterpreter { for table_option in table_meta.options.iter() { let key = table_option.0.to_lowercase(); - if !is_valid_create_opt(&key, &self.plan.engine) { + // PARTITION BY is normalized and inserted by the binder as internal metadata. + if key != OPT_KEY_PARTITION_BY && !is_valid_create_opt(&key, &self.plan.engine) { let msg = format!( "table option {key} is invalid for create table statement with engine {}", self.plan.engine diff --git a/src/query/service/src/interpreters/interpreter_table_drop_column.rs b/src/query/service/src/interpreters/interpreter_table_drop_column.rs index d0d8b71781df8..65e4d345b2ad0 100644 --- a/src/query/service/src/interpreters/interpreter_table_drop_column.rs +++ b/src/query/service/src/interpreters/interpreter_table_drop_column.rs @@ -29,6 +29,7 @@ use databend_common_storages_stream::stream_table::STREAM_ENGINE; use databend_storages_common_table_meta::table::OPT_KEY_ANALYZE_FREQUENCY_COLUMNS; use databend_storages_common_table_meta::table::OPT_KEY_APPROX_DISTINCT_COLUMNS; use databend_storages_common_table_meta::table::OPT_KEY_BLOOM_INDEX_COLUMNS; +use databend_storages_common_table_meta::table::OPT_KEY_PARTITION_BY; use crate::interpreters::Interpreter; use crate::interpreters::common::check_referenced_computed_columns; @@ -105,6 +106,15 @@ impl Interpreter for DropTableColumnInterpreter { ))); } } + if let Some(partition_key) = table_info.options().get(OPT_KEY_PARTITION_BY) { + let referenced = cluster_key_referenced_columns(partition_key)?; + if referenced.contains(self.plan.column.as_str()) { + return Err(ErrorCode::AlterTableError(format!( + "Cannot drop column '{}' because it is referenced by partition key {}", + self.plan.column, partition_key + ))); + } + } if table_info.meta.is_column_reference_policy(&field.column_id) { return Err(ErrorCode::AlterTableError(format!( diff --git a/src/query/service/src/interpreters/interpreter_table_modify_column.rs b/src/query/service/src/interpreters/interpreter_table_modify_column.rs index 7f53f0af83766..a341cb68b66bd 100644 --- a/src/query/service/src/interpreters/interpreter_table_modify_column.rs +++ b/src/query/service/src/interpreters/interpreter_table_modify_column.rs @@ -64,6 +64,7 @@ use databend_storages_common_table_meta::readers::snapshot_reader::TableSnapshot use databend_storages_common_table_meta::table::OPT_KEY_ANALYZE_FREQUENCY_COLUMNS; use databend_storages_common_table_meta::table::OPT_KEY_APPROX_DISTINCT_COLUMNS; use databend_storages_common_table_meta::table::OPT_KEY_BLOOM_INDEX_COLUMNS; +use databend_storages_common_table_meta::table::OPT_KEY_PARTITION_BY; use crate::interpreters::Interpreter; use crate::interpreters::common::check_referenced_computed_columns; @@ -271,6 +272,15 @@ impl ModifyTableColumnInterpreter { } } } + if let Some(partition_key) = table.options().get(OPT_KEY_PARTITION_BY) { + let referenced = cluster_key_referenced_columns(partition_key)?; + if referenced.iter().any(|v| modified_cols.contains(v)) { + return Err(ErrorCode::AlterTableError(format!( + "Cannot modify column data type because it is referenced by partition key '{}'", + partition_key + ))); + } + } } let catalog_name = table_info.catalog(); diff --git a/src/query/service/src/interpreters/interpreter_table_rename_column.rs b/src/query/service/src/interpreters/interpreter_table_rename_column.rs index ac985165aaa18..abfa779ff6cf8 100644 --- a/src/query/service/src/interpreters/interpreter_table_rename_column.rs +++ b/src/query/service/src/interpreters/interpreter_table_rename_column.rs @@ -26,6 +26,7 @@ use databend_common_storages_stream::stream_table::STREAM_ENGINE; use databend_storages_common_table_meta::table::OPT_KEY_ANALYZE_FREQUENCY_COLUMNS; use databend_storages_common_table_meta::table::OPT_KEY_APPROX_DISTINCT_COLUMNS; use databend_storages_common_table_meta::table::OPT_KEY_BLOOM_INDEX_COLUMNS; +use databend_storages_common_table_meta::table::OPT_KEY_PARTITION_BY; use crate::interpreters::Interpreter; use crate::interpreters::common::check_referenced_computed_columns; @@ -130,6 +131,16 @@ impl Interpreter for RenameTableColumnInterpreter { &self.plan.new_column, )?; } + if let Some(value) = opts.get_mut(OPT_KEY_PARTITION_BY) + && let Some(new_partition_key) = rename_column_in_cluster_key( + self.ctx.as_ref(), + value, + &self.plan.old_column, + &self.plan.new_column, + )? + { + *value = new_partition_key; + } if let Some(value) = opts.get_mut(OPT_KEY_ANALYZE_FREQUENCY_COLUMNS) { rename_column_in_comma_separated_ident( self.ctx.as_ref(), diff --git a/src/query/service/src/interpreters/interpreter_table_set_options.rs b/src/query/service/src/interpreters/interpreter_table_set_options.rs index 56c03defced62..2131237519cfc 100644 --- a/src/query/service/src/interpreters/interpreter_table_set_options.rs +++ b/src/query/service/src/interpreters/interpreter_table_set_options.rs @@ -48,10 +48,13 @@ use databend_storages_common_table_meta::table::OPT_KEY_CHANGE_TRACKING; use databend_storages_common_table_meta::table::OPT_KEY_CHANGE_TRACKING_BEGIN_VER; use databend_storages_common_table_meta::table::OPT_KEY_CLUSTER_TYPE; use databend_storages_common_table_meta::table::OPT_KEY_DATABASE_ID; +use databend_storages_common_table_meta::table::OPT_KEY_PARTITION_BY; use databend_storages_common_table_meta::table::OPT_KEY_SEGMENT_FORMAT; use databend_storages_common_table_meta::table::OPT_KEY_SNAPSHOT_LOCATION; use databend_storages_common_table_meta::table::OPT_KEY_STORAGE_FORMAT; use databend_storages_common_table_meta::table::OPT_KEY_TEMP_PREFIX; +use databend_storages_common_table_meta::table::OPT_KEY_WRITE_DISTRIBUTION_MODE; +use databend_storages_common_table_meta::table::WriteDistributionMode; use log::error; use crate::interpreters::Interpreter; @@ -155,6 +158,13 @@ impl Interpreter for SetOptionsInterpreter { OPT_KEY_CLUSTER_TYPE ))); } + if self.plan.set_options.contains_key(OPT_KEY_PARTITION_BY) { + error!("{}", &error_str); + return Err(ErrorCode::TableOptionInvalid(format!( + "can't change {} for alter table statement", + OPT_KEY_PARTITION_BY + ))); + } // Same as settings of FUSE_OPT_KEY_ENABLE_AUTO_VACUUM, expect value type is unsigned integer is_valid_option_of_type::(&self.plan.set_options, FUSE_OPT_KEY_ENABLE_AUTO_VACUUM)?; @@ -171,6 +181,17 @@ impl Interpreter for SetOptionsInterpreter { .get_table(&self.ctx.get_tenant(), database, table_name) .await?; + if let Some(mode) = self.plan.set_options.get(OPT_KEY_WRITE_DISTRIBUTION_MODE) { + let mode = mode.parse::()?; + if mode == WriteDistributionMode::Hash + && !table.options().contains_key(OPT_KEY_PARTITION_BY) + { + return Err(ErrorCode::TableOptionInvalid(format!( + "{OPT_KEY_WRITE_DISTRIBUTION_MODE}='hash' requires PARTITION BY" + ))); + } + } + let engine = Engine::from(table.engine()); for table_option in self.plan.set_options.iter() { let key = table_option.0.to_lowercase(); diff --git a/src/query/service/src/interpreters/interpreter_table_show_create.rs b/src/query/service/src/interpreters/interpreter_table_show_create.rs index cc4e40004942c..1dd59f200daa4 100644 --- a/src/query/service/src/interpreters/interpreter_table_show_create.rs +++ b/src/query/service/src/interpreters/interpreter_table_show_create.rs @@ -34,6 +34,7 @@ use databend_common_storages_basic::view_table::VIEW_ENGINE; use databend_common_storages_fuse::FUSE_OPT_KEY_ATTACH_COLUMN_IDS; use databend_common_storages_stream::stream_table::STREAM_ENGINE; use databend_common_storages_stream::stream_table::StreamTable; +use databend_storages_common_table_meta::table::OPT_KEY_PARTITION_BY; use databend_storages_common_table_meta::table::OPT_KEY_STORAGE_PREFIX; use databend_storages_common_table_meta::table::OPT_KEY_TABLE_ATTACHED_DATA_URI; use databend_storages_common_table_meta::table::OPT_KEY_TEMP_PREFIX; @@ -286,6 +287,21 @@ impl ShowCreateTableInterpreter { let table_engine = format!(") ENGINE={}", engine); table_create_sql.push_str(table_engine.as_str()); + if let Some(partition_keys_str) = table_info.options().get(OPT_KEY_PARTITION_BY) { + let mut exprs = parse_cluster_key_exprs(partition_keys_str)?; + let mut normalizer = ClusterKeyNormalizer { + force_quoted_ident, + unquoted_ident_case_sensitive, + quoted_ident_case_sensitive, + sql_dialect, + }; + for expr in exprs.iter_mut() { + expr.drive_mut(&mut normalizer); + } + let partition_keys_str = exprs.into_iter().map(|e| format!("{e:#}")).join(", "); + table_create_sql.push_str(format!(" PARTITION BY ({partition_keys_str})").as_str()); + } + if let Some(cluster_keys_str) = table_info.meta.cluster_key_str() { let mut exprs = parse_cluster_key_exprs(cluster_keys_str)?; let mut normalizer = ClusterKeyNormalizer { diff --git a/src/query/service/src/physical_plans/mod.rs b/src/query/service/src/physical_plans/mod.rs index 8f3273c3d7dde..62e4ca2df0b0c 100644 --- a/src/query/service/src/physical_plans/mod.rs +++ b/src/query/service/src/physical_plans/mod.rs @@ -103,6 +103,7 @@ pub use physical_replace_into::ReplaceInto; pub use physical_row_fetch::RowFetch; pub use physical_sequence::*; pub use physical_sort::Sort; +pub use physical_sort::SortStep; pub use physical_spatial_join::PhysicalSpatialJoin; pub use physical_table_scan::TableScan; pub use physical_udf::UdfFunctionDesc; diff --git a/src/query/service/src/physical_plans/physical_column_mutation.rs b/src/query/service/src/physical_plans/physical_column_mutation.rs index b4c90345ce8cd..af4fe752ec01f 100644 --- a/src/query/service/src/physical_plans/physical_column_mutation.rs +++ b/src/query/service/src/physical_plans/physical_column_mutation.rs @@ -193,7 +193,7 @@ impl IPhysicalPlan for ColumnMutation { let input_schema = DataSchema::from(table.schema_with_stream()).into(); table.get_cluster_stats_gen(builder.ctx.clone(), 0, block_thresholds, input_schema)? } else { - table.cluster_gen_for_append( + table.cluster_gen_for_update( builder.ctx.clone(), &mut builder.main_pipeline, block_thresholds, diff --git a/src/query/service/src/physical_plans/physical_commit_sink.rs b/src/query/service/src/physical_plans/physical_commit_sink.rs index cac8d733ce4f1..fde0e57a43bf8 100644 --- a/src/query/service/src/physical_plans/physical_commit_sink.rs +++ b/src/query/service/src/physical_plans/physical_commit_sink.rs @@ -149,7 +149,7 @@ impl IPhysicalPlan for CommitSink { } CommitType::Mutation { kind, merge_meta } => { if *merge_meta { - let cluster_key_id = table.cluster_key_id(); + let cluster_key_id = table.physical_cluster_key_id(); builder.main_pipeline.add_accumulating_transformer(|| { TransformMergeCommitMeta::create(cluster_key_id) }); diff --git a/src/query/service/src/physical_plans/physical_compact_source.rs b/src/query/service/src/physical_plans/physical_compact_source.rs index db28e9857f532..53ffd233f46cd 100644 --- a/src/query/service/src/physical_plans/physical_compact_source.rs +++ b/src/query/service/src/physical_plans/physical_compact_source.rs @@ -97,7 +97,8 @@ impl IPhysicalPlan for CompactSource { let is_lazy = self.parts.partitions_type() == PartInfoType::LazyLevel; let thresholds = table.get_block_thresholds(); - let cluster_key_id = table.cluster_key_id(); + let cluster_key_id = table.physical_cluster_key_id(); + let partition_key_count = table.partition_key_count(); let mut max_threads = builder.settings.get_max_threads()? as usize; if is_lazy { @@ -125,6 +126,7 @@ impl IPhysicalPlan for CompactSource { ctx.clone(), dal.clone(), cluster_key_id, + partition_key_count, thresholds, lazy_parts, ) diff --git a/src/query/service/src/physical_plans/physical_multi_table_insert.rs b/src/query/service/src/physical_plans/physical_multi_table_insert.rs index 22b4091a8ebf6..a8da67f856f56 100644 --- a/src/query/service/src/physical_plans/physical_multi_table_insert.rs +++ b/src/query/service/src/physical_plans/physical_multi_table_insert.rs @@ -42,6 +42,7 @@ use databend_common_pipeline_transforms::sorts::TransformSortPartial; use databend_common_sql::DefaultExprBinder; use databend_common_storages_fuse::FuseTable; use databend_common_storages_fuse::operations::CommitMultiTableInsert; +use databend_common_storages_fuse::operations::TransformPartitionBy; use databend_common_storages_fuse::operations::TransformVectorCluster; use databend_storages_common_table_meta::meta::TableMetaTimestamps; @@ -685,6 +686,9 @@ impl IPhysicalPlan for ChunkAppendData { let mut sort_builders: Vec = Vec::with_capacity(self.target_tables.len()); let mut sort_num = 0; + let mut partition_builders: Vec = + Vec::with_capacity(self.target_tables.len()); + let mut partition_num = 0; for append_data in self.target_tables.iter() { let table = builder @@ -757,6 +761,20 @@ impl IPhysicalPlan for ChunkAppendData { } else { sort_builders.push(Box::new(builder.dummy_transform_builder())); } + let partition_key_indices: Arc<[_]> = + cluster_stats_gen.cluster_key_index[..cluster_stats_gen.partition_key_count].into(); + if !partition_key_indices.is_empty() { + partition_builders.push(Box::new(move |input, output| { + Ok(ProcessorPtr::create(AccumulatingTransformer::create( + input, + output, + TransformPartitionBy::new(partition_key_indices.clone()), + ))) + })); + partition_num += 1; + } else { + partition_builders.push(Box::new(builder.dummy_transform_builder())); + } serialize_block_builders.push(Box::new( builder.with_tid_serialize_block_transform_builder( table, @@ -791,6 +809,12 @@ impl IPhysicalPlan for ChunkAppendData { .add_transforms_by_chunk(sort_builders)?; } + if partition_num > 0 { + builder + .main_pipeline + .add_transforms_by_chunk(partition_builders)?; + } + builder .main_pipeline .add_transforms_by_chunk(serialize_block_builders) diff --git a/src/query/service/tests/it/indexes/inverted_index/pruning.rs b/src/query/service/tests/it/indexes/inverted_index/pruning.rs index f9b0c4db13050..b906e470f7d32 100644 --- a/src/query/service/tests/it/indexes/inverted_index/pruning.rs +++ b/src/query/service/tests/it/indexes/inverted_index/pruning.rs @@ -76,6 +76,7 @@ async fn apply_block_pruning( dal, schema, push_down, + None, bloom_index_cols, vec![], HashSet::new(), diff --git a/src/query/service/tests/it/indexes/spatial_index/pruning.rs b/src/query/service/tests/it/indexes/spatial_index/pruning.rs index 0ac78e317339b..9247b57dfaf63 100644 --- a/src/query/service/tests/it/indexes/spatial_index/pruning.rs +++ b/src/query/service/tests/it/indexes/spatial_index/pruning.rs @@ -74,6 +74,7 @@ async fn apply_block_pruning( dal, schema, push_down, + None, bloom_index_cols, vec![], spatial_index_columns, diff --git a/src/query/service/tests/it/indexes/vector_index/pruning.rs b/src/query/service/tests/it/indexes/vector_index/pruning.rs index ddeaf949b99bb..035b2be52a88f 100644 --- a/src/query/service/tests/it/indexes/vector_index/pruning.rs +++ b/src/query/service/tests/it/indexes/vector_index/pruning.rs @@ -85,6 +85,7 @@ async fn apply_block_pruning( dal, schema, push_down, + None, bloom_index_cols, vec![], HashSet::new(), diff --git a/src/query/service/tests/it/storages/fuse/pruning.rs b/src/query/service/tests/it/storages/fuse/pruning.rs index 653a8728c7099..6366532625f41 100644 --- a/src/query/service/tests/it/storages/fuse/pruning.rs +++ b/src/query/service/tests/it/storages/fuse/pruning.rs @@ -69,6 +69,7 @@ async fn apply_block_pruning( op, schema, push_down, + None, bloom_index_cols, vec![], HashSet::new(), diff --git a/src/query/service/tests/it/storages/fuse/pruning_column_oriented_segment.rs b/src/query/service/tests/it/storages/fuse/pruning_column_oriented_segment.rs index 06401e6dcd23f..723cd1773c5cf 100644 --- a/src/query/service/tests/it/storages/fuse/pruning_column_oriented_segment.rs +++ b/src/query/service/tests/it/storages/fuse/pruning_column_oriented_segment.rs @@ -79,6 +79,7 @@ async fn apply_snapshot_pruning( op, schema.clone(), push_down, + None, bloom_index_cols, vec![], HashSet::new(), diff --git a/src/query/service/tests/it/storages/fuse/pruning_pipeline.rs b/src/query/service/tests/it/storages/fuse/pruning_pipeline.rs index b8384480c0b6d..3e6bb403e4635 100644 --- a/src/query/service/tests/it/storages/fuse/pruning_pipeline.rs +++ b/src/query/service/tests/it/storages/fuse/pruning_pipeline.rs @@ -84,6 +84,7 @@ async fn apply_snapshot_pruning( op, schema, push_down, + None, bloom_index_cols, vec![], HashSet::new(), diff --git a/src/query/sql/src/planner/binder/bind_query/bind.rs b/src/query/sql/src/planner/binder/bind_query/bind.rs index 8d8eb781eab6a..3584a96f5ee5f 100644 --- a/src/query/sql/src/planner/binder/bind_query/bind.rs +++ b/src/query/sql/src/planner/binder/bind_query/bind.rs @@ -331,7 +331,7 @@ impl Binder { uri_location: None, cluster_by: None, table_options: Default::default(), - iceberg_table_partition: None, + partition_by: None, table_properties: Default::default(), as_query: Some(as_query), table_type: TableType::Temporary, diff --git a/src/query/sql/src/planner/binder/ddl/table.rs b/src/query/sql/src/planner/binder/ddl/table.rs index 9e9f3e8f7a45a..0db2b4c3f4add 100644 --- a/src/query/sql/src/planner/binder/ddl/table.rs +++ b/src/query/sql/src/planner/binder/ddl/table.rs @@ -34,7 +34,10 @@ use databend_common_ast::ast::DescribeTableStmt; use databend_common_ast::ast::DropTableStmt; use databend_common_ast::ast::Engine; use databend_common_ast::ast::ExistsTableStmt; +use databend_common_ast::ast::Expr as AstExpr; +use databend_common_ast::ast::FunctionCall; use databend_common_ast::ast::Identifier; +use databend_common_ast::ast::Literal; use databend_common_ast::ast::ModifyColumnAction; use databend_common_ast::ast::OptimizeTableAction as AstOptimizeTableAction; use databend_common_ast::ast::OptimizeTableStmt; @@ -98,14 +101,19 @@ use databend_common_users::UserApiProvider; use databend_storages_common_table_meta::meta::VectorDistanceType; use databend_storages_common_table_meta::table::OPT_KEY_DATABASE_ID; use databend_storages_common_table_meta::table::OPT_KEY_ENGINE_META; +use databend_storages_common_table_meta::table::OPT_KEY_PARTITION_BY; use databend_storages_common_table_meta::table::OPT_KEY_STORAGE_FORMAT; use databend_storages_common_table_meta::table::OPT_KEY_STORAGE_PREFIX; use databend_storages_common_table_meta::table::OPT_KEY_TABLE_ATTACHED_DATA_URI; use databend_storages_common_table_meta::table::OPT_KEY_TABLE_COMPRESSION; use databend_storages_common_table_meta::table::OPT_KEY_TEMP_PREFIX; +use databend_storages_common_table_meta::table::OPT_KEY_WRITE_DISTRIBUTION_MODE; use databend_storages_common_table_meta::table::TableCompression; +use databend_storages_common_table_meta::table::WriteDistributionMode; use databend_storages_common_table_meta::table::is_reserved_opt_key; +use derive_visitor::Drive; use derive_visitor::DriveMut; +use derive_visitor::Visitor; use log::debug; use opendal::Operator; use parking_lot::RwLock; @@ -175,6 +183,29 @@ use crate::plans::VacuumTemporaryFilesPlan; const FUSE_OPT_KEY_AGGRESSIVE_RECLUSTER: &str = "aggressive_recluster"; +#[derive(Visitor)] +#[visitor(FunctionCall(enter))] +struct PartitionBucketValidator { + invalid: bool, +} + +impl PartitionBucketValidator { + fn enter_function_call(&mut self, func: &FunctionCall) { + if func.name.name.eq_ignore_ascii_case("bucket") { + self.invalid |= !matches!( + func.args.as_slice(), + [ + AstExpr::Literal { + value: Literal::UInt64(buckets), + .. + }, + _ + ] if (1..=i32::MAX as u64).contains(buckets) + ); + } + } +} + pub(in crate::planner::binder) struct AnalyzeCreateTableResult { pub(in crate::planner::binder) schema: TableSchemaRef, pub(in crate::planner::binder) field_comments: Vec, @@ -567,7 +598,7 @@ impl Binder { table_type, engine, uri_location, - iceberg_table_partition, + partition_by, table_properties, } = stmt; @@ -669,12 +700,36 @@ impl Binder { None => None, }; - let table_partition = iceberg_table_partition.as_ref().map(|partitions| { - partitions - .iter() - .map(|p| p.to_string()) - .collect::>() - }); + if partition_by.is_some() && !matches!(engine, Engine::Fuse | Engine::Iceberg) { + return Err(ErrorCode::UnsupportedEngineParams(format!( + "PARTITION BY is not supported for engine {engine}" + ))); + } + + // Iceberg currently supports identity partition columns. Fuse partition expressions + // are normalized after the table schema has been bound below. + let table_partition = if matches!(engine, Engine::Iceberg) { + partition_by + .as_ref() + .map(|partitions| { + partitions + .iter() + .map(|partition| match partition { + AstExpr::ColumnRef { column, .. } + if column.database.is_none() && column.table.is_none() => + { + Ok(column.column.to_string()) + } + _ => Err(ErrorCode::BadArguments(format!( + "Iceberg PARTITION BY only supports column identifiers, got `{partition:#}`" + ))), + }) + .collect::>>() + }) + .transpose()? + } else { + None + }; let mut storage_params = match (uri_location_to_use.as_ref(), engine) { (Some(uri), Engine::Fuse) => { @@ -928,6 +983,32 @@ impl Binder { ))); } + if matches!(engine, Engine::Fuse) { + if let Some(mode) = options.get(OPT_KEY_WRITE_DISTRIBUTION_MODE) { + let mode = mode.parse::()?; + if mode == WriteDistributionMode::Hash && partition_by.is_none() { + return Err(ErrorCode::TableOptionInvalid(format!( + "{OPT_KEY_WRITE_DISTRIBUTION_MODE}='hash' requires PARTITION BY" + ))); + } + } + } + + if matches!(engine, Engine::Fuse) + && let Some(partition_exprs) = partition_by + { + let partition_opt = ClusterOption { + cluster_exprs: partition_exprs.clone(), + }; + let keys = self + .analyze_table_keys(&partition_opt, schema.clone(), None, "PARTITION BY", false) + .await?; + options.insert( + OPT_KEY_PARTITION_BY.to_owned(), + format!("({})", keys.join(", ")), + ); + } + let mut cluster_key = None; if let Some(cluster_opt) = cluster_by { let keys = self @@ -2316,6 +2397,18 @@ impl Binder { cluster_opt: &ClusterOption, schema: TableSchemaRef, table_indexes: Option<&BTreeMap>, + ) -> Result> { + self.analyze_table_keys(cluster_opt, schema, table_indexes, "Cluster by", true) + .await + } + + async fn analyze_table_keys( + &mut self, + cluster_opt: &ClusterOption, + schema: TableSchemaRef, + table_indexes: Option<&BTreeMap>, + key_name: &str, + allow_vector: bool, ) -> Result> { let ClusterOption { cluster_exprs } = cluster_opt; @@ -2357,28 +2450,36 @@ impl Binder { let mut cluster_keys = Vec::with_capacity(expr_len); let mut vector_cluster_key_num = 0; for cluster_expr in cluster_exprs.iter() { + if !allow_vector { + let mut validator = PartitionBucketValidator { invalid: false }; + cluster_expr.drive(&mut validator); + if validator.invalid { + return Err(ErrorCode::InvalidClusterKeys(format!( + "{key_name} expression `{cluster_expr:#}` must use bucket with a constant count between 1 and {}", + i32::MAX + ))); + } + } + let (cluster_key, _) = scalar_binder.bind(cluster_expr)?; if cluster_key.used_columns().len() != 1 || !cluster_key.evaluable() { return Err(ErrorCode::InvalidClusterKeys(format!( - "Cluster by expression `{:#}` is invalid", - cluster_expr + "{key_name} expression `{cluster_expr:#}` is invalid" ))); } let expr = cluster_key.as_expr()?; if !expr.is_deterministic(&BUILTIN_FUNCTIONS) { return Err(ErrorCode::InvalidClusterKeys(format!( - "Cluster by expression `{:#}` is not deterministic", - cluster_expr + "{key_name} expression `{cluster_expr:#}` is not deterministic" ))); } let data_type = expr.data_type(); let (is_valid_type, is_vector_type) = Self::valid_cluster_key_type(data_type); - if !is_valid_type { + if !is_valid_type || is_vector_type && !allow_vector { return Err(ErrorCode::InvalidClusterKeys(format!( - "Unsupported data type '{}' for cluster by expression `{:#}`", - data_type, cluster_expr + "Unsupported data type '{data_type}' for {key_name} expression `{cluster_expr:#}`" ))); } if is_vector_type { @@ -2396,8 +2497,7 @@ impl Binder { }; let Ok(field) = schema.field_with_name(&id.column_name) else { return Err(ErrorCode::InvalidClusterKeys(format!( - "Cluster by expression `{:#}` is invalid", - cluster_expr + "{key_name} expression `{cluster_expr:#}` is invalid" ))); }; let distances = table_indexes diff --git a/src/query/sql/src/planner/optimizer/ir/stats/constraint.rs b/src/query/sql/src/planner/optimizer/ir/stats/constraint.rs index 7bc1edcffb026..0ec5b683e8f16 100644 --- a/src/query/sql/src/planner/optimizer/ir/stats/constraint.rs +++ b/src/query/sql/src/planner/optimizer/ir/stats/constraint.rs @@ -52,22 +52,10 @@ impl ValueConstraint { match op { ComparisonOp::Equal => ValueConstraint::Eq(datum), ComparisonOp::NotEqual => ValueConstraint::NotEq, - ComparisonOp::GT => ValueConstraint::Range { - lower: Bound::Excluded(datum), - upper: Bound::Unbounded, - }, - ComparisonOp::GTE => ValueConstraint::Range { - lower: Bound::Included(datum), - upper: Bound::Unbounded, - }, - ComparisonOp::LT => ValueConstraint::Range { - lower: Bound::Unbounded, - upper: Bound::Excluded(datum), - }, - ComparisonOp::LTE => ValueConstraint::Range { - lower: Bound::Unbounded, - upper: Bound::Included(datum), - }, + op => { + let (lower, upper) = op.range_bounds(datum).unwrap(); + ValueConstraint::Range { lower, upper } + } } } diff --git a/src/query/sql/src/planner/plans/scalar_expr.rs b/src/query/sql/src/planner/plans/scalar_expr.rs index be25600569470..5b5030d312bef 100644 --- a/src/query/sql/src/planner/plans/scalar_expr.rs +++ b/src/query/sql/src/planner/plans/scalar_expr.rs @@ -19,6 +19,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; +use std::ops::Bound; use std::str::FromStr; use std::sync::Arc; @@ -887,6 +888,18 @@ impl ComparisonOp { ComparisonOp::LTE => ComparisonOp::GTE, } } + + pub fn range_bounds(&self, value: T) -> Option<(Bound, Bound)> + where T: Clone { + match self { + ComparisonOp::Equal => Some((Bound::Included(value.clone()), Bound::Included(value))), + ComparisonOp::NotEqual => None, + ComparisonOp::GT => Some((Bound::Excluded(value), Bound::Unbounded)), + ComparisonOp::LT => Some((Bound::Unbounded, Bound::Excluded(value))), + ComparisonOp::GTE => Some((Bound::Included(value), Bound::Unbounded)), + ComparisonOp::LTE => Some((Bound::Unbounded, Bound::Included(value))), + } + } } impl<'a> TryFrom<&'a BinaryOperator> for ComparisonOp { diff --git a/src/query/storages/common/table_meta/src/table/table_keys.rs b/src/query/storages/common/table_meta/src/table/table_keys.rs index 6df391599ce42..f18e847b1e587 100644 --- a/src/query/storages/common/table_meta/src/table/table_keys.rs +++ b/src/query/storages/common/table_meta/src/table/table_keys.rs @@ -16,6 +16,7 @@ use std::collections::BTreeMap; use std::collections::HashSet; use std::fmt::Display; use std::fmt::Formatter; +use std::str::FromStr; use std::sync::LazyLock; use databend_common_exception::ErrorCode; @@ -73,10 +74,34 @@ pub const OPT_KEY_RANDOM_MAX_STRING_LEN: &str = "max_string_len"; pub const OPT_KEY_RANDOM_MAX_ARRAY_LEN: &str = "max_array_len"; pub const OPT_KEY_CLUSTER_TYPE: &str = "cluster_type"; +/// The normalized Fuse partition expressions. This is set by CREATE TABLE +/// PARTITION BY and is not a user-settable table option. +pub const OPT_KEY_PARTITION_BY: &str = "partition_by"; +pub const OPT_KEY_WRITE_DISTRIBUTION_MODE: &str = "write_distribution_mode"; pub const OPT_KEY_ENABLE_COPY_DEDUP_FULL_PATH: &str = "copy_dedup_full_path"; pub const LINEAR_CLUSTER_TYPE: &str = "linear"; pub const HILBERT_CLUSTER_TYPE: &str = "hilbert"; +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum WriteDistributionMode { + None, + Hash, +} + +impl FromStr for WriteDistributionMode { + type Err = ErrorCode; + + fn from_str(value: &str) -> Result { + match value.to_ascii_lowercase().as_str() { + "none" => Ok(Self::None), + "hash" => Ok(Self::Hash), + _ => Err(ErrorCode::TableOptionInvalid(format!( + "{OPT_KEY_WRITE_DISTRIBUTION_MODE} must be either 'none' or 'hash', got: {value}" + ))), + } + } +} + /// Table option keys that reserved for internal usage only /// - Users are not allowed to specify this option keys in DDL /// - Should not be shown in `show create table` statement @@ -85,6 +110,7 @@ pub static RESERVED_TABLE_OPTION_KEYS: LazyLock> = LazyLoc r.insert(OPT_KEY_DATABASE_ID); r.insert(OPT_KEY_LEGACY_SNAPSHOT_LOC); r.insert(OPT_KEY_RECURSIVE_CTE); + r.insert(OPT_KEY_PARTITION_BY); r.insert(OPT_KEY_CLUSTER_TYPE); r }); @@ -98,6 +124,7 @@ pub static INTERNAL_TABLE_OPTION_KEYS: LazyLock> = LazyLoc r.insert(OPT_KEY_CHANGE_TRACKING_BEGIN_VER); r.insert(OPT_KEY_TEMP_PREFIX); r.insert(OPT_KEY_RECURSIVE_CTE); + r.insert(OPT_KEY_PARTITION_BY); r.insert(OPT_KEY_CLUSTER_TYPE); r }); diff --git a/src/query/storages/fuse/Cargo.toml b/src/query/storages/fuse/Cargo.toml index 0c6d44822b9f0..9484f2d628d50 100644 --- a/src/query/storages/fuse/Cargo.toml +++ b/src/query/storages/fuse/Cargo.toml @@ -8,6 +8,7 @@ edition = { workspace = true } [dependencies] +databend-common-ast = { workspace = true } databend-common-base = { workspace = true } databend-common-catalog = { workspace = true } databend-common-column = { workspace = true } diff --git a/src/query/storages/fuse/src/fuse_table.rs b/src/query/storages/fuse/src/fuse_table.rs index 0eb78700c3376..464a8f2132ecd 100644 --- a/src/query/storages/fuse/src/fuse_table.rs +++ b/src/query/storages/fuse/src/fuse_table.rs @@ -25,6 +25,8 @@ use std::time::Instant; use async_channel::Receiver; use chrono::Duration; use chrono::TimeDelta; +use databend_common_ast::ast::Expr as AstExpr; +use databend_common_ast::parser::parse_cluster_key_exprs; use databend_common_base::runtime::GlobalIORuntime; use databend_common_catalog::catalog::StorageDescription; use databend_common_catalog::plan::DataSourcePlan; @@ -105,12 +107,14 @@ use databend_storages_common_table_meta::table::OPT_KEY_BLOOM_INDEX_TYPE; use databend_storages_common_table_meta::table::OPT_KEY_CHANGE_TRACKING; use databend_storages_common_table_meta::table::OPT_KEY_CLUSTER_TYPE; use databend_storages_common_table_meta::table::OPT_KEY_LEGACY_SNAPSHOT_LOC; +use databend_storages_common_table_meta::table::OPT_KEY_PARTITION_BY; use databend_storages_common_table_meta::table::OPT_KEY_SEGMENT_FORMAT; use databend_storages_common_table_meta::table::OPT_KEY_SNAPSHOT_LOCATION; use databend_storages_common_table_meta::table::OPT_KEY_SNAPSHOT_LOCATION_FIXED_FLAG; use databend_storages_common_table_meta::table::OPT_KEY_STORAGE_FORMAT; use databend_storages_common_table_meta::table::OPT_KEY_TABLE_ATTACHED_DATA_URI; use databend_storages_common_table_meta::table::OPT_KEY_TABLE_COMPRESSION; +use databend_storages_common_table_meta::table::OPT_KEY_WRITE_DISTRIBUTION_MODE; use databend_storages_common_table_meta::table::TableCompression; use databend_storages_common_table_meta::table::analyze_top_n_size_from_options; use futures_util::TryStreamExt; @@ -147,6 +151,7 @@ use crate::io::WriteSettings; use crate::operations::ChangesDesc; use crate::operations::SnapshotHint; use crate::operations::load_last_snapshot_hint; +use crate::pruning::PartitionPruningInfo; use crate::statistics::STATS_STRING_PREFIX_LEN; use crate::statistics::reduce_block_statistics; @@ -593,8 +598,67 @@ impl FuseTable { self.table_info.meta.cluster_key_id() } + fn partition_key_str(&self) -> Option<&str> { + self.table_info + .meta + .options + .get(OPT_KEY_PARTITION_BY) + .map(String::as_str) + } + + fn resolve_partition_keys(&self) -> Option> { + self.partition_key_str() + .map(|partition_key| parse_cluster_key_exprs(partition_key).unwrap()) + } + + pub fn partition_key_count(&self) -> usize { + self.resolve_partition_keys().map_or(0, |keys| keys.len()) + } + + pub fn use_hash_write_distribution(&self) -> bool { + self.partition_key_count() != 0 + && self + .table_info + .options() + .get(OPT_KEY_WRITE_DISTRIBUTION_MODE) + .is_some_and(|mode| mode.eq_ignore_ascii_case("hash")) + } + + /// The key used for physical Fuse layout. Partition expressions are always the + /// leading dimensions, followed by the user-visible linear cluster key. + pub fn resolve_physical_cluster_keys(&self) -> Option> { + let mut keys = self.resolve_partition_keys().unwrap_or_default(); + if let Some(cluster_keys) = self.resolve_cluster_keys() { + keys.extend(cluster_keys); + } + (!keys.is_empty()).then_some(keys) + } + + pub fn physical_cluster_key_id(&self) -> Option { + self.cluster_key_id() + .or_else(|| self.partition_key_str().map(|_| 0)) + } + + pub fn partition_pruning_info( + &self, + ctx: Arc, + ) -> Option { + let partition_key_count = self.partition_key_count(); + if partition_key_count == 0 { + return None; + } + Some(PartitionPruningInfo { + cluster_key_id: self.physical_cluster_key_id().unwrap(), + partition_keys: self + .linear_cluster_keys(ctx) + .into_iter() + .take(partition_key_count) + .collect(), + }) + } + pub fn linear_cluster_keys(&self, ctx: Arc) -> Vec> { - let Some(cluster_key_exprs) = self.resolve_cluster_keys() else { + let Some(cluster_key_exprs) = self.resolve_physical_cluster_keys() else { return vec![]; }; @@ -633,7 +697,7 @@ impl FuseTable { } pub fn cluster_key_types(&self, ctx: Arc) -> Vec { - let Some(ast_exprs) = self.resolve_cluster_keys() else { + let Some(ast_exprs) = self.resolve_physical_cluster_keys() else { return vec![]; }; let cluster_keys = parse_cluster_keys(ctx, Arc::new(self.clone()), ast_exprs).unwrap(); @@ -937,7 +1001,7 @@ impl FuseTable { pub fn enable_stream_block_write(&self, ctx: Arc) -> Result { Ok(ctx.get_settings().get_enable_block_stream_write()? && matches!(self.storage_format, FuseStorageFormat::Parquet) - && self.cluster_key_meta().is_none()) + && self.resolve_physical_cluster_keys().is_none()) } pub fn with_schema(&self, schema: Arc) -> Arc { diff --git a/src/query/storages/fuse/src/operations/analyze/analyze_state_sink.rs b/src/query/storages/fuse/src/operations/analyze/analyze_state_sink.rs index 819bfd68457b8..61b81ad4793c4 100644 --- a/src/query/storages/fuse/src/operations/analyze/analyze_state_sink.rs +++ b/src/query/storages/fuse/src/operations/analyze/analyze_state_sink.rs @@ -550,7 +550,7 @@ impl SinkAnalyzeState { Option, )> { // 1. Read table snapshot. - let default_cluster_key_id = table.cluster_key_id(); + let default_cluster_key_id = table.physical_cluster_key_id(); // 2. Iterator segments and blocks to estimate statistics. let mut read_segment_count = 0; diff --git a/src/query/storages/fuse/src/operations/append.rs b/src/query/storages/fuse/src/operations/append.rs index 1fe850b291995..c65d75b5f556d 100644 --- a/src/query/storages/fuse/src/operations/append.rs +++ b/src/query/storages/fuse/src/operations/append.rs @@ -44,6 +44,7 @@ use crate::FuseTable; use crate::io::StreamBlockProperties; use crate::operations::TransformBlockBuilder; use crate::operations::TransformBlockWriter; +use crate::operations::TransformPartitionBy; use crate::operations::TransformSerializeBlock; use crate::operations::TransformVectorCluster; use crate::statistics::ClusterStatsGenerator; @@ -57,7 +58,10 @@ impl FuseTable { pipeline: &mut Pipeline, table_meta_timestamps: TableMetaTimestamps, ) -> Result<()> { - let enable_stream_block_write = self.enable_stream_block_write(ctx.clone())?; + // Stream block writing does not expose a partition-boundary split point. + // Use the regular append pipeline for partitioned tables. + let enable_stream_block_write = + self.partition_key_count() == 0 && self.enable_stream_block_write(ctx.clone())?; if enable_stream_block_write { let properties = StreamBlockProperties::try_create( ctx.clone(), @@ -168,6 +172,24 @@ impl FuseTable { } pipeline.add_pipe(builder.finalize()); } + let partition_key_indices: Arc<[_]> = + cluster_stats_gen.cluster_key_index[..cluster_stats_gen.partition_key_count].into(); + if !partition_key_indices.is_empty() { + let mut builder = pipeline.add_transform_with_specified_len( + move |input, output| { + Ok(ProcessorPtr::create(AccumulatingTransformer::create( + input, + output, + TransformPartitionBy::new(partition_key_indices.clone()), + ))) + }, + transform_len, + )?; + if need_match { + builder.add_items_prepend(vec![create_dummy_item()]); + } + pipeline.add_pipe(builder.finalize()); + } Ok(cluster_stats_gen) } @@ -177,6 +199,27 @@ impl FuseTable { pipeline: &mut Pipeline, block_thresholds: BlockThresholds, modified_schema: Option>, + ) -> Result { + self.cluster_gen_for_append_impl(ctx, pipeline, block_thresholds, modified_schema, false) + } + + pub fn cluster_gen_for_update( + &self, + ctx: Arc, + pipeline: &mut Pipeline, + block_thresholds: BlockThresholds, + modified_schema: Option>, + ) -> Result { + self.cluster_gen_for_append_impl(ctx, pipeline, block_thresholds, modified_schema, true) + } + + fn cluster_gen_for_append_impl( + &self, + ctx: Arc, + pipeline: &mut Pipeline, + block_thresholds: BlockThresholds, + modified_schema: Option>, + rewrite_replaced_block: bool, ) -> Result { let input_schema = modified_schema.unwrap_or(DataSchema::from(self.schema_with_stream()).into()); @@ -212,6 +255,19 @@ impl FuseTable { move || TransformSortPartial::new(LimitType::None, sort_desc.clone()) }); } + let partition_key_indices: Arc<[_]> = + cluster_stats_gen.cluster_key_index[..cluster_stats_gen.partition_key_count].into(); + if !partition_key_indices.is_empty() { + if rewrite_replaced_block { + pipeline.add_accumulating_transformer(move || { + TransformPartitionBy::new_for_update(partition_key_indices.clone()) + }); + } else { + pipeline.add_accumulating_transformer(move || { + TransformPartitionBy::new(partition_key_indices.clone()) + }); + } + } Ok(cluster_stats_gen) } @@ -222,7 +278,7 @@ impl FuseTable { block_thresholds: BlockThresholds, input_schema: Arc, ) -> Result { - if self.cluster_key_meta().is_none() { + if self.physical_cluster_key_id().is_none() { return Ok(ClusterStatsGenerator::default()); } @@ -320,8 +376,8 @@ impl FuseTable { } } - Ok(ClusterStatsGenerator::new( - self.cluster_key_id().unwrap(), + let mut generator = ClusterStatsGenerator::new( + self.physical_cluster_key_id().unwrap(), cluster_key_index, extra_key_num, level, @@ -330,7 +386,9 @@ impl FuseTable { vector_operator, merged, ctx.get_function_context()?, - )) + ); + generator.partition_key_count = self.partition_key_count(); + Ok(generator) } pub fn get_option(&self, opt_key: &str, default: T) -> T { diff --git a/src/query/storages/fuse/src/operations/changes.rs b/src/query/storages/fuse/src/operations/changes.rs index 04043b2f08b94..78687945a559f 100644 --- a/src/query/storages/fuse/src/operations/changes.rs +++ b/src/query/storages/fuse/src/operations/changes.rs @@ -310,6 +310,7 @@ impl FuseTable { self.get_operator(), table_schema.clone(), &push_downs, + None, bloom_index_cols, ngram_args, spatial_index_columns, diff --git a/src/query/storages/fuse/src/operations/common/meta/mutation_log.rs b/src/query/storages/fuse/src/operations/common/meta/mutation_log.rs index a5d900ae431ed..cd8659adee524 100644 --- a/src/query/storages/fuse/src/operations/common/meta/mutation_log.rs +++ b/src/query/storages/fuse/src/operations/common/meta/mutation_log.rs @@ -51,6 +51,7 @@ pub enum MutationLogEntry { }, AppendBlock { block_meta: Arc, + merge_hll: bool, }, DeletedBlock { index: BlockMetaIndex, diff --git a/src/query/storages/fuse/src/operations/common/processors/mod.rs b/src/query/storages/fuse/src/operations/common/processors/mod.rs index 3b0f398c29a6f..e9f41f911236e 100644 --- a/src/query/storages/fuse/src/operations/common/processors/mod.rs +++ b/src/query/storages/fuse/src/operations/common/processors/mod.rs @@ -18,6 +18,7 @@ mod transform_block_writer; mod transform_constraint_verify; mod transform_merge_commit_meta; mod transform_mutation_aggregator; +mod transform_partition_by; mod transform_serialize_block; mod transform_serialize_segment; mod transform_vector_cluster; @@ -29,6 +30,7 @@ pub use transform_block_writer::TransformBlockWriter; pub use transform_constraint_verify::TransformConstraintVerify; pub use transform_merge_commit_meta::TransformMergeCommitMeta; pub use transform_mutation_aggregator::TableMutationAggregator; +pub use transform_partition_by::TransformPartitionBy; pub use transform_serialize_block::TransformSerializeBlock; pub use transform_serialize_segment::TransformSerializeSegment; pub use transform_serialize_segment::new_serialize_segment_processor; diff --git a/src/query/storages/fuse/src/operations/common/processors/multi_table_insert_commit.rs b/src/query/storages/fuse/src/operations/common/processors/multi_table_insert_commit.rs index a6b15a653b502..a97aa6e9fe394 100644 --- a/src/query/storages/fuse/src/operations/common/processors/multi_table_insert_commit.rs +++ b/src/query/storages/fuse/src/operations/common/processors/multi_table_insert_commit.rs @@ -306,7 +306,7 @@ impl AsyncSink for CommitMultiTableInsert { *m = TransformMergeCommitMeta::merge_commit_meta( m.clone(), meta, - table.cluster_key_id(), + table.physical_cluster_key_id(), )?; } None => { diff --git a/src/query/storages/fuse/src/operations/common/processors/transform_block_writer.rs b/src/query/storages/fuse/src/operations/common/processors/transform_block_writer.rs index 07fcc16db4ad1..56348372a6628 100644 --- a/src/query/storages/fuse/src/operations/common/processors/transform_block_writer.rs +++ b/src/query/storages/fuse/src/operations/common/processors/transform_block_writer.rs @@ -294,6 +294,7 @@ impl AsyncAccumulatingTransform for TransformBlockWriter { DataBlock::empty_with_meta(Box::new(MutationLogs { entries: vec![MutationLogEntry::AppendBlock { block_meta: Arc::new(extended_block_meta), + merge_hll: false, }], ..Default::default() })) diff --git a/src/query/storages/fuse/src/operations/common/processors/transform_mutation_aggregator.rs b/src/query/storages/fuse/src/operations/common/processors/transform_mutation_aggregator.rs index 2a96a261847e0..3cf8f498a9e2f 100644 --- a/src/query/storages/fuse/src/operations/common/processors/transform_mutation_aggregator.rs +++ b/src/query/storages/fuse/src/operations/common/processors/transform_mutation_aggregator.rs @@ -76,6 +76,7 @@ use crate::statistics::get_min_max_stats; use crate::statistics::prepare_cluster_key_exprs; use crate::statistics::reducers::merge_statistics_mut; use crate::statistics::reducers::reduce_block_metas; +use crate::statistics::same_partition; use crate::statistics::sort_by_cluster_stats; pub struct TableMutationAggregator { @@ -196,12 +197,12 @@ impl TableMutationAggregator { kind: MutationKind, table_meta_timestamps: TableMetaTimestamps, ) -> Self { - let fill_missing_cluster_stats = table.cluster_key_meta().is_some(); + let fill_missing_cluster_stats = table.resolve_physical_cluster_keys().is_some(); let virtual_schema = table.table_info.meta.virtual_schema.clone(); let cluster_key_exprs = if fill_missing_cluster_stats { table - .resolve_cluster_keys() + .resolve_physical_cluster_keys() .map(|cluster_keys| { parse_cluster_keys(ctx.clone(), Arc::new(table.clone()), cluster_keys) }) @@ -215,7 +216,8 @@ impl TableMutationAggregator { dal: table.get_operator(), location_gen: table.meta_location_generator().clone(), thresholds: table.get_block_thresholds(), - default_cluster_key: table.cluster_key_id(), + default_cluster_key: table.physical_cluster_key_id(), + partition_key_count: table.partition_key_count(), cluster_key_exprs: Arc::from(cluster_key_exprs.clone()), schema: table.schema(), kind, @@ -274,12 +276,17 @@ impl TableMutationAggregator { } } } - MutationLogEntry::AppendBlock { block_meta } => { + MutationLogEntry::AppendBlock { + block_meta, + merge_hll, + } => { // MERGE and REPLACE append logical INSERT/UPDATE after-images. - if matches!( - self.write_segment_ctx.kind, - MutationKind::MergeInto | MutationKind::Replace - ) { + if merge_hll + || matches!( + self.write_segment_ctx.kind, + MutationKind::MergeInto | MutationKind::Replace + ) + { BlockHLLState::merge_column_hll(&mut self.hll, &block_meta.column_hlls); } self.accumulate_top_n(block_meta.column_top_n.clone())?; @@ -372,35 +379,57 @@ impl TableMutationAggregator { }); } + let mut partition_groups: Vec> = Vec::new(); + for block in merged_blocks { + if partition_groups + .last() + .and_then(|group| group.last()) + .is_none_or(|previous| { + !same_partition( + previous.0.cluster_stats.as_ref(), + block.0.cluster_stats.as_ref(), + self.write_segment_ctx.default_cluster_key, + self.write_segment_ctx.partition_key_count, + ) + }) + { + partition_groups.push(Vec::new()); + } + partition_groups.last_mut().unwrap().push(block); + } + let mut tasks = Vec::new(); - let segments_num = - (merged_blocks.len() / self.write_segment_ctx.thresholds.block_per_segment).max(1); - let chunk_size = merged_blocks.len().div_ceil(segments_num); - for chunk in &merged_blocks.into_iter().chunks(chunk_size) { - let (new_blocks, new_hlls): (Vec>, Vec>) = - chunk.unzip(); - let new_hlls = if new_hlls.iter().all(|v| v.is_none()) { - None - } else { - let hlls = new_hlls - .into_iter() - .map(|x| x.unwrap_or_default()) - .collect::>(); - Some(SegmentStatistics::new(hlls, Vec::new()).to_bytes()?) - }; - // Only compaction/reclustering output may be force-marked perfect to keep - // those operations at a fixed point. REPLACE/MERGE append after-images - // must retain the physical perfect-block count from reduce_block_metas. - let force_all_blocks_perfect = matches!( - self.write_segment_ctx.kind, - MutationKind::Compact | MutationKind::Recluster - ) && new_blocks.len() > 1; - - let ctx = self.write_segment_ctx.clone(); - tasks.push(async move { - ctx.write_segment(new_blocks, new_hlls, force_all_blocks_perfect) - .await - }); + for partition_blocks in partition_groups { + let segments_num = (partition_blocks.len() + / self.write_segment_ctx.thresholds.block_per_segment) + .max(1); + let chunk_size = partition_blocks.len().div_ceil(segments_num); + for chunk in &partition_blocks.into_iter().chunks(chunk_size) { + let (new_blocks, new_hlls): (Vec>, Vec>) = + chunk.unzip(); + let new_hlls = if new_hlls.iter().all(|v| v.is_none()) { + None + } else { + let hlls = new_hlls + .into_iter() + .map(|x| x.unwrap_or_default()) + .collect::>(); + Some(SegmentStatistics::new(hlls, Vec::new()).to_bytes()?) + }; + // Only compaction/reclustering output may be force-marked perfect to keep + // those operations at a fixed point. REPLACE/MERGE append after-images + // must retain the physical perfect-block count from reduce_block_metas. + let force_all_blocks_perfect = matches!( + self.write_segment_ctx.kind, + MutationKind::Compact | MutationKind::Recluster + ) && new_blocks.len() > 1; + + let ctx = self.write_segment_ctx.clone(); + tasks.push(async move { + ctx.write_segment(new_blocks, new_hlls, force_all_blocks_perfect) + .await + }); + } } let threads_nums = self.ctx.get_settings().get_max_threads()? as usize; @@ -837,6 +866,7 @@ struct WriteSegmentCtx { location_gen: TableMetaLocationGenerator, thresholds: BlockThresholds, default_cluster_key: Option, + partition_key_count: usize, cluster_key_exprs: Arc<[Expr]>, schema: TableSchemaRef, kind: MutationKind, diff --git a/src/query/storages/fuse/src/operations/common/processors/transform_partition_by.rs b/src/query/storages/fuse/src/operations/common/processors/transform_partition_by.rs new file mode 100644 index 0000000000000..29ef8a9a60f03 --- /dev/null +++ b/src/query/storages/fuse/src/operations/common/processors/transform_partition_by.rs @@ -0,0 +1,300 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::BlockEntry; +use databend_common_expression::BlockMetaInfoDowncast; +use databend_common_expression::Column; +use databend_common_expression::DataBlock; +use databend_common_expression::types::AnyType; +use databend_common_expression::types::DecimalColumn; +use databend_common_expression::types::NullableColumn; +use databend_common_expression::types::NumberColumn; +use databend_common_pipeline_transforms::AccumulatingTransform; + +use crate::operations::mutation::SerializeDataMeta; + +/// Splits a block that is already sorted by the physical Fuse key into blocks +/// whose partition-key prefix is constant. +/// +/// This is a per-block boundary invariant; fragments of the same partition from +/// different input blocks are not merged here. +pub struct TransformPartitionBy { + partition_key_indices: Arc<[usize]>, + rewrite_replaced_block: bool, +} + +impl TransformPartitionBy { + pub fn new(partition_key_indices: Arc<[usize]>) -> Self { + Self { + partition_key_indices, + rewrite_replaced_block: false, + } + } + + pub fn new_for_update(partition_key_indices: Arc<[usize]>) -> Self { + Self { + partition_key_indices, + // A direct UPDATE normally replaces one old block with one new block. Partition + // splitting makes that one-to-many, so rewrite it as one delete plus appends. + rewrite_replaced_block: true, + } + } +} + +fn partition_boundaries(block: &DataBlock, indices: &[usize]) -> Vec { + let rows = block.num_rows(); + if rows <= 1 || indices.is_empty() { + return Vec::new(); + } + + let mut changed = vec![false; rows]; + for index in indices { + mark_column_boundaries(block.get_by_offset(*index), &mut changed); + } + + changed + .into_iter() + .enumerate() + .skip(1) + .filter_map(|(row, changed)| changed.then_some(row)) + .collect() +} + +fn mark_column_boundaries(entry: &BlockEntry, changed: &mut [bool]) { + if let BlockEntry::Column(column) = entry { + mark_column_boundaries_inner(column, changed); + } +} + +fn mark_column_boundaries_inner(column: &Column, changed: &mut [bool]) { + match column { + Column::Null { .. } | Column::EmptyArray { .. } | Column::EmptyMap { .. } => {} + Column::Number(column) => mark_number_boundaries(column, changed), + Column::Decimal(column) => mark_decimal_boundaries(column, changed), + Column::Boolean(column) => mark_iter_boundaries(column.iter(), changed), + Column::String(column) => mark_iter_boundaries(column.iter(), changed), + Column::Timestamp(column) => mark_iter_boundaries(column.iter(), changed), + Column::TimestampTz(column) => mark_iter_boundaries(column.iter(), changed), + Column::Date(column) => mark_iter_boundaries(column.iter(), changed), + Column::Nullable(column) => mark_nullable_boundaries(column, changed), + _ => mark_iter_boundaries(column.iter(), changed), + } +} + +fn mark_number_boundaries(column: &NumberColumn, changed: &mut [bool]) { + match column { + NumberColumn::UInt8(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::UInt16(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::UInt32(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::UInt64(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::Int8(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::Int16(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::Int32(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::Int64(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::Float32(column) => mark_iter_boundaries(column.iter(), changed), + NumberColumn::Float64(column) => mark_iter_boundaries(column.iter(), changed), + } +} + +fn mark_decimal_boundaries(column: &DecimalColumn, changed: &mut [bool]) { + match column { + DecimalColumn::Decimal64(column, _) => mark_iter_boundaries(column.iter(), changed), + DecimalColumn::Decimal128(column, _) => mark_iter_boundaries(column.iter(), changed), + DecimalColumn::Decimal256(column, _) => mark_iter_boundaries(column.iter(), changed), + } +} + +fn mark_nullable_boundaries(column: &NullableColumn, changed: &mut [bool]) { + let mut value_changed = vec![false; changed.len()]; + mark_column_boundaries_inner(&column.column, &mut value_changed); + + let mut validity = column.validity.iter(); + let Some(mut previous_valid) = validity.next() else { + return; + }; + for (offset, valid) in validity.enumerate() { + let row = offset + 1; + changed[row] |= valid != previous_valid || valid && value_changed[row]; + previous_valid = valid; + } +} + +fn mark_iter_boundaries(values: impl Iterator, changed: &mut [bool]) { + let mut values = values; + let Some(mut previous) = values.next() else { + return; + }; + for (offset, value) in values.enumerate() { + let row = offset + 1; + changed[row] |= value != previous; + previous = value; + } +} + +impl AccumulatingTransform for TransformPartitionBy { + const NAME: &'static str = "TransformPartitionBy"; + + fn transform(&mut self, mut block: DataBlock) -> Result> { + if self.partition_key_indices.is_empty() + || block.num_rows() == 0 + || block.num_rows() == 1 && !self.rewrite_replaced_block + { + return Ok(vec![block]); + } + + let replaced_block_meta = if self.rewrite_replaced_block { + block.take_meta() + } else { + None + }; + let boundaries = partition_boundaries(&block, &self.partition_key_indices); + let mut blocks = Vec::with_capacity(boundaries.len() + 1); + let mut start = 0; + for end in boundaries { + blocks.push(block.slice(start..end)); + start = end; + } + blocks.push(block.slice(start..block.num_rows())); + + if let Some(meta) = replaced_block_meta { + let meta = SerializeDataMeta::downcast_from(meta) + .ok_or_else(|| ErrorCode::Internal("Invalid partitioned UPDATE block metadata"))?; + let SerializeDataMeta::SerializeBlock(serialize_block) = meta else { + return Err(ErrorCode::Internal( + "Invalid partitioned UPDATE block metadata", + )); + }; + for block in &mut blocks { + block.replace_meta(Box::new(SerializeDataMeta::SerializeAppend)); + } + blocks.insert( + 0, + DataBlock::empty_with_meta(Box::new(SerializeDataMeta::SerializeBlock( + serialize_block, + ))), + ); + } + Ok(blocks) + } +} + +#[cfg(test)] +mod tests { + use databend_common_expression::DataBlock; + use databend_common_expression::FromData; + use databend_common_expression::types::Int32Type; + use databend_common_expression::types::StringType; + + use super::*; + use crate::operations::common::BlockMetaIndex; + use crate::operations::mutation::ClusterStatsGenType; + use crate::operations::mutation::SerializeBlock; + + #[test] + fn test_partition_boundaries() { + let block = DataBlock::new_from_columns(vec![ + Int32Type::from_data(vec![0, 0, 1, 1, 1]), + StringType::from_data(vec!["a", "a", "a", "b", "b"]), + ]); + assert_eq!(partition_boundaries(&block, &[0, 1]), vec![2, 3]); + + let block = DataBlock::new_from_columns(vec![Int32Type::from_opt_data(vec![ + None, + None, + Some(1), + Some(1), + Some(2), + ])]); + assert_eq!(partition_boundaries(&block, &[0]), vec![2, 4]); + } + + #[test] + fn test_split_sorted_block_by_partition_prefix() -> Result<()> { + let block = DataBlock::new_from_columns(vec![ + Int32Type::from_data(vec![0, 0, 1, 1, 2]), + Int32Type::from_data(vec![1, 2, 1, 2, 1]), + ]); + let mut transform = TransformPartitionBy::new(Arc::from([0])); + let blocks = transform.transform(block)?; + + assert_eq!(blocks.len(), 3); + assert_eq!( + blocks.iter().map(DataBlock::num_rows).collect::>(), + vec![2, 2, 1] + ); + Ok(()) + } + + #[test] + fn test_split_replaced_block_as_delete_and_append() -> Result<()> { + let meta = SerializeDataMeta::SerializeBlock(SerializeBlock::create( + BlockMetaIndex::default(), + ClusterStatsGenType::Generally, + 2, + 0, + )); + let block = DataBlock::new_from_columns(vec![ + Int32Type::from_data(vec![0, 0, 1, 1, 2]), + Int32Type::from_data(vec![1, 2, 1, 2, 1]), + ]) + .add_meta(Some(Box::new(meta)))?; + let mut transform = TransformPartitionBy::new_for_update(Arc::from([0])); + let mut blocks = transform.transform(block)?; + + assert_eq!(blocks.len(), 4); + assert!(blocks[0].is_empty()); + assert!(matches!( + SerializeDataMeta::downcast_from(blocks[0].take_meta().unwrap()), + Some(SerializeDataMeta::SerializeBlock(SerializeBlock { + logical_updated_rows: 2, + logical_deleted_rows: 0, + .. + })) + )); + assert!(blocks[1..].iter_mut().all(|block| matches!( + SerializeDataMeta::downcast_from(block.take_meta().unwrap()), + Some(SerializeDataMeta::SerializeAppend) + ))); + assert_eq!( + blocks[1..] + .iter() + .map(DataBlock::num_rows) + .collect::>(), + vec![2, 2, 1] + ); + + let meta = SerializeDataMeta::SerializeBlock(SerializeBlock::create( + BlockMetaIndex::default(), + ClusterStatsGenType::Generally, + 1, + 0, + )); + let block = DataBlock::new_from_columns(vec![Int32Type::from_data(vec![1])]) + .add_meta(Some(Box::new(meta)))?; + let mut transform = TransformPartitionBy::new_for_update(Arc::from([0])); + let mut blocks = transform.transform(block)?; + assert_eq!(blocks.len(), 2); + assert!(blocks[0].is_empty()); + assert_eq!(blocks[1].num_rows(), 1); + assert!(matches!( + SerializeDataMeta::downcast_from(blocks[1].take_meta().unwrap()), + Some(SerializeDataMeta::SerializeAppend) + )); + Ok(()) + } +} diff --git a/src/query/storages/fuse/src/operations/common/processors/transform_serialize_block.rs b/src/query/storages/fuse/src/operations/common/processors/transform_serialize_block.rs index 4484514682fe0..869c13339a849 100644 --- a/src/query/storages/fuse/src/operations/common/processors/transform_serialize_block.rs +++ b/src/query/storages/fuse/src/operations/common/processors/transform_serialize_block.rs @@ -77,6 +77,7 @@ pub struct TransformSerializeBlock { dal: Operator, table_id: Option, // Only used in multi table insert kind: MutationKind, + pending_merge_hll: bool, pending_logical_change: (u64, u64), } @@ -222,6 +223,7 @@ impl TransformSerializeBlock { dal: table.get_operator(), table_id: if with_tid { Some(table.get_id()) } else { None }, kind, + pending_merge_hll: false, pending_logical_change: (0, 0), }) } @@ -340,6 +342,15 @@ impl Processor for TransformSerializeBlock { Ok(Event::Sync) } } + SerializeDataMeta::SerializeAppend => { + self.pending_merge_hll = true; + self.state = State::NeedSerialize { + block: input_data, + stats_type: ClusterStatsGenType::Generally, + index: None, + }; + Ok(Event::Sync) + } SerializeDataMeta::CompactExtras(compact_extras) => { // compact extras let data_block = Self::mutation_logs( @@ -400,6 +411,7 @@ impl Processor for TransformSerializeBlock { async fn async_process(&mut self) -> Result<()> { match std::mem::replace(&mut self.state, State::Consume) { State::Serialized { serialized, index } => { + let merge_hll = std::mem::take(&mut self.pending_merge_hll); let (logical_updated_rows, logical_deleted_rows) = std::mem::take(&mut self.pending_logical_change); let extended_block_meta = BlockWriter::write_down(&self.dal, serialized).await?; @@ -454,6 +466,7 @@ impl Processor for TransformSerializeBlock { Self::mutation_logs( MutationLogEntry::AppendBlock { block_meta: Arc::new(extended_block_meta), + merge_hll, }, logical_updated_rows, logical_deleted_rows, diff --git a/src/query/storages/fuse/src/operations/common/processors/transform_serialize_segment.rs b/src/query/storages/fuse/src/operations/common/processors/transform_serialize_segment.rs index 62fce13bd91ca..d924ca9adc488 100644 --- a/src/query/storages/fuse/src/operations/common/processors/transform_serialize_segment.rs +++ b/src/query/storages/fuse/src/operations/common/processors/transform_serialize_segment.rs @@ -23,6 +23,7 @@ use databend_common_exception::Result; use databend_common_expression::BlockMetaInfoDowncast; use databend_common_expression::BlockThresholds; use databend_common_expression::DataBlock; +use databend_common_expression::Scalar; use databend_common_pipeline::core::Event; use databend_common_pipeline::core::InputPort; use databend_common_pipeline::core::OutputPort; @@ -52,6 +53,7 @@ use crate::operations::common::MutationLogs; use crate::statistics::ColumnHLLAccumulator; use crate::statistics::RowOrientedSegmentBuilder; use crate::statistics::VirtualColumnAccumulator; +use crate::statistics::partition_values; enum State { None, @@ -87,9 +89,12 @@ pub struct TransformSerializeSegment { input: Arc, output: Arc, output_data: Option, + pending_block: Option, + current_partition: Option>, thresholds: BlockThresholds, default_cluster_key_id: Option, + partition_key_count: usize, table_meta_timestamps: TableMetaTimestamps, is_column_oriented: bool, } @@ -107,7 +112,7 @@ impl TransformSerializeSegment { let virtual_column_accumulator = VirtualColumnAccumulator::try_create(&table_meta.schema, &table_meta.virtual_schema); - let default_cluster_key_id = table.cluster_key_id(); + let default_cluster_key_id = table.physical_cluster_key_id(); let block_top_n_template = table @@ -123,6 +128,8 @@ impl TransformSerializeSegment { input, output, output_data: None, + pending_block: None, + current_partition: None, data_accessor: table.get_operator(), meta_locations: table.meta_location_generator().clone(), state: State::None, @@ -134,6 +141,7 @@ impl TransformSerializeSegment { top_n: HashMap::new(), thresholds, default_cluster_key_id, + partition_key_count: table.partition_key_count(), table_meta_timestamps, is_column_oriented: table.is_column_oriented(), }) @@ -208,7 +216,7 @@ impl Processor for TransformSerializeSegment { return Ok(Event::NeedConsume); } - if self.input.is_finished() { + if self.pending_block.is_none() && self.input.is_finished() { if self.segment_builder.block_count() != 0 { self.state = State::GenerateSegment; return Ok(Event::Sync); @@ -239,17 +247,38 @@ impl Processor for TransformSerializeSegment { return Ok(Event::Finished); } - if self.input.has_data() { - let input_meta = self - .input - .pull_data() - .unwrap()? - .get_meta() - .cloned() - .ok_or_else(|| ErrorCode::Internal("No block meta. It's a bug"))?; - let extended_block_meta = ExtendedBlockMeta::downcast_ref_from(&input_meta) - .ok_or_else(|| ErrorCode::Internal("No block meta. It's a bug"))? - .clone(); + if self.pending_block.is_some() || self.input.has_data() { + let extended_block_meta = if let Some(block) = self.pending_block.take() { + block + } else { + let input_meta = self + .input + .pull_data() + .unwrap()? + .get_meta() + .cloned() + .ok_or_else(|| ErrorCode::Internal("No block meta. It's a bug"))?; + ExtendedBlockMeta::downcast_ref_from(&input_meta) + .ok_or_else(|| ErrorCode::Internal("No block meta. It's a bug"))? + .clone() + }; + + let next_partition = partition_values( + extended_block_meta.block_meta.cluster_stats.as_ref(), + self.default_cluster_key_id, + self.partition_key_count, + ); + if self.segment_builder.block_count() != 0 + && self.partition_key_count != 0 + && (next_partition.is_none() || self.current_partition.as_deref() != next_partition) + { + self.pending_block = Some(extended_block_meta); + self.state = State::GenerateSegment; + return Ok(Event::Sync); + } + if self.current_partition.is_none() { + self.current_partition = next_partition.map(<[Scalar]>::to_vec); + } if let Some(draft_virtual_block_meta) = extended_block_meta.draft_virtual_block_meta { let mut block_meta = extended_block_meta.block_meta.clone(); @@ -285,6 +314,14 @@ impl Processor for TransformSerializeSegment { self.state = State::GenerateSegment; return Ok(Event::Sync); } + + // The last input block may have been held in `pending_block` while the previous + // partition was serialized. If the input finished in the meantime, no port event + // remains to wake this processor, so flush the final partition immediately. + if self.input.is_finished() { + self.state = State::GenerateSegment; + return Ok(Event::Sync); + } } self.input.set_need_data(); @@ -325,6 +362,7 @@ impl Processor for TransformSerializeSegment { self.default_cluster_key_id, additional_stats_meta, )?; + self.current_partition = None; self.state = State::SerializedSegment { data: segment_info.serialize()?, diff --git a/src/query/storages/fuse/src/operations/compact.rs b/src/query/storages/fuse/src/operations/compact.rs index 90141672d4f51..67f4d3d668873 100644 --- a/src/query/storages/fuse/src/operations/compact.rs +++ b/src/query/storages/fuse/src/operations/compact.rs @@ -75,9 +75,10 @@ impl FuseTable { compact_options, self.meta_location_generator().clone(), self.operator.clone(), - self.cluster_key_id(), + self.physical_cluster_key_id(), table_meta_timestamps, )?; + segment_compactor.partition_key_count = self.partition_key_count(); if !segment_compactor.target_select().await? { return Ok(()); @@ -165,8 +166,9 @@ impl FuseTable { thresholds, compact_options, self.operator.clone(), - self.cluster_key_id(), + self.physical_cluster_key_id(), ); + mutator.partition_key_count = self.partition_key_count(); let partitions = mutator.target_select().await?; if partitions.is_empty() { diff --git a/src/query/storages/fuse/src/operations/mutation/meta/mutation_meta.rs b/src/query/storages/fuse/src/operations/mutation/meta/mutation_meta.rs index 98fd63f4fd80d..fb2b7b11ab140 100644 --- a/src/query/storages/fuse/src/operations/mutation/meta/mutation_meta.rs +++ b/src/query/storages/fuse/src/operations/mutation/meta/mutation_meta.rs @@ -28,6 +28,7 @@ use crate::operations::mutation::DeletedSegmentInfo; #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq)] pub enum SerializeDataMeta { SerializeBlock(SerializeBlock), + SerializeAppend, DeletedSegment(DeletedSegmentInfo), CompactExtras(CompactExtraInfo), } diff --git a/src/query/storages/fuse/src/operations/mutation/mutator/block_compact_mutator.rs b/src/query/storages/fuse/src/operations/mutation/mutator/block_compact_mutator.rs index 4962fdf27f0bb..21ff66c483de4 100644 --- a/src/query/storages/fuse/src/operations/mutation/mutator/block_compact_mutator.rs +++ b/src/query/storages/fuse/src/operations/mutation/mutator/block_compact_mutator.rs @@ -25,6 +25,7 @@ use databend_common_catalog::plan::PartitionsShuffleKind; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_expression::BlockThresholds; +use databend_common_expression::Scalar; use databend_common_metrics::storage::*; use databend_storages_common_table_meta::meta::BlockMeta; use databend_storages_common_table_meta::meta::CompactSegmentInfo; @@ -46,7 +47,9 @@ use crate::operations::mutation::CompactExtraInfo; use crate::operations::mutation::CompactLazyPartInfo; use crate::operations::mutation::CompactTaskInfo; use crate::operations::mutation::SegmentIndex; +use crate::statistics::partition_values; use crate::statistics::reducers::merge_statistics_mut; +use crate::statistics::same_partition; use crate::statistics::sort_by_cluster_stats; #[derive(Clone)] @@ -57,6 +60,7 @@ pub struct BlockCompactMutator { pub thresholds: BlockThresholds, pub compact_params: CompactOptions, pub cluster_key_id: Option, + pub partition_key_count: usize, } impl BlockCompactMutator { @@ -73,6 +77,7 @@ impl BlockCompactMutator { thresholds, compact_params, cluster_key_id, + partition_key_count: 0, } } @@ -111,6 +116,8 @@ impl BlockCompactMutator { Arc::new(self.compact_params.base_snapshot.schema.clone()), ); let mut checker = SegmentCompactChecker::new(self.thresholds); + checker.cluster_key_id = self.cluster_key_id; + checker.partition_key_count = self.partition_key_count; let mut segment_idx = 0; let mut is_end = false; @@ -213,6 +220,7 @@ impl BlockCompactMutator { self.ctx.clone(), self.operator.clone(), self.cluster_key_id, + self.partition_key_count, self.thresholds, lazy_parts, ) @@ -237,6 +245,7 @@ impl BlockCompactMutator { ctx: Arc, dal: Operator, cluster_key_id: Option, + partition_key_count: usize, thresholds: BlockThresholds, lazy_parts: Vec, ) -> Result> { @@ -268,8 +277,12 @@ impl BlockCompactMutator { works.push(async move { let mut res = vec![]; for lazy_part in batch { - let mut builder = - CompactTaskBuilder::new(dal.clone(), cluster_key_id, thresholds); + let mut builder = CompactTaskBuilder::new( + dal.clone(), + cluster_key_id, + partition_key_count, + thresholds, + ); let parts = builder .build_tasks( lazy_part.segment_indices, @@ -320,6 +333,8 @@ pub enum CompactLimitState { pub struct SegmentCompactChecker { thresholds: BlockThresholds, + cluster_key_id: Option, + partition_key_count: usize, segments: Vec<(SegmentIndex, Arc)>, total_block_count: u64, @@ -333,6 +348,8 @@ impl SegmentCompactChecker { segments: vec![], total_block_count: 0, thresholds, + cluster_key_id: None, + partition_key_count: 0, compacted_segment_cnt: 0, compacted_imperfect_block_cnt: 0, } @@ -377,14 +394,27 @@ impl SegmentCompactChecker { ) -> Vec)>> { let block_per_segment = self.thresholds.block_per_segment as u64; + let mut output = Vec::new(); + if let Some((_, previous)) = self.segments.last() + && !same_partition( + previous.summary.cluster_stats.as_ref(), + segment.summary.cluster_stats.as_ref(), + self.cluster_key_id, + self.partition_key_count, + ) + { + output.push(std::mem::take(&mut self.segments)); + self.total_block_count = 0; + } + self.total_block_count += segment.summary.block_count; self.segments.push((idx, segment)); if self.total_block_count < block_per_segment { - return vec![]; + return output; } - let output = if self.total_block_count >= 2 * block_per_segment { + let mut threshold_output = if self.total_block_count >= 2 * block_per_segment { let trivial = vec![self.segments.pop().unwrap()]; if self.segments.is_empty() { vec![trivial] @@ -396,6 +426,7 @@ impl SegmentCompactChecker { }; self.total_block_count = 0; + output.append(&mut threshold_output); output } @@ -449,6 +480,7 @@ impl SegmentCompactChecker { struct CompactTaskBuilder { dal: Operator, cluster_key_id: Option, + partition_key_count: usize, thresholds: BlockThresholds, blocks: Vec, @@ -464,10 +496,16 @@ enum TailMergeSource { } impl CompactTaskBuilder { - fn new(dal: Operator, cluster_key_id: Option, thresholds: BlockThresholds) -> Self { + fn new( + dal: Operator, + cluster_key_id: Option, + partition_key_count: usize, + thresholds: BlockThresholds, + ) -> Self { Self { dal, cluster_key_id, + partition_key_count, thresholds, blocks: vec![], total_rows: 0, @@ -637,7 +675,26 @@ impl CompactTaskBuilder { } let mut tasks = VecDeque::new(); + let mut previous_partition: Option> = None; + let mut seen_block = false; for (block_meta, hlls) in blocks.iter() { + let current_partition = partition_values( + block_meta.cluster_stats.as_ref(), + self.cluster_key_id, + self.partition_key_count, + ); + if seen_block + && self.partition_key_count != 0 + && (current_partition.is_none() + || previous_partition.as_deref() != current_partition) + { + self.flush_pending(&mut tasks, &mut unchanged_blocks, &mut block_idx); + // The tail-merging optimization is local to one partition. + tail_merge_source = TailMergeSource::None; + } + previous_partition = current_partition.map(<[Scalar]>::to_vec); + seen_block = true; + if self.is_empty() { if !self.can_start_compact(block_meta) { // Reclustered blocks can be carried by an existing compact group, diff --git a/src/query/storages/fuse/src/operations/mutation/mutator/recluster_mutator.rs b/src/query/storages/fuse/src/operations/mutation/mutator/recluster_mutator.rs index debb348602878..b0bcc8d1a535d 100644 --- a/src/query/storages/fuse/src/operations/mutation/mutator/recluster_mutator.rs +++ b/src/query/storages/fuse/src/operations/mutation/mutator/recluster_mutator.rs @@ -71,6 +71,7 @@ use crate::statistics::PreparedClusterKeyExpr; use crate::statistics::RangeMaxTree; use crate::statistics::VectorClusterInfo; use crate::statistics::get_min_max_stats; +use crate::statistics::partition_values; use crate::statistics::prepare_cluster_key_exprs; use crate::statistics::reducers::merge_statistics_mut; use crate::statistics::sort_by_cluster_stats; @@ -310,6 +311,7 @@ pub struct ReclusterMutator { pub(crate) depth_threshold: f64, pub(crate) block_thresholds: BlockThresholds, pub(crate) cluster_key_id: u32, + pub(crate) partition_key_count: usize, pub(crate) schema: TableSchemaRef, pub(crate) max_tasks: usize, pub(crate) memory_threshold: usize, @@ -326,7 +328,7 @@ impl ReclusterMutator { mode: ReclusterMode, ) -> Result { let schema = table.schema_with_stream(); - let cluster_key_id = table.cluster_key_id().unwrap(); + let cluster_key_id = table.physical_cluster_key_id().unwrap(); let block_thresholds = table.get_block_thresholds(); let depth_threshold = table @@ -351,7 +353,7 @@ impl ReclusterMutator { } // safe to unwrap - let cluster_keys = table.resolve_cluster_keys().unwrap(); + let cluster_keys = table.resolve_physical_cluster_keys().unwrap(); let full_cluster_key_exprs = parse_cluster_keys(ctx.clone(), Arc::new(table.clone()), cluster_keys)?; let vector_cluster_info = vector_cluster_info_from_exprs(table, &full_cluster_key_exprs)?; @@ -376,6 +378,7 @@ impl ReclusterMutator { depth_threshold, block_thresholds, cluster_key_id, + partition_key_count: table.partition_key_count(), max_tasks, memory_threshold, prepared_cluster_key_exprs, @@ -422,6 +425,7 @@ impl ReclusterMutator { depth_threshold, block_thresholds, cluster_key_id, + partition_key_count: 0, max_tasks, memory_threshold, prepared_cluster_key_exprs, @@ -545,7 +549,7 @@ impl ReclusterMutator { blocks: &[&ReclusterBlock], task_budget: usize, ) -> Result> { - let mut blocks_map: BTreeMap> = BTreeMap::new(); + let mut blocks_map: BTreeMap<(ReclusterGroup, Vec), Vec> = BTreeMap::new(); for (idx, block) in blocks.iter().enumerate() { let level = block.stats().level; if level < 0 { @@ -555,8 +559,19 @@ impl ReclusterMutator { // Terminal-level blocks are excluded from further rewrite tasks. continue; } + let partition = if self.partition_key_count == 0 { + Vec::new() + } else if let Some(partition) = partition_values( + Some(block.stats()), + Some(self.cluster_key_id), + self.partition_key_count, + ) { + partition.to_vec() + } else { + continue; + }; blocks_map - .entry(ReclusterGroup::assign(level, self.mode)) + .entry((ReclusterGroup::assign(level, self.mode), partition)) .or_default() .push(idx); } @@ -564,7 +579,7 @@ impl ReclusterMutator { let mut tasks: Vec = Vec::new(); let mut deferred_group_candidates = None; let large_task_bytes_threshold = self.large_task_bytes_threshold(); - for (group, indices) in blocks_map { + for ((group, _), indices) in blocks_map { if tasks.len() >= task_budget { break; } @@ -1170,8 +1185,33 @@ impl ReclusterMutator { let mut current_window_max_depth = 0usize; let (keys, values): (Vec<_>, Vec<_>) = segment_points.into_iter().unzip(); let sorted_indices = compare_scalars(&keys, &self.cluster_key_types)?; + let mut previous_partition: Option> = None; for idx in sorted_indices { + let point = &keys[idx as usize]; + let point_partition = (self.partition_key_count != 0 + && point.len() >= self.partition_key_count) + .then(|| point[..self.partition_key_count].to_vec()); + if previous_partition.is_some() + && (point_partition.is_none() || previous_partition != point_partition) + { + // A sweep window is never allowed to span physical partitions. In + // particular, this prevents the repack-only path from combining + // otherwise disjoint partition segments. + if let Some((segs, depth)) = prev_window.take() { + windows.push((segs, depth)); + } + if !current_window.is_empty() { + windows.push(( + std::mem::take(&mut current_window), + current_window_max_depth, + )); + } + current_window_max_depth = 0; + unfinished_intervals.clear(); + } + previous_partition = point_partition; + let start = &values[idx as usize].0; let end = &values[idx as usize].1; let point_depth = Self::calc_point_depth(unfinished_intervals.len(), start, end); diff --git a/src/query/storages/fuse/src/operations/mutation/mutator/segment_compact_mutator.rs b/src/query/storages/fuse/src/operations/mutation/mutator/segment_compact_mutator.rs index ef5c5be28a316..ce2fc38754cab 100644 --- a/src/query/storages/fuse/src/operations/mutation/mutator/segment_compact_mutator.rs +++ b/src/query/storages/fuse/src/operations/mutation/mutator/segment_compact_mutator.rs @@ -36,6 +36,7 @@ use crate::io::TableMetaLocationGenerator; use crate::io::read::read_segment_stats_in_parallel; use crate::operations::CompactOptions; use crate::statistics::reducers::merge_statistics_mut; +use crate::statistics::same_partition; use crate::statistics::sort_by_cluster_stats; #[derive(Default)] @@ -61,6 +62,7 @@ pub struct SegmentCompactMutator { location_generator: TableMetaLocationGenerator, compaction: SegmentCompactionState, default_cluster_key_id: Option, + pub(crate) partition_key_count: usize, table_meta_timestamps: TableMetaTimestamps, } @@ -80,6 +82,7 @@ impl SegmentCompactMutator { location_generator, compaction: Default::default(), default_cluster_key_id, + partition_key_count: 0, table_meta_timestamps, }) } @@ -131,7 +134,7 @@ impl SegmentCompactMutator { let fuse_segment_io = SegmentsIO::create(self.ctx.clone(), self.data_accessor.clone(), schema); let chunk_size = self.ctx.get_settings().get_max_threads()? as usize * 4; - let compactor = SegmentCompactor::new( + let mut compactor = SegmentCompactor::new( self.compact_params.block_per_seg as u64, self.default_cluster_key_id, chunk_size, @@ -140,6 +143,7 @@ impl SegmentCompactMutator { &self.location_generator, self.table_meta_timestamps, ); + compactor.partition_key_count = self.partition_key_count; self.compaction = compactor .compact(base_segment_locations, limit, |status| { @@ -168,6 +172,7 @@ pub struct SegmentCompactor<'a> { // within R, smaller one is preferred threshold: u64, default_cluster_key_id: Option, + partition_key_count: usize, // fragmented segment collected so far, it will be reset to empty if compaction occurs fragmented_segments: Vec<(usize, SegmentInfo, Location)>, // state which keep the number of blocks of all the fragmented segment collected so far, @@ -195,6 +200,7 @@ impl<'a> SegmentCompactor<'a> { Self { threshold, default_cluster_key_id, + partition_key_count: 0, accumulated_num_blocks: 0, fragmented_segments: vec![], chunk_size, @@ -322,6 +328,17 @@ impl<'a> SegmentCompactor<'a> { return Ok(()); } + if let Some((_, previous, _)) = self.fragmented_segments.last() + && !same_partition( + previous.summary.cluster_stats.as_ref(), + segment_info.summary.cluster_stats.as_ref(), + self.default_cluster_key_id, + self.partition_key_count, + ) + { + self.compact_fragments().await?; + } + let s = self.accumulated_num_blocks + num_blocks_current_segment; if s < self.threshold { diff --git a/src/query/storages/fuse/src/operations/mutation_source.rs b/src/query/storages/fuse/src/operations/mutation_source.rs index 378e98345dcd3..2be438e301461 100644 --- a/src/query/storages/fuse/src/operations/mutation_source.rs +++ b/src/query/storages/fuse/src/operations/mutation_source.rs @@ -191,6 +191,7 @@ impl FuseTable { self.operator.clone(), self.schema_with_stream(), &push_down, + self.partition_pruning_info(ctx.clone()), self.bloom_index_cols(), Self::create_ngram_index_args(&self.table_info.meta.indexes, &self.schema(), false)?, spatial_index_columns, diff --git a/src/query/storages/fuse/src/operations/read_partitions.rs b/src/query/storages/fuse/src/operations/read_partitions.rs index 69887236129af..e29fbfc6ee72b 100644 --- a/src/query/storages/fuse/src/operations/read_partitions.rs +++ b/src/query/storages/fuse/src/operations/read_partitions.rs @@ -913,6 +913,7 @@ impl FuseTable { dal, table_schema.clone(), &push_downs, + self.partition_pruning_info(ctx.clone()), self.bloom_index_cols(), ngram_args, spatial_index_columns, diff --git a/src/query/storages/fuse/src/operations/recluster.rs b/src/query/storages/fuse/src/operations/recluster.rs index ed2f71b4348f3..73ad8473edb17 100644 --- a/src/query/storages/fuse/src/operations/recluster.rs +++ b/src/query/storages/fuse/src/operations/recluster.rs @@ -65,7 +65,7 @@ impl FuseTable { ctx.set_status_info("[FUSE-RECLUSTER] Starting recluster operation"); - if self.cluster_key_meta().is_none() { + if self.physical_cluster_key_id().is_none() { return Ok(None); } @@ -452,6 +452,7 @@ impl FuseTable { dal, schema.clone(), push_down, + None, BloomIndexColumns::None, vec![], HashSet::new(), diff --git a/src/query/storages/fuse/src/pruning/fuse_pruner.rs b/src/query/storages/fuse/src/pruning/fuse_pruner.rs index cc8c91611aff4..a488d03fcd3a7 100644 --- a/src/query/storages/fuse/src/pruning/fuse_pruner.rs +++ b/src/query/storages/fuse/src/pruning/fuse_pruner.rs @@ -67,6 +67,8 @@ use crate::pruning::BloomPruner; use crate::pruning::BloomPrunerCreator; use crate::pruning::FusePruningStatistics; use crate::pruning::InvertedIndexPruner; +use crate::pruning::PartitionPruner; +use crate::pruning::PartitionPruningInfo; use crate::pruning::PruningCostController; use crate::pruning::PruningCostKind; use crate::pruning::SegmentLocation; @@ -86,6 +88,7 @@ pub struct PruningContext { pub limit_pruner: Arc, pub range_pruner: Arc, pub bloom_pruner: Option>, + pub partition_pruner: Option, pub internal_column_pruner: Option>, pub inverted_index_pruner: Option>, pub virtual_column_pruner: Option>, @@ -102,6 +105,7 @@ impl PruningContext { dal: Operator, table_schema: TableSchemaRef, push_down: &Option, + partition_pruning_info: Option, bloom_index_cols: BloomIndexColumns, ngram_args: Vec, spatial_index_columns: HashSet, @@ -117,6 +121,11 @@ impl PruningContext { .effective_filters(&BUILTIN_FUNCTIONS) .map(|f| f.filter.as_expr(&BUILTIN_FUNCTIONS)) }); + let partition_pruner = PartitionPruner::try_create( + func_ctx.clone(), + filter_expr.as_ref(), + partition_pruning_info, + ); // Limit pruner. // if there are ordering/filter clause, ignore limit, even it has been pushed down @@ -229,6 +238,7 @@ impl PruningContext { limit_pruner, range_pruner, bloom_pruner, + partition_pruner, internal_column_pruner, inverted_index_pruner, virtual_column_pruner, @@ -257,6 +267,7 @@ impl FusePruner { dal: Operator, table_schema: TableSchemaRef, push_down: &Option, + partition_pruning_info: Option, bloom_index_cols: BloomIndexColumns, ngram_args: Vec, spatial_index_columns: HashSet, @@ -267,6 +278,7 @@ impl FusePruner { dal, table_schema, push_down, + partition_pruning_info, bloom_index_cols, ngram_args, spatial_index_columns, @@ -280,6 +292,7 @@ impl FusePruner { dal: Operator, table_schema: TableSchemaRef, push_down: &Option, + partition_pruning_info: Option, bloom_index_cols: BloomIndexColumns, ngram_args: Vec, spatial_index_columns: HashSet, @@ -308,6 +321,7 @@ impl FusePruner { dal, table_schema.clone(), push_down, + partition_pruning_info, bloom_index_cols, ngram_args, spatial_index_columns, diff --git a/src/query/storages/fuse/src/pruning/mod.rs b/src/query/storages/fuse/src/pruning/mod.rs index 89864254d5d43..0ada116e90c23 100644 --- a/src/query/storages/fuse/src/pruning/mod.rs +++ b/src/query/storages/fuse/src/pruning/mod.rs @@ -18,6 +18,7 @@ mod expr_bloom_filter; mod expr_runtime_pruner; mod fuse_pruner; mod inverted_index_pruner; +mod partition_pruner; mod profile_guard; mod pruner_location; mod pruning_statistics; @@ -39,6 +40,8 @@ pub use fuse_pruner::PruningContext; pub use fuse_pruner::table_sample; pub use inverted_index_pruner::InvertedIndexPruner; pub use inverted_index_pruner::create_inverted_index_query; +pub use partition_pruner::PartitionPruner; +pub use partition_pruner::PartitionPruningInfo; pub use pruner_location::SegmentLocation; pub use pruner_location::create_segment_location_vector; pub use pruning_statistics::FusePruningStatistics; diff --git a/src/query/storages/fuse/src/pruning/partition_pruner.rs b/src/query/storages/fuse/src/pruning/partition_pruner.rs new file mode 100644 index 0000000000000..dd5e1effd9c75 --- /dev/null +++ b/src/query/storages/fuse/src/pruning/partition_pruner.rs @@ -0,0 +1,546 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use databend_common_expression::Constant; +use databend_common_expression::ConstantFolder; +use databend_common_expression::Domain; +use databend_common_expression::Expr; +use databend_common_expression::ExprVisitor; +use databend_common_expression::FunctionContext; +use databend_common_expression::RangeConstraint; +use databend_common_expression::RemoteExpr; +use databend_common_expression::Scalar; +use databend_common_expression::expr::Cast; +use databend_common_expression::expr::ColumnRef; +use databend_common_expression::expr::FunctionCall; +use databend_common_expression::type_check::check_function; +use databend_common_expression::visit_expr; +use databend_common_functions::BUILTIN_FUNCTIONS; +use databend_common_sql::plans::ComparisonOp; +use databend_common_statistics::HistogramBounds; +use databend_common_statistics::HistogramRangeBounds; +use databend_storages_common_table_meta::meta::ClusterStatistics; + +use crate::statistics::partition_values; + +#[derive(Clone)] +pub struct PartitionPruningInfo { + pub cluster_key_id: u32, + pub partition_keys: Vec>, +} + +pub struct PartitionPruner { + cluster_key_id: u32, + partition_keys: Vec>, + filter: Expr, + func_ctx: FunctionContext, +} + +impl PartitionPruner { + pub fn try_create( + func_ctx: FunctionContext, + filter: Option<&Expr>, + info: Option, + ) -> Option { + let filter = filter?.clone(); + let info = info?; + if info.partition_keys.is_empty() { + return None; + } + + let partition_keys = info + .partition_keys + .into_iter() + .map(|expr| { + ConstantFolder::fold( + &expr.as_expr(&BUILTIN_FUNCTIONS), + &func_ctx, + &BUILTIN_FUNCTIONS, + ) + .0 + }) + .collect(); + + Some(Self { + cluster_key_id: info.cluster_key_id, + partition_keys, + filter, + func_ctx, + }) + } + + pub fn should_keep(&self, stats: Option<&ClusterStatistics>) -> bool { + let Some(values) = + partition_values(stats, Some(self.cluster_key_id), self.partition_keys.len()) + else { + return true; + }; + + let replacements = self.partition_keys.iter().zip(values).collect::>(); + let mut visitor = PartitionPredicateRewriter { + replacements, + func_ctx: &self.func_ctx, + }; + if !visitor.conjunctive_predicate_possible(&self.filter) { + return false; + } + let filter = visit_expr(&self.filter, &mut visitor) + .unwrap() + .unwrap_or_else(|| self.filter.clone()); + let (filter, _) = ConstantFolder::fold(&filter, &self.func_ctx, &BUILTIN_FUNCTIONS); + + !matches!( + filter, + Expr::Constant(Constant { + scalar: Scalar::Boolean(false), + .. + }) + ) + } +} + +struct PartitionPredicateRewriter<'a> { + replacements: Vec<(&'a Expr, &'a Scalar)>, + func_ctx: &'a FunctionContext, +} + +impl PartitionPredicateRewriter<'_> { + fn replace(&self, expr: Expr) -> Option> { + self.replacements + .iter() + .find(|(partition_expr, _)| *partition_expr == &expr) + .map(|(partition_expr, value)| { + Expr::Constant(Constant { + span: None, + scalar: (*value).clone(), + data_type: partition_expr.data_type().clone(), + }) + }) + } + + fn replace_comparison(&self, call: &FunctionCall) -> Option> { + let constraint = RangeConstraint::try_from_function_call(call)?; + let predicate_domain = predicate_domain(std::slice::from_ref(&constraint))?; + (!self.column_domain_possible(&constraint.column_id, &predicate_domain)).then(|| { + Expr::Constant(Constant { + span: call.span, + scalar: Scalar::Boolean(false), + data_type: call.return_type.clone(), + }) + }) + } + + fn conjunctive_predicate_possible(&self, filter: &Expr) -> bool { + let mut ranges = HashMap::new(); + collect_conjunctive_ranges(filter, &mut ranges); + ranges.values().all(|constraints| { + let Some(domain) = predicate_domain(constraints) else { + return true; + }; + self.column_domain_possible(&constraints[0].column_id, &domain) + }) + } + + fn column_domain_possible(&self, column_id: &str, domain: &Domain) -> bool { + self.replacements + .iter() + .all(|(partition_expr, partition_value)| { + let mut input_domains = ConstantFolder::full_input_domains(partition_expr); + if input_domains + .insert(column_id.to_string(), domain.clone()) + .is_none() + { + return true; + } + self.partition_value_possible(partition_expr, partition_value, &input_domains) + }) + } + + fn partition_value_possible( + &self, + partition_expr: &Expr, + partition_value: &Scalar, + input_domains: &HashMap, + ) -> bool { + let expected = Expr::Constant(Constant { + span: None, + scalar: partition_value.clone(), + data_type: partition_expr.data_type().clone(), + }); + let Ok(equality) = check_function( + None, + "eq", + &[], + &[partition_expr.clone(), expected], + &BUILTIN_FUNCTIONS, + ) else { + return true; + }; + let (equality, _) = ConstantFolder::fold_with_domain( + &equality, + input_domains, + self.func_ctx, + &BUILTIN_FUNCTIONS, + ); + !matches!( + equality, + Expr::Constant(Constant { + scalar: Scalar::Boolean(false), + .. + }) + ) + } +} + +fn predicate_domain(constraints: &[RangeConstraint]) -> Option { + let data_type = &constraints.first()?.data_type; + if constraints + .iter() + .any(|constraint| constraint.data_type != *data_type || constraint.constant.is_null()) + { + return None; + } + + if let Some(constraint) = constraints + .iter() + .find(|constraint| constraint.operator == "eq") + { + return Some(constraint.constant.as_ref().domain(data_type)); + } + + let (min, max) = Domain::full(&data_type.remove_nullable()).to_minmax(); + let mut bounds = HistogramBounds::new(min.to_datum()?, max.to_datum()?); + for constraint in constraints { + let op = ComparisonOp::try_from_func_name(&constraint.operator)?; + let value = constraint.constant.clone().to_datum()?; + let (lower, upper) = op.range_bounds(value)?; + bounds = match HistogramBounds::from_range_constraint( + bounds.lower_bound(), + bounds.upper_bound(), + &lower, + &upper, + ) + .ok()? + { + HistogramRangeBounds::Bounds(bounds) => bounds, + HistogramRangeBounds::Empty | HistogramRangeBounds::Imprecise => return None, + }; + } + Domain::from_datum( + data_type, + bounds.lower_bound().clone(), + bounds.upper_bound().clone(), + false, + ) + .ok() +} + +fn collect_conjunctive_ranges( + expr: &Expr, + ranges: &mut HashMap>>, +) { + let Expr::FunctionCall(call) = expr else { + return; + }; + match call.function.signature.name.as_str() { + "and" | "and_filters" => { + for arg in &call.args { + collect_conjunctive_ranges(arg, ranges); + } + } + "is_true" if call.args.len() == 1 => collect_conjunctive_ranges(&call.args[0], ranges), + _ => { + let Some(constraint) = RangeConstraint::try_from_function_call(call) else { + return; + }; + if constraint.operator != "noteq" { + ranges + .entry(constraint.column_id.clone()) + .or_default() + .push(constraint); + } + } + } +} + +impl ExprVisitor for PartitionPredicateRewriter<'_> { + fn enter_column_ref( + &mut self, + column: &ColumnRef, + ) -> Result>, Self::Error> { + Ok(self.replace(Expr::ColumnRef(column.clone()))) + } + + fn enter_cast(&mut self, cast: &Cast) -> Result>, Self::Error> { + if let Some(expr) = self.replace(Expr::Cast(cast.clone())) { + Ok(Some(expr)) + } else { + Self::visit_cast(cast, self) + } + } + + fn enter_function_call( + &mut self, + call: &FunctionCall, + ) -> Result>, Self::Error> { + if let Some(expr) = self.replace(Expr::FunctionCall(call.clone())) { + Ok(Some(expr)) + } else if let Some(expr) = self.replace_comparison(call) { + Ok(Some(expr)) + } else { + Self::visit_function_call(call, self) + } + } +} + +#[cfg(test)] +mod tests { + use databend_common_expression::types::DataType; + use databend_common_expression::types::NumberDataType; + + use super::*; + + fn column_ref(id: &str, data_type: DataType) -> Expr { + Expr::ColumnRef(ColumnRef { + span: None, + id: id.to_string(), + data_type, + display_name: "p".to_string(), + }) + } + + fn int64(value: i64) -> Expr { + Expr::Constant(Constant { + span: None, + scalar: Scalar::Number(value.into()), + data_type: DataType::Number(NumberDataType::Int64), + }) + } + + fn timestamp(value: i64) -> Expr { + Expr::Constant(Constant { + span: None, + scalar: Scalar::Timestamp(value), + data_type: DataType::Timestamp, + }) + } + + fn call(name: &str, args: Vec>) -> Expr { + check_function(None, name, &[], &args, &BUILTIN_FUNCTIONS).unwrap() + } + + fn assert_comparison_pruned( + partition_expr: &Expr, + partition_value: &Scalar, + filter: Expr, + ) { + let Expr::FunctionCall(call) = filter else { + panic!("expected function call"); + }; + let func_ctx = FunctionContext::default(); + let rewriter = PartitionPredicateRewriter { + replacements: vec![(partition_expr, partition_value)], + func_ctx: &func_ctx, + }; + assert!(matches!( + rewriter.replace_comparison(&call), + Some(Expr::Constant(Constant { + scalar: Scalar::Boolean(false), + .. + })) + )); + } + + fn assert_timestamp_range_partition_projection( + partition_expr: Expr, + matching_value: Scalar, + non_matching_value: Scalar, + ) { + const MICROS_PER_DAY: i64 = 86_400_000_000; + + let column = column_ref("order_time", DataType::Timestamp); + let filter = call("and_filters", vec![ + call("gte", vec![column.clone(), timestamp(0)]), + call("lt", vec![column, timestamp(MICROS_PER_DAY)]), + ]); + let func_ctx = FunctionContext::default(); + + let rewriter = PartitionPredicateRewriter { + replacements: vec![(&partition_expr, &matching_value)], + func_ctx: &func_ctx, + }; + assert!(rewriter.conjunctive_predicate_possible(&filter)); + + let rewriter = PartitionPredicateRewriter { + replacements: vec![(&partition_expr, &non_matching_value)], + func_ctx: &func_ctx, + }; + assert!(!rewriter.conjunctive_predicate_possible(&filter)); + } + + #[test] + fn test_replace_partition_expr_by_structural_identity() { + let partition_expr = column_ref("p", DataType::Number(NumberDataType::Int64)); + let partition_value = Scalar::Number(1_i64.into()); + let func_ctx = FunctionContext::default(); + let pruner = PartitionPredicateRewriter { + replacements: vec![(&partition_expr, &partition_value)], + func_ctx: &func_ctx, + }; + + let replaced = pruner.replace(partition_expr.clone()); + assert_eq!( + replaced, + Some(Expr::Constant(Constant { + span: None, + scalar: partition_value, + data_type: DataType::Number(NumberDataType::Int64), + })) + ); + } + + #[test] + fn test_does_not_replace_different_expr_with_same_sql_display() { + let partition_expr = column_ref("p", DataType::Number(NumberDataType::Int64)); + let filter_expr = column_ref("p", DataType::Number(NumberDataType::UInt64)); + assert_eq!(partition_expr.sql_display(), filter_expr.sql_display()); + + let partition_value = Scalar::Number(1_i64.into()); + let func_ctx = FunctionContext::default(); + let pruner = PartitionPredicateRewriter { + replacements: vec![(&partition_expr, &partition_value)], + func_ctx: &func_ctx, + }; + + assert_eq!(pruner.replace(filter_expr), None); + } + + #[test] + fn test_replace_source_column_range_comparison() { + let data_type = DataType::Number(NumberDataType::Int64); + let column = column_ref("p", data_type); + + let upper_partition_expr = call("plus", vec![column.clone(), int64(1)]); + let upper_partition_value = Scalar::Number(20_i64.into()); + assert_comparison_pruned( + &upper_partition_expr, + &upper_partition_value, + call("lte", vec![column.clone(), int64(10)]), + ); + assert_comparison_pruned( + &upper_partition_expr, + &upper_partition_value, + call("gte", vec![int64(10), column.clone()]), + ); + + let lower_partition_expr = call("minus", vec![column.clone(), int64(1)]); + let lower_partition_value = Scalar::Number(0_i64.into()); + assert_comparison_pruned( + &lower_partition_expr, + &lower_partition_value, + call("gte", vec![column.clone(), int64(10)]), + ); + assert_comparison_pruned( + &lower_partition_expr, + &lower_partition_value, + call("lte", vec![int64(10), column]), + ); + } + + #[test] + fn test_combines_conjunctive_source_column_range() { + let data_type = DataType::Number(NumberDataType::Int64); + let column = column_ref("p", data_type); + let partition_expr = call("modulo", vec![column.clone(), int64(3)]); + let partition_value = Scalar::Number(2_i64.into()); + let filter = call("and_filters", vec![ + call("gte", vec![column.clone(), int64(4)]), + call("lte", vec![column, int64(4)]), + ]); + let func_ctx = FunctionContext::default(); + let rewriter = PartitionPredicateRewriter { + replacements: vec![(&partition_expr, &partition_value)], + func_ctx: &func_ctx, + }; + + assert!(!rewriter.conjunctive_predicate_possible(&filter)); + + let matching_partition_value = Scalar::Number(1_i64.into()); + let rewriter = PartitionPredicateRewriter { + replacements: vec![(&partition_expr, &matching_partition_value)], + func_ctx: &func_ctx, + }; + assert!(rewriter.conjunctive_predicate_possible(&filter)); + } + + #[test] + fn test_does_not_combine_ranges_across_disjunction() { + let data_type = DataType::Number(NumberDataType::Int64); + let column = column_ref("p", data_type); + let partition_expr = call("modulo", vec![column.clone(), int64(3)]); + let partition_value = Scalar::Number(2_i64.into()); + let filter = call("or", vec![ + call("eq", vec![column.clone(), int64(4)]), + call("eq", vec![column, int64(5)]), + ]); + let func_ctx = FunctionContext::default(); + let rewriter = PartitionPredicateRewriter { + replacements: vec![(&partition_expr, &partition_value)], + func_ctx: &func_ctx, + }; + + assert!(rewriter.conjunctive_predicate_possible(&filter)); + } + + #[test] + fn test_projects_timestamp_range_through_day_partition() { + let column = column_ref("order_time", DataType::Timestamp); + assert_timestamp_range_partition_projection( + call("to_date", vec![column.clone()]), + Scalar::Date(0), + Scalar::Date(1), + ); + assert_timestamp_range_partition_projection( + call("to_start_of_day", vec![column]), + Scalar::Timestamp(0), + Scalar::Timestamp(86_400_000_000), + ); + } + + #[test] + fn test_should_keep_returns_true_for_segment_without_partition_metadata() { + // Segments written before PARTITION BY was added have no cluster statistics. + // The pruner must keep them conservatively rather than silently dropping rows. + let partition_expr = column_ref("p", DataType::Number(NumberDataType::Int64)); + let filter = call("eq", vec![partition_expr.clone(), { + Expr::Constant(Constant { + span: None, + scalar: Scalar::Number(1_i64.into()), + data_type: DataType::Number(NumberDataType::Int64), + }) + }]); + + let pruner = PartitionPruner { + cluster_key_id: 0, + partition_keys: vec![partition_expr], + filter, + func_ctx: FunctionContext::default(), + }; + + // No cluster stats at all — must keep the segment. + assert!(pruner.should_keep(None)); + } +} diff --git a/src/query/storages/fuse/src/pruning/segment_pruner.rs b/src/query/storages/fuse/src/pruning/segment_pruner.rs index 5e924ba0bdac2..9897b2f59f728 100644 --- a/src/query/storages/fuse/src/pruning/segment_pruner.rs +++ b/src/query/storages/fuse/src/pruning/segment_pruner.rs @@ -90,7 +90,11 @@ impl SegmentPruner { info.summary().spatial_stats.as_ref(), ); if pruning_cost.measure(PruningCostKind::SegmentsRange, || { - range_pruner.should_keep(&range_input, None) + self.pruning_ctx + .partition_pruner + .as_ref() + .is_none_or(|pruner| pruner.should_keep(info.summary().cluster_stats.as_ref())) + && range_pruner.should_keep(&range_input, None) }) { // Perf. { diff --git a/src/query/storages/fuse/src/retry/commit.rs b/src/query/storages/fuse/src/retry/commit.rs index e5c22bcd3301d..46104ed46ea6c 100644 --- a/src/query/storages/fuse/src/retry/commit.rs +++ b/src/query/storages/fuse/src/retry/commit.rs @@ -186,7 +186,7 @@ async fn try_rebuild_req( storage_class, table_info.desc.as_str(), )?; - let default_cluster_key_id = latest_table.cluster_key_id(); + let default_cluster_key_id = latest_table.physical_cluster_key_id(); let latest_snapshot = latest_table.read_table_snapshot().await?; let (update_table_meta_req, _) = req .update_table_metas diff --git a/src/query/storages/fuse/src/statistics/cluster_statistics.rs b/src/query/storages/fuse/src/statistics/cluster_statistics.rs index 742c933730d87..74934deff3067 100644 --- a/src/query/storages/fuse/src/statistics/cluster_statistics.rs +++ b/src/query/storages/fuse/src/statistics/cluster_statistics.rs @@ -89,6 +89,7 @@ pub struct ClusterStatsGenerator { block_thresholds: BlockThresholds, pub extra_key_num: usize, + pub partition_key_count: usize, pub cluster_key_index: Vec, pub operators: Vec, pub vector_operator: Option, @@ -112,6 +113,7 @@ impl ClusterStatsGenerator { Self { cluster_key_id, cluster_key_index, + partition_key_count: 0, extra_key_num, level, block_thresholds, @@ -294,6 +296,46 @@ pub fn sort_by_cluster_stats( } } +/// Returns the exact partition prefix stored in physical cluster statistics. +/// A partitioned block/segment is valid only when every prefix dimension has +/// identical min and max values. +pub(crate) fn partition_values( + stats: Option<&ClusterStatistics>, + cluster_key_id: Option, + partition_key_count: usize, +) -> Option<&[Scalar]> { + if partition_key_count == 0 { + return None; + } + let stats = stats?; + if Some(stats.cluster_key_id) != cluster_key_id + || stats.min.len() < partition_key_count + || stats.max.len() < partition_key_count + || stats.min[..partition_key_count] != stats.max[..partition_key_count] + { + return None; + } + Some(&stats.min[..partition_key_count]) +} + +pub(crate) fn same_partition( + left: Option<&ClusterStatistics>, + right: Option<&ClusterStatistics>, + cluster_key_id: Option, + partition_key_count: usize, +) -> bool { + if partition_key_count == 0 { + return true; + } + match ( + partition_values(left, cluster_key_id, partition_key_count), + partition_values(right, cluster_key_id, partition_key_count), + ) { + (Some(left), Some(right)) => left == right, + _ => false, + } +} + pub fn aggregate_cluster_key_min_max( data_block: &DataBlock, key: usize, @@ -603,4 +645,60 @@ mod tests { assert_eq!(min, vec![int32_scalar(i32::MIN)]); assert_eq!(max, vec![int32_scalar(i32::MAX)]); } + + #[test] + fn test_partition_values_returns_none_for_missing_stats() { + // When segment has no cluster statistics, pruner must keep it conservatively. + assert_eq!(partition_values(None, Some(0), 1), None); + } + + #[test] + fn test_partition_values_returns_none_for_mismatched_cluster_key_id() { + let stats = ClusterStatistics::new(0, vec![int32_scalar(1)], vec![int32_scalar(1)], 0); + // Segment was written with a different cluster key id; treat as unpartitioned. + assert_eq!(partition_values(Some(&stats), Some(1), 1), None); + } + + #[test] + fn test_partition_values_returns_none_when_prefix_min_max_differ() { + // A block that spans two partition values is not a valid single-partition block. + let stats = ClusterStatistics::new(0, vec![int32_scalar(1)], vec![int32_scalar(2)], 0); + assert_eq!(partition_values(Some(&stats), Some(0), 1), None); + } + + #[test] + fn test_partition_values_returns_prefix_when_valid() { + let stats = ClusterStatistics::new( + 0, + vec![int32_scalar(3), int32_scalar(10)], + vec![int32_scalar(3), int32_scalar(20)], + 0, + ); + // Only the partition prefix (first dimension) is returned. + let values = partition_values(Some(&stats), Some(0), 1).unwrap(); + assert_eq!(values, &[int32_scalar(3)]); + } + + #[test] + fn test_same_partition_returns_false_when_either_side_has_no_stats() { + // Without partition metadata on either side, compact must not merge the segments. + let stats = ClusterStatistics::new(0, vec![int32_scalar(1)], vec![int32_scalar(1)], 0); + assert!(!same_partition(None, None, Some(0), 1)); + assert!(!same_partition(Some(&stats), None, Some(0), 1)); + assert!(!same_partition(None, Some(&stats), Some(0), 1)); + } + + #[test] + fn test_same_partition_returns_true_when_partition_key_count_is_zero() { + // Tables without PARTITION BY always treat segments as the same partition. + assert!(same_partition(None, None, Some(0), 0)); + } + + #[test] + fn test_same_partition_distinguishes_different_partition_values() { + let left = ClusterStatistics::new(0, vec![int32_scalar(0)], vec![int32_scalar(0)], 0); + let right = ClusterStatistics::new(0, vec![int32_scalar(1)], vec![int32_scalar(1)], 0); + assert!(!same_partition(Some(&left), Some(&right), Some(0), 1)); + assert!(same_partition(Some(&left), Some(&left), Some(0), 1)); + } } diff --git a/src/query/storages/fuse/src/statistics/mod.rs b/src/query/storages/fuse/src/statistics/mod.rs index b3ef78e6fce0e..2710049840ec5 100644 --- a/src/query/storages/fuse/src/statistics/mod.rs +++ b/src/query/storages/fuse/src/statistics/mod.rs @@ -32,7 +32,9 @@ pub use cluster_statistics::VectorClusterOperator; pub use cluster_statistics::aggregate_cluster_key_min_max; pub use cluster_statistics::calculate_block_overlap_depths; pub(crate) use cluster_statistics::get_min_max_stats; +pub(crate) use cluster_statistics::partition_values; pub(crate) use cluster_statistics::prepare_cluster_key_exprs; +pub(crate) use cluster_statistics::same_partition; pub use cluster_statistics::sort_by_cluster_stats; pub use cluster_statistics::vector_cluster_info_from_column; pub use column_statistic::END_OF_UNICODE_RANGE; diff --git a/src/query/storages/fuse/src/table_functions/clustering_information.rs b/src/query/storages/fuse/src/table_functions/clustering_information.rs index 6f7208b8a5df2..31f80ccf8597a 100644 --- a/src/query/storages/fuse/src/table_functions/clustering_information.rs +++ b/src/query/storages/fuse/src/table_functions/clustering_information.rs @@ -174,7 +174,7 @@ impl<'a> ClusteringInformationImpl<'a> { cluster_key: &Option, ) -> Result { let mut default_cluster_key_id = None; - let (cluster_key, exprs) = match (self.table.cluster_key_str(), cluster_key) { + let (cluster_key, mut exprs) = match (self.table.cluster_key_str(), cluster_key) { (a, Some(b)) => { let (cluster_key, exprs) = analyze_cluster_keys(self.ctx.clone(), Arc::new(self.table.clone()), b)?; @@ -183,24 +183,31 @@ impl<'a> ClusteringInformationImpl<'a> { .map(|expr| expr.project_column_ref(|index| Ok(index.as_usize()))) .collect::>>()?; if a.is_some() && a.unwrap() == cluster_key { - default_cluster_key_id = self.table.cluster_key_id(); + default_cluster_key_id = self.table.physical_cluster_key_id(); } (cluster_key, exprs) } (Some(a), None) => { - let cluster_keys = self.table.resolve_cluster_keys().unwrap(); + let cluster_keys = self.table.resolve_physical_cluster_keys().unwrap(); let exprs = parse_cluster_keys( self.ctx.clone(), Arc::new(self.table.clone()), cluster_keys, )?; - default_cluster_key_id = self.table.cluster_key_id(); + default_cluster_key_id = self.table.physical_cluster_key_id(); (a.to_string(), exprs) } _ => { unreachable!("Unclustered table {}", self.table.table_info.desc); } }; + if default_cluster_key_id.is_some() && self.table.partition_key_count() != 0 { + exprs = parse_cluster_keys( + self.ctx.clone(), + Arc::new(self.table.clone()), + self.table.resolve_physical_cluster_keys().unwrap(), + )?; + } let cluster_type = "linear".to_string(); diff --git a/src/query/storages/fuse/src/table_functions/clustering_statistics.rs b/src/query/storages/fuse/src/table_functions/clustering_statistics.rs index 24ff779b2f4f6..2b27b80d22f94 100644 --- a/src/query/storages/fuse/src/table_functions/clustering_statistics.rs +++ b/src/query/storages/fuse/src/table_functions/clustering_statistics.rs @@ -95,7 +95,7 @@ impl TableMetaFunc for ClusteringStatistics { return Ok(DataBlock::empty_with_schema(&Self::schema().into())); } - let cluster_keys = tbl.resolve_cluster_keys().unwrap(); + let cluster_keys = tbl.resolve_physical_cluster_keys().unwrap(); let exprs = parse_cluster_keys(ctx.clone(), Arc::new(tbl.clone()), cluster_keys)?; let scalar_exprs = exprs .into_iter() diff --git a/tests/sqllogictests/suites/base/09_fuse_engine/09_0054_partition_by.test b/tests/sqllogictests/suites/base/09_fuse_engine/09_0054_partition_by.test new file mode 100644 index 0000000000000..a7947e82251b0 --- /dev/null +++ b/tests/sqllogictests/suites/base/09_fuse_engine/09_0054_partition_by.test @@ -0,0 +1,506 @@ +statement ok +DROP DATABASE IF EXISTS db_09_0054 + +statement ok +CREATE DATABASE db_09_0054 + +statement ok +USE db_09_0054 + +# Keep accepting table options before PARTITION BY for both explicit and default Fuse. +statement ok +CREATE TABLE explicit_fuse_option_order(a INT) ENGINE=FUSE ROW_PER_BLOCK=1 PARTITION BY (a) + +statement ok +CREATE TABLE default_fuse_option_order(a INT) ROW_PER_BLOCK=1 PARTITION BY (a) + +# partition_by is internal metadata and cannot be supplied as a raw table option. +statement error table option partition_by reserved +CREATE TABLE raw_partition_option(a INT) PARTITION_BY='(a)' + +statement ok +CREATE TABLE t(p INT, c INT, v STRING) PARTITION BY (p % 3) CLUSTER BY (c) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +SET hide_options_in_show_create_table=0 + +# SHOW CREATE TABLE must emit clauses in the same order accepted by CREATE TABLE. +query TT +SHOW CREATE TABLE t +---- +t PARTITION BY (p % 3)CLUSTER BY (c)BLOCK_PER_SEGMENT='1000'ROW_PER_BLOCK='1000' + +statement ok +SET hide_options_in_show_create_table=1 + +statement ok +INSERT INTO t VALUES (0, 30, 'a0'), (1, 20, 'a1'), (2, 10, 'a2') + +statement ok +INSERT INTO t VALUES (3, 31, 'b0'), (4, 21, 'b1'), (5, 11, 'b2') + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 't') LIMIT 1 +---- +6 6 + +# Segment compaction may combine metadata only inside one physical partition. +statement ok +OPTIMIZE TABLE t COMPACT SEGMENT + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 't') LIMIT 1 +---- +3 6 + +# Block compaction may merge the two small blocks in each partition, but never +# merge blocks from different partition values. +statement ok +OPTIMIZE TABLE t COMPACT + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 't') LIMIT 1 +---- +3 3 + +# The exact partition expression is evaluated against the partition prefix in +# Segment ClusterStatistics before BlockMeta is read. +query T +EXPLAIN SELECT p, c, v FROM t WHERE p % 3 = 1 +---- +partitions total: 3partitions scanned: 1range pruning: 3 to 1 + +# An equality on the source column is projected through the partition expression. +query T +EXPLAIN SELECT p, c, v FROM t WHERE p = 4 +---- +partitions total: 3partitions scanned: 1range pruning: 3 to 1 + +# Conjunctive source-column ranges are combined before being projected through +# the partition expression. +statement ok +CREATE TABLE range_partitioned(p INT, v STRING) PARTITION BY (abs(p)) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO range_partitioned VALUES (-5, 'n5'), (5, 'p5'), (-4, 'n4'), (4, 'p4'), (-3, 'n3'), (3, 'p3'), (-2, 'n2'), (2, 'p2'), (-1, 'n1'), (1, 'p1'), (0, 'z') + +query T +EXPLAIN SELECT p, v FROM range_partitioned WHERE p >= 3 AND p <= 4 +---- +partitions total: 6partitions scanned: 2range pruning: 6 to 2 + +# Range pruning must not remove matching rows. +query II rowsort +SELECT p, length(v) FROM range_partitioned WHERE p >= 3 AND p <= 4 +---- +3 2 +4 2 + +# Strict bounds with constants on the left are normalized before projection. +query T +EXPLAIN SELECT p, v FROM range_partitioned WHERE 4 > p AND 2 < p +---- +partitions total: 6partitions scanned: 1range pruning: 6 to 1 + +query I +SELECT p FROM range_partitioned WHERE 4 > p AND 2 < p +---- +3 + +# Comparisons inside OR are simplified independently, never merged as an AND range. +query T +EXPLAIN SELECT p, v FROM range_partitioned WHERE p = 3 OR p = -5 +---- +partitions total: 6partitions scanned: 2range pruning: 6 to 2 + +query IS rowsort +SELECT p, v FROM range_partitioned WHERE p = 3 OR p = -5 +---- +-5 n5 +3 p3 + +# A timestamp range on the source column is projected through its day expression. +statement ok +CREATE TABLE timestamp_partitioned(order_time TIMESTAMP, label STRING) PARTITION BY (date_trunc(day, order_time)) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO timestamp_partitioned VALUES ('2026-07-07 12:00:00', 'previous'), ('2026-07-08 00:00:00', 'start'), ('2026-07-08 23:59:59', 'end'), ('2026-07-09 00:00:00', 'next') + +query T +EXPLAIN SELECT label FROM timestamp_partitioned WHERE order_time >= TIMESTAMP '2026-07-08' AND order_time < TIMESTAMP '2026-07-09' +---- +partitions total: 3partitions scanned: 1range pruning: 3 to 1 + +query T rowsort +SELECT label FROM timestamp_partitioned WHERE order_time >= TIMESTAMP '2026-07-08' AND order_time < TIMESTAMP '2026-07-09' +---- +end +start + +# Each source-column predicate is projected through its matching partition key. +statement ok +CREATE TABLE multi_partitioned(a INT, b INT) PARTITION BY (a % 2, b % 2) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO multi_partitioned VALUES (0, 0), (2, 2), (0, 1), (2, 3), (1, 0), (3, 2), (1, 1), (3, 3) + +query T +EXPLAIN SELECT a, b FROM multi_partitioned WHERE a = 2 AND b = 3 +---- +partitions total: 4partitions scanned: 1range pruning: 4 to 1 + +query II +SELECT a, b FROM multi_partitioned WHERE a = 2 AND b = 3 +---- +2 3 + +# OR keeps every partition that can satisfy either partition-key dimension. +query T +EXPLAIN SELECT a, b FROM multi_partitioned WHERE a = 2 OR b = 3 +---- +partitions total: 4partitions scanned: 3range pruning: 4 to 3 + +query II rowsort +SELECT a, b FROM multi_partitioned WHERE a = 2 OR b = 3 +---- +2 2 +2 3 +3 3 + +query IIT rowsort +SELECT p, c, v FROM t +---- +0 30 a0 +1 20 a1 +2 10 a2 +3 31 b0 +4 21 b1 +5 11 b2 + +statement ok +UPDATE t SET p = p + 1 WHERE p = 0 + +query IIT rowsort +SELECT p, c, v FROM t +---- +1 20 a1 +1 30 a0 +2 10 a2 +3 31 b0 +4 21 b1 +5 11 b2 + +# Moving a row to another partition must preserve physical partition boundaries. +statement ok +OPTIMIZE TABLE t COMPACT + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 't') LIMIT 1 +---- +3 3 + +statement error +ALTER TABLE t PARTITION BY (c) + +statement error referenced by partition key +ALTER TABLE t DROP COLUMN p + +statement ok +ALTER TABLE t RENAME COLUMN p TO part + +statement ok +INSERT INTO t VALUES (6, 32, 'c0') + +query I +SELECT count(*) FROM t +---- +7 + +statement ok +CREATE TABLE partition_only(p INT, v STRING) PARTITION BY (p % 2) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO partition_only VALUES (0, 'a'), (1, 'b'), (2, 'c'), (3, 'd') + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 'partition_only') LIMIT 1 +---- +2 2 + +query II rowsort +SELECT p, length(v) FROM partition_only +---- +0 1 +1 1 +2 1 +3 1 + +statement ok +CREATE TABLE replace_partitioned(id INT, p INT) PARTITION BY (p % 2) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +REPLACE INTO replace_partitioned ON(id) VALUES (1, 0), (2, 1), (3, 2), (4, 3) + +# REPLACE INTO is sorted and split before serialization, so a mixed input block +# produces one block and segment per physical partition. +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 'replace_partitioned') LIMIT 1 +---- +2 2 + +query II rowsort +SELECT id, p FROM replace_partitioned +---- +1 0 +2 1 +3 2 +4 3 + +statement ok +CREATE TABLE merge_partitioned(id INT, p INT) PARTITION BY (p % 2) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO merge_partitioned VALUES (1, 0), (2, 0), (3, 1) + +statement ok +MERGE INTO merge_partitioned AS t USING (SELECT 1 AS id, 1 AS p) AS s ON t.id = s.id WHEN MATCHED THEN UPDATE SET t.p = s.p + +query II rowsort +SELECT id, p FROM merge_partitioned +---- +1 1 +2 0 +3 1 + +statement ok +OPTIMIZE TABLE merge_partitioned COMPACT + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 'merge_partitioned') LIMIT 1 +---- +2 2 + +# Recluster overlapping cluster-key ranges inside each partition without ever +# selecting or rewriting blocks across the physical partition boundary. +statement ok +CREATE TABLE recluster_partitioned(p INT, c INT) PARTITION BY (p % 2) CLUSTER BY (c) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO recluster_partitioned VALUES (0, 30), (1, 30), (2, 10), (3, 10) + +statement ok +INSERT INTO recluster_partitioned VALUES (4, 31), (5, 31), (6, 11), (7, 11) + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 'recluster_partitioned') LIMIT 1 +---- +4 4 + +statement ok +ALTER TABLE recluster_partitioned RECLUSTER FINAL + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 'recluster_partitioned') LIMIT 1 +---- +2 2 + +query II rowsort +SELECT p, c FROM recluster_partitioned +---- +0 30 +1 30 +2 10 +3 10 +4 31 +5 31 +6 11 +7 11 + +query T +EXPLAIN SELECT p, c FROM recluster_partitioned WHERE p % 2 = 1 +---- +partitions total: 2partitions scanned: 1range pruning: 2 to 1 + +# The last block can be held pending while the preceding partition segment is +# serialized. If the input finishes in the meantime, the pending block must +# immediately flush its own segment instead of waiting for more input forever. +statement ok +SET max_threads=2 + +statement ok +SET max_block_size=10000 + +statement ok +CREATE TABLE segment_finish_race(event_time TIMESTAMP, id BIGINT) PARTITION BY (date_trunc(day, event_time)) ROW_PER_BLOCK=10000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO segment_finish_race SELECT to_timestamp(1767225600000000 + (number DIV 25000) * 86400000000 + number % 25000), number FROM numbers(200000) + +query I +SELECT row_count FROM fuse_snapshot('db_09_0054', 'segment_finish_race') LIMIT 1 +---- +200000 + +# The number of blocks depends on pipeline and cluster scheduling. Assert the +# physical invariant directly: no block may span two day partitions. +query I +SELECT count_if(date_trunc(day, statistics:min::timestamp) != date_trunc(day, statistics:max::timestamp)) +FROM fuse_block_statistics('db_09_0054', 'segment_finish_race') +WHERE column_name = 'event_time' +---- +0 + +# Partition splitting is local to each sorted input block. Fragmented multi-block +# input may therefore produce several blocks for the same physical partition. +statement ok +SET max_threads=4 + +statement ok +SET max_block_size=2 + +statement ok +CREATE TABLE fragmented_input(p INT) PARTITION BY (p % 2) ROW_PER_BLOCK=4 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO fragmented_input SELECT number FROM numbers(12) + +query I +SELECT block_count FROM fuse_snapshot('db_09_0054', 'fragmented_input') LIMIT 1 +---- +6 + +# INSERT SELECT evaluates the physical partition transform, applies GlobalHash, +# and sorts each receiving node before Fuse serialization. +statement ok +SET max_threads=4 + +statement ok +SET max_block_size=1000 + +statement ok +CREATE TABLE hash_partitioned(id BIGINT) PARTITION BY (bucket(16, id)) WRITE_DISTRIBUTION_MODE='hash' ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO hash_partitioned SELECT number::BIGINT FROM numbers(40000) + +query I +SELECT count(*) FROM hash_partitioned +---- +40000 + +# Each bucket contains 2500 rows. GlobalHash plus sort should +# keep the layout close to the intrinsic three blocks per bucket instead of one +# fragment per incoming block. Allow partition-boundary tails. +query I +SELECT to_int64(block_count <= 64) FROM fuse_snapshot('db_09_0054', 'hash_partitioned') LIMIT 1 +---- +1 + +# Other insert sources keep using the existing append path. +statement ok +INSERT INTO hash_partitioned VALUES (40000), (40001) + +query I +SELECT count(*) FROM hash_partitioned WHERE id >= 40000 +---- +2 + +statement error write_distribution_mode must be either 'none' or 'hash' +CREATE TABLE invalid_write_distribution(p INT) PARTITION BY (p) WRITE_DISTRIBUTION_MODE='random' + +statement error write_distribution_mode='hash' requires PARTITION BY +CREATE TABLE missing_write_partition(p INT) WRITE_DISTRIBUTION_MODE='hash' + +statement error must use bucket with a constant count between 1 and 2147483647 +CREATE TABLE invalid_bucket_count(id BIGINT) PARTITION BY (bucket(id, id)) + +statement error must use bucket with a constant count between 1 and 2147483647 +CREATE TABLE zero_bucket_count(id BIGINT) PARTITION BY (bucket(0, id)) + +# Like CLUSTER BY, each physical key expression may reference only one source +# column. Composite partitioning is expressed as PARTITION BY (a, b). +statement error PARTITION BY expression `a \+ b` is invalid +CREATE TABLE multi_column_expr(a INT, b INT) PARTITION BY (a + b) + +# SET OPTIONS cannot overwrite the internal partition_by metadata. +statement error can't change partition_by +ALTER TABLE t SET OPTIONS (partition_by='(part % 2)') + +# MODIFY COLUMN on a partition-key-referenced column is rejected. +statement error referenced by partition key +ALTER TABLE t MODIFY COLUMN part BIGINT + +# MODIFY COLUMN on an unrelated column is allowed. +statement ok +ALTER TABLE t MODIFY COLUMN c BIGINT + +# A filter with no correlation to the partition key must not prune any partition. +statement ok +CREATE TABLE unrelated_filter(p INT, v STRING) PARTITION BY (p % 2) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO unrelated_filter VALUES (0, 'a'), (2, 'm'), (1, 'm'), (3, 'z') + +# v is not a partition key, and both partitions contain the matching value, so +# ordinary range and bloom pruning cannot mask incorrect partition pruning. +query T +EXPLAIN SELECT p, v FROM unrelated_filter WHERE v = 'm' +---- +partitions total: 2partitions scanned: 2 + +query IS rowsort +SELECT p, v FROM unrelated_filter WHERE v = 'm' +---- +1 m +2 m + +# DELETE preserves physical partition boundaries after compaction. +statement ok +CREATE TABLE delete_partitioned(p INT, v STRING) PARTITION BY (p % 2) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 + +statement ok +INSERT INTO delete_partitioned VALUES (0, 'a'), (1, 'b'), (2, 'c'), (3, 'd') + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 'delete_partitioned') LIMIT 1 +---- +2 2 + +statement ok +DELETE FROM delete_partitioned WHERE v = 'b' + +# After deleting one row from partition 1, compact must not merge the two partitions. +statement ok +OPTIMIZE TABLE delete_partitioned COMPACT + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 'delete_partitioned') LIMIT 1 +---- +2 2 + +query II rowsort +SELECT p, length(v) FROM delete_partitioned +---- +0 1 +2 1 +3 1 + +# CREATE TABLE AS SELECT with PARTITION BY applies partition boundaries during the initial write. +statement ok +CREATE TABLE ctas_partitioned PARTITION BY (p % 2) ROW_PER_BLOCK=1000 BLOCK_PER_SEGMENT=1000 +AS SELECT p, v FROM delete_partitioned + +query II +SELECT segment_count, block_count FROM fuse_snapshot('db_09_0054', 'ctas_partitioned') LIMIT 1 +---- +2 2 + +query II rowsort +SELECT p, length(v) FROM ctas_partitioned +---- +0 1 +2 1 +3 1 + +statement ok +DROP DATABASE db_09_0054