feat(fuse): add hash-distributed partition writes

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

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::hash_map::DefaultHasher;
 use std::hash::Hash;
 use std::hash::Hasher;
 
@@ -48,6 +47,7 @@ use md5::Digest;
 use md5::Md5 as Md5Hasher;
 use naive_cityhash::cityhash64_with_seed;
 use num_traits::AsPrimitive;
+use siphasher::sip::SipHasher13;
 use twox_hash::XxHash32;
 use twox_hash::XxHash64;
 
@@ -70,7 +70,7 @@ pub fn register(registry: &mut FunctionRegistry) {
             NumberClass::Decimal128 => {
                 struct Siphash64;
                 impl HashFunction for Siphash64 {
-                    type Hasher = DefaultHasher;
+                    type Hasher = SipHasher13;
                     fn name() -> &'static str {
                         "siphash64"
                     }
@@ -102,6 +102,8 @@ pub fn register(registry: &mut FunctionRegistry) {
         });
     }
 
+    register_bucket(registry);
+
     // Could be used in exchange
     registry
         .scalar_builder("siphash64")
@@ -113,9 +115,7 @@ pub fn register(registry: &mut FunctionRegistry) {
             GenericType<0>,
             NumberType<u64>,
         >(|val, output, _| {
-            let mut hasher = DefaultHasher::default();
-            DFHash::hash(&val.to_owned(), &mut hasher);
-            output.push(hasher.finish());
+            output.push(siphash64(&val.to_owned()));
         }))
         .register();
 
@@ -222,6 +222,93 @@ pub fn register(registry: &mut FunctionRegistry) {
     );
 }
 
+fn register_bucket(registry: &mut FunctionRegistry) {
+    registry
+        .register_passthrough_nullable_2_arg::<NumberType<u64>, StringType, NumberType<u32>, _, _>(
+            "bucket",
+            |_, _, _| FunctionDomain::MayThrow,
+            vectorize_with_builder_2_arg::<NumberType<u64>, StringType, NumberType<u32>>(
+                |buckets, value, output, ctx| {
+                    output.push(bucket(buckets, value, output.len(), ctx));
+                },
+            ),
+        );
+    registry
+        .register_passthrough_nullable_2_arg::<NumberType<u64>, DateType, NumberType<u32>, _, _>(
+            "bucket",
+            |_, _, _| FunctionDomain::MayThrow,
+            vectorize_with_builder_2_arg::<NumberType<u64>, DateType, NumberType<u32>>(
+                |buckets, value, output, ctx| {
+                    output.push(bucket(buckets, &value, output.len(), ctx));
+                },
+            ),
+        );
+    registry.register_passthrough_nullable_2_arg::<
+        NumberType<u64>,
+        TimestampType,
+        NumberType<u32>,
+        _,
+        _,
+    >(
+        "bucket",
+        |_, _, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_2_arg::<NumberType<u64>, TimestampType, NumberType<u32>>(
+            |buckets, value, output, ctx| {
+                output.push(bucket(buckets, &value, 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<u32>,
+                    _,
+                    _,
+                >(
+                    "bucket",
+                    |_, _, _| FunctionDomain::MayThrow,
+                    vectorize_with_builder_2_arg::<
+                        NumberType<u64>,
+                        NumberType<NUM_TYPE>,
+                        NumberType<u32>,
+                    >(|buckets, value, output, ctx| {
+                        output.push(bucket(buckets, &value, output.len(), ctx));
+                    }),
+                );
+            }
+            _ => unreachable!(),
+        });
+    }
+}
+
+fn bucket<T: DFHash + ?Sized>(
+    buckets: u64,
+    value: &T,
+    row: usize,
+    ctx: &mut databend_common_expression::EvalContext,
+) -> u32 {
+    let Ok(buckets) = u32::try_from(buckets) else {
+        ctx.set_error(row, "bucket count must be between 1 and 4294967295");
+        return 0;
+    };
+    if buckets == 0 {
+        ctx.set_error(row, "bucket count must be greater than zero");
+        return 0;
+    }
+
+    (siphash64(value) % buckets as u64) as u32
+}
+
+fn siphash64<T: DFHash + ?Sized>(value: &T) -> u64 {
+    // Keep SQL hash results independent of Rust's unspecified DefaultHasher.
+    let mut hasher = SipHasher13::new_with_keys(0, 0);
+    DFHash::hash(value, &mut hasher);
+    hasher.finish()
+}
+
 fn register_simple_domain_type_hash<T: ArgType>(registry: &mut FunctionRegistry)
 where for<'a> T::ScalarRef<'a>: DFHash {
     registry
@@ -232,9 +319,7 @@ where for<'a> T::ScalarRef<'a>: DFHash {
         .calc_domain(|_, _| FunctionDomain::Full)
         .vectorized(vectorize_with_builder_1_arg::<T, NumberType<u64>>(
             |val, output, _| {
-                let mut hasher = DefaultHasher::default();
-                DFHash::hash(&val, &mut hasher);
-                output.push(hasher.finish());
+                output.push(siphash64(&val));
             },
         ))
         .register();

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) {

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) :: UInt32
+1 bucket(UInt64 NULL, String NULL) :: UInt32 NULL
+2 bucket(UInt64, Date) :: UInt32
+3 bucket(UInt64 NULL, Date NULL) :: UInt32 NULL
+4 bucket(UInt64, Timestamp) :: UInt32
+5 bucket(UInt64 NULL, Timestamp NULL) :: UInt32 NULL
+6 bucket(UInt64, UInt8) :: UInt32
+7 bucket(UInt64 NULL, UInt8 NULL) :: UInt32 NULL
+8 bucket(UInt64, UInt16) :: UInt32
+9 bucket(UInt64 NULL, UInt16 NULL) :: UInt32 NULL
+10 bucket(UInt64, UInt32) :: UInt32
+11 bucket(UInt64 NULL, UInt32 NULL) :: UInt32 NULL
+12 bucket(UInt64, UInt64) :: UInt32
+13 bucket(UInt64 NULL, UInt64 NULL) :: UInt32 NULL
+14 bucket(UInt64, Int8) :: UInt32
+15 bucket(UInt64 NULL, Int8 NULL) :: UInt32 NULL
+16 bucket(UInt64, Int16) :: UInt32
+17 bucket(UInt64 NULL, Int16 NULL) :: UInt32 NULL
+18 bucket(UInt64, Int32) :: UInt32
+19 bucket(UInt64 NULL, Int32 NULL) :: UInt32 NULL
+20 bucket(UInt64, Int64) :: UInt32
+21 bucket(UInt64 NULL, Int64 NULL) :: UInt32 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/hash.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 : 5_u32
+output type    : UInt32
+output domain  : {5..=5}
+output         : 5
+
+
+ast            : bucket(16, 'iceberg')
+raw expr       : bucket(16, 'iceberg')
+checked expr   : bucket<UInt64, String>(CAST<UInt8>(16_u8 AS UInt64), "iceberg")
+optimized expr : 6_u32
+output type    : UInt32
+output domain  : {6..=6}
+output         : 6
+
+
+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 : 0_u32
+output type    : UInt32
+output domain  : {0..=0}
+output         : 0
+
+
+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 : 11_u32
+output type    : UInt32
+output domain  : {11..=11}
+output         : 11
+
+
+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   | UInt32  |
+| Domain | {1..=4} | Unknown |
+| Row 0  | 1       | 1       |
+| Row 1  | 2       | 6       |
+| Row 2  | 3       | 2       |
+| Row 3  | 4       | 1       |
++--------+---------+---------+
+evaluation (internal):
++--------+-----------------------------+
+| Column | Data                        |
++--------+-----------------------------+
+| a      | Column(Int64([1, 2, 3, 4])) |
+| Output | UInt32([1, 6, 2, 1])        |
++--------+-----------------------------+
+
+

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<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);