fix

Author: zhyass

src/query/functions/src/scalars/hilbert.rs

@@ -24,9 +24,11 @@ use databend_common_expression::types::ALL_NUMERICS_TYPES;
 use databend_common_expression::vectorize_with_builder_1_arg;
 use databend_common_expression::vectorize_with_builder_2_arg;
 use databend_common_expression::with_number_mapped_type;
+use databend_common_expression::Column;
 use databend_common_expression::FixedLengthEncoding;
 use databend_common_expression::FunctionDomain;
 use databend_common_expression::FunctionRegistry;
+use databend_common_expression::ScalarRef;
 
 pub fn register(registry: &mut FunctionRegistry) {
     registry.register_passthrough_nullable_1_arg::<StringType, BinaryType, _, _>(
@@ -92,22 +94,45 @@ pub fn register(registry: &mut FunctionRegistry) {
     // The column values are conceptually divided into multiple partitions defined by the range_bounds.
     // For example, given the column values (0, 1, 3, 6, 8) and a partition configuration with 3 partitions,
     // the range_bounds might be [1, 6]. The function would then return partition IDs as (0, 0, 1, 1, 2).
-    registry.register_passthrough_nullable_2_arg::<GenericType<0>, ArrayType<GenericType<0>>, NumberType<u64>, _, _>(
+    registry
+        .register_2_arg_core::<GenericType<0>, ArrayType<GenericType<0>>, NumberType<u64>, _, _>(
+            "range_partition_id",
+            |_, _, _| FunctionDomain::Full,
+            vectorize_with_builder_2_arg::<
+                GenericType<0>,
+                ArrayType<GenericType<0>>,
+                NumberType<u64>,
+            >(|val, arr, builder, _| {
+                let id = calc_range_partition_id(val, arr);
+                builder.push(id);
+            }),
+        );
+
+    registry.register_2_arg_core::<NullableType<GenericType<0>>, NullableType<ArrayType<GenericType<0>>>, NumberType<u64>, _, _>(
         "range_partition_id",
         |_, _, _| FunctionDomain::Full,
-        vectorize_with_builder_2_arg::<GenericType<0>, ArrayType<GenericType<0>>, NumberType<u64>>(|val, arr, builder, _| {
-            let mut low = 0;
-            let mut high = arr.len();
-            while low < high {
-                let mid = low + ((high - low) / 2);
-                let bound = unsafe {arr.index_unchecked(mid)};
-                if val > bound {
-                    low = mid + 1;
-                } else {
-                    high = mid;
-                }
-            }
-            builder.push(low as u64);
+        vectorize_with_builder_2_arg::<NullableType<GenericType<0>>, NullableType<ArrayType<GenericType<0>>>, NumberType<u64>>(|val, arr, builder, _| {
+            let id = match (val, arr) {
+                (Some(val), Some(arr)) => calc_range_partition_id(val, arr),
+                (None, Some(arr)) => arr.len() as u64 + 1,
+                _ => 0,
+            };
+            builder.push(id);
         }),
     );
 }
+
+fn calc_range_partition_id(val: ScalarRef, arr: Column) -> u64 {
+    let mut low = 0;
+    let mut high = arr.len();
+    while low < high {
+        let mid = low + ((high - low) / 2);
+        let bound = unsafe { arr.index_unchecked(mid) };
+        if val > bound {
+            low = mid + 1;
+        } else {
+            high = mid;
+        }
+    }
+    low as u64
+}

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

@@ -3211,7 +3211,7 @@ Functions overloads:
 0 range(UInt64, UInt64) :: Array(UInt64)
 1 range(UInt64 NULL, UInt64 NULL) :: Array(UInt64) NULL
 0 range_partition_id(T0, Array(T0)) :: UInt64
-1 range_partition_id(T0 NULL, Array(T0) NULL) :: UInt64 NULL
+1 range_partition_id(T0 NULL, Array(T0) NULL) :: UInt64
 0 regexp(String, String) :: Boolean
 1 regexp(String NULL, String NULL) :: Boolean NULL
 0 regexp_instr FACTORY

src/query/service/src/interpreters/interpreter_table_recluster.rs

@@ -279,8 +279,7 @@ impl ReclusterTableInterpreter {
                         "range_bound({partitions}, {sample_size})({cluster_key_str}) AS bound_{index}"
                     ));
                     hilbert_keys.push(format!(
-                        "hilbert_key(cast(ifnull(range_partition_id({table}.{cluster_key_str}, \
-                        _keys_bound.bound_{index}), {partitions}) as uint16))"
+                        "hilbert_key(cast(range_partition_id({table}.{cluster_key_str}, _keys_bound.bound_{index}) as uint16))"
                     ));
                 }
                 let keys_bounds_str = keys_bounds.join(", ");

src/query/service/src/pipelines/processors/transforms/transform_partition_collect.rs

@@ -54,8 +54,12 @@ impl Exchange for HilbertPartitionExchange {
     fn partition(&self, data_block: DataBlock, n: usize) -> Result<Vec<DataBlock>> {
         // Extract the columns used for hash computation.
         let mut data_block = data_block.consume_convert_to_full();
-        let last = data_block.get_last_column().as_nullable().unwrap();
-        let range_ids = last.column.as_number().unwrap().as_u_int64().unwrap();
+        let range_ids = data_block
+            .get_last_column()
+            .as_number()
+            .unwrap()
+            .as_u_int64()
+            .unwrap();
 
         // Scatter the data block to different partitions.
         let indices = range_ids