fix

Author: zhyass

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

@@ -257,17 +257,17 @@ pub fn register(registry: &mut FunctionRegistry) {
 ///
 /// # Example
 /// For boundaries [10, 20, 30]:
-/// - Values < 10 get partition ID 0
-/// - Values >= 10 and < 20 get partition ID 1
-/// - Values >= 20 and < 30 get partition ID 2
-/// - Values >= 30 get partition ID 3
+/// - Values <= 10 get partition ID 0
+/// - Values > 10 and <= 20 get partition ID 1
+/// - Values > 20 and <= 30 get partition ID 2
+/// - Values > 30 get partition ID 3
 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 {
+        if val > bound {
             low = mid + 1;
         } else {
             high = mid;

src/query/service/src/pipelines/processors/transforms/recluster/recluster_sample_state.rs

@@ -35,6 +35,7 @@ impl SampleState {
                 completed_inputs: 0,
                 values: vec![],
                 bounds: vec![],
+                max_value: None,
             }),
             done: Arc::new(WatchNotify::new()),
         })
@@ -56,17 +57,22 @@ impl SampleState {
         Ok(())
     }
 
-    pub fn get_bounds<T>(&self) -> Vec<T::Scalar>
+    pub fn get_bounds<T>(&self) -> (Vec<T::Scalar>, Option<T::Scalar>)
     where
         T: ArgType,
         T::Scalar: Ord,
     {
         let inner = self.inner.read().unwrap();
-        inner
+        let bounds = inner
             .bounds
             .iter()
             .map(|v| T::to_owned_scalar(T::try_downcast_scalar(&v.as_ref()).unwrap()))
-            .collect()
+            .collect();
+        let max_value = inner
+            .max_value
+            .as_ref()
+            .map(|v| T::to_owned_scalar(T::try_downcast_scalar(&v.as_ref()).unwrap()));
+        (bounds, max_value)
     }
 }
 
@@ -76,6 +82,7 @@ pub struct SampleStateInner {
 
     completed_inputs: usize,
     bounds: Vec<Scalar>,
+    max_value: Option<Scalar>,
 
     values: Vec<(u64, Vec<Scalar>)>,
 }
@@ -112,6 +119,9 @@ impl SampleStateInner {
         let col = T::upcast_column(T::column_from_vec(data.clone(), &[]));
         let indices = compare_columns(vec![col], total_samples)?;
 
+        let max_index = indices[total_samples - 1] as usize;
+        let max_val = data[max_index].clone();
+
         let mut cum_weight = 0.0;
         let mut target = step;
         let mut bounds = Vec::with_capacity(self.partitions - 1);
@@ -126,7 +136,13 @@ impl SampleStateInner {
             if cum_weight >= target {
                 let data = &data[idx];
                 if previous_bound.as_ref().is_none_or(|prev| data > prev) {
-                    bounds.push(T::upcast_scalar(data.clone()));
+                    if data == &max_val {
+                        self.max_value = Some(T::upcast_scalar(max_val));
+                        break;
+                    }
+
+                    let bound = T::upcast_scalar(data.clone());
+                    bounds.push(bound);
                     target += step;
                     j += 1;
                     previous_bound = Some(data.clone());

src/query/service/src/pipelines/processors/transforms/recluster/transform_range_partition_indexer.rs

@@ -44,6 +44,7 @@ where T: ArgType
     input_data: Vec<DataBlock>,
     output_data: VecDeque<DataBlock>,
     bounds: Vec<T::Scalar>,
+    max_value: Option<T::Scalar>,
 }
 
 impl<T> TransformRangePartitionIndexer<T>
@@ -63,6 +64,7 @@ where
             input_data: vec![],
             output_data: VecDeque::new(),
             bounds: vec![],
+            max_value: None,
         })
     }
 }
@@ -124,6 +126,7 @@ where
 
     fn process(&mut self) -> Result<()> {
         if let Some(mut block) = self.input_data.pop() {
+            let bound_len = self.bounds.len();
             let num_rows = block.num_rows();
             let last = block.get_last_column().clone();
             block.pop_columns(1);
@@ -132,12 +135,18 @@ where
             for index in 0..num_rows {
                 let val =
                     T::to_owned_scalar(unsafe { T::index_column_unchecked(&last_col, index) });
+                if self.max_value.as_ref().is_some_and(|v| val >= *v) {
+                    let range_id = bound_len + 1;
+                    builder.push(range_id as u64);
+                    continue;
+                }
+
                 let mut low = 0;
-                let mut high = self.bounds.len();
+                let mut high = bound_len;
                 while low < high {
                     let mid = low + ((high - low) / 2);
                     let bound = unsafe { self.bounds.get_unchecked(mid) }.clone();
-                    if val >= bound {
+                    if val > bound {
                         low = mid + 1;
                     } else {
                         high = mid;
@@ -158,7 +167,7 @@ where
     #[async_backtrace::framed]
     async fn async_process(&mut self) -> Result<()> {
         self.state.done.notified().await;
-        self.bounds = self.state.get_bounds::<T>();
+        (self.bounds, self.max_value) = self.state.get_bounds::<T>();
         Ok(())
     }
 }