feat(fuse): add hash-distributed partition writes

Author: KKould

Cargo.lock

@@ -358,6 +358,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"

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 }

src/query/functions/Cargo.toml

@@ -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::<NumberType<u64>, StringType, NumberType<i32>, _, _>(
+            "bucket",
+            |_, _, _| FunctionDomain::MayThrow,
+            vectorize_with_builder_2_arg::<NumberType<u64>, StringType, NumberType<i32>>(
+                |buckets, value, output, ctx| {
+                    output.push(bucket(buckets, value.as_bytes(), output.len(), ctx));
+                },
+            ),
+        );
+    registry
+        .register_passthrough_nullable_2_arg::<NumberType<u64>, DateType, NumberType<i32>, _, _>(
+            "bucket",
+            |_, _, _| FunctionDomain::MayThrow,
+            vectorize_with_builder_2_arg::<NumberType<u64>, DateType, NumberType<i32>>(
+                |buckets, value, output, ctx| {
+                    output.push(bucket(buckets, &value.to_le_bytes(), output.len(), ctx));
+                },
+            ),
+        );
+    registry.register_passthrough_nullable_2_arg::<
+        NumberType<u64>,
+        TimestampType,
+        NumberType<i32>,
+        _,
+        _,
+    >(
+        "bucket",
+        |_, _, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_2_arg::<NumberType<u64>, TimestampType, NumberType<i32>>(
+            |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<u64>,
+                    NumberType<NUM_TYPE>,
+                    NumberType<i32>,
+                    _,
+                    _,
+                >(
+                    "bucket",
+                    |_, _, _| FunctionDomain::MayThrow,
+                    vectorize_with_builder_2_arg::<
+                        NumberType<u64>,
+                        NumberType<NUM_TYPE>,
+                        NumberType<i32>,
+                    >(|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<T: ArgType>(registry: &mut FunctionRegistry)
 where for<'a> T::ScalarRef<'a>: DFHash {
     registry

src/query/functions/src/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) {

src/query/functions/tests/it/scalars/hash.rs

@@ -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

src/query/functions/tests/it/scalars/testdata/function_list.txt

@@ -851,3 +851,71 @@ evaluation (internal):
 +--------+----------------------------------+
 
 
+ast            : bucket(16, 34)
+raw expr       : bucket(16, 34)
+checked expr   : bucket<UInt64, UInt8>(CAST<UInt8>(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<UInt64, String>(CAST<UInt8>(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<UInt64, Date>(CAST<UInt8>(16_u8 AS UInt64), CAST<Int64>(CAST<UInt16>(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<UInt64, Timestamp>(CAST<UInt8>(16_u8 AS UInt64), CAST<Int64>(CAST<UInt32>(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<UInt64, Int64>(CAST<UInt8>(8_u8 AS UInt64), a)
+optimized expr : bucket<UInt64, Int64>(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])         |
++--------+-----------------------------+
+
+

src/query/functions/tests/it/scalars/testdata/hash.txt

@@ -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<HashSet<&'static str>> = 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);