update

Author: zhyass

src/common/base/src/base/watch_notify.rs

@@ -44,6 +44,10 @@ impl WatchNotify {
     pub fn notify_waiters(&self) {
         let _ = self.tx.send_replace(true);
     }
+
+    pub fn is_notified(&self) -> bool {
+        *self.rx.borrow()
+    }
 }
 
 #[cfg(test)]

src/query/functions/src/aggregates/aggregate_range_bound.rs

@@ -326,9 +326,7 @@ pub fn try_create_aggregate_range_bound_function(
 /// For a column with values `(0, 1, 3, 6, 8)` and `partition_num = 3`, the function calculates the
 /// partition boundaries based on the distribution of the data. The boundaries might be `[1, 6]`.
 pub fn aggregate_range_bound_function_desc() -> AggregateFunctionDescription {
-    AggregateFunctionDescription::creator(Box::new(
-        crate::aggregates::try_create_aggregate_range_bound_function,
-    ))
+    AggregateFunctionDescription::creator(Box::new(try_create_aggregate_range_bound_function))
 }
 
 fn get_partitions(

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

@@ -0,0 +1,198 @@
+// 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::Result;
+use databend_common_expression::row::RowConverter as CommonConverter;
+use databend_common_expression::types::AccessType;
+use databend_common_expression::types::ArgType;
+use databend_common_expression::types::DataType;
+use databend_common_expression::types::DateType;
+use databend_common_expression::types::NumberDataType;
+use databend_common_expression::types::NumberType;
+use databend_common_expression::types::StringType;
+use databend_common_expression::types::TimestampType;
+use databend_common_expression::with_number_mapped_type;
+use databend_common_expression::DataSchemaRef;
+use databend_common_expression::SortColumnDescription;
+use databend_common_pipeline_core::processors::InputPort;
+use databend_common_pipeline_core::processors::OutputPort;
+use databend_common_pipeline_core::processors::Processor;
+use databend_common_pipeline_transforms::sort::CommonRows;
+use databend_common_pipeline_transforms::sort::RowConverter;
+use databend_common_pipeline_transforms::sort::Rows;
+use databend_common_pipeline_transforms::sort::SimpleRowConverter;
+use databend_common_pipeline_transforms::sort::SimpleRowsAsc;
+use databend_common_pipeline_transforms::AccumulatingTransformer;
+use databend_common_pipeline_transforms::Transformer;
+use match_template::match_template;
+
+use crate::pipelines::processors::transforms::recluster::transform_add_order_column::TransformAddOrderColumn;
+use crate::pipelines::processors::transforms::recluster::TransformRangePartitionIndexer;
+use crate::pipelines::processors::transforms::SampleState;
+use crate::pipelines::processors::transforms::TransformReclusterCollect;
+
+pub struct TransformReclusterBuilder {
+    schema: DataSchemaRef,
+    sort_desc: Arc<[SortColumnDescription]>,
+    sample_rate: f64,
+    seed: u64,
+}
+
+impl TransformReclusterBuilder {
+    pub fn build_recluster_sample(
+        &self,
+        input: Arc<InputPort>,
+        output: Arc<OutputPort>,
+    ) -> Result<Box<dyn Processor>> {
+        self.build_inner(BuilderType::ReclusterSample, input, output, None)
+    }
+
+    pub fn build_range_partition_indexer(
+        &self,
+        input: Arc<InputPort>,
+        output: Arc<OutputPort>,
+        state: Arc<SampleState>,
+    ) -> Result<Box<dyn Processor>> {
+        self.build_inner(
+            BuilderType::RangePartitionIndexer,
+            input,
+            output,
+            Some(state),
+        )
+    }
+
+    pub fn build_add_order_column(
+        &self,
+        input: Arc<InputPort>,
+        output: Arc<OutputPort>,
+    ) -> Result<Box<dyn Processor>> {
+        self.build_inner(BuilderType::AddOrderColumn, input, output, None)
+    }
+
+    fn build_inner(
+        &self,
+        typ: BuilderType,
+        input: Arc<InputPort>,
+        output: Arc<OutputPort>,
+        state: Option<Arc<SampleState>>,
+    ) -> Result<Box<dyn Processor>> {
+        let mut build = BuilderInner {
+            input,
+            output,
+            typ,
+            base: self,
+            state,
+        };
+        build.select_row_type()
+    }
+}
+
+enum BuilderType {
+    AddOrderColumn,
+    ReclusterSample,
+    RangePartitionIndexer,
+}
+
+struct BuilderInner<'a> {
+    input: Arc<InputPort>,
+    output: Arc<OutputPort>,
+    typ: BuilderType,
+    base: &'a TransformReclusterBuilder,
+    state: Option<Arc<SampleState>>,
+}
+
+impl BuilderInner<'_> {
+    pub fn select_row_type(&mut self) -> Result<Box<dyn Processor>> {
+        match self.base.sort_desc.as_ref() {
+            [desc] => {
+                let schema = self.base.schema.clone();
+                let sort_type = schema.field(desc.offset).data_type();
+                assert!(desc.asc);
+
+                match_template! {
+                    T = [ Date => DateType, Timestamp => TimestampType, String => StringType ],
+                    match sort_type {
+                        DataType::T => {
+                            self.visit_type::<SimpleRowsAsc<T>, SimpleRowConverter<T>>()
+                        },
+                        DataType::Number(num_ty) => with_number_mapped_type!(|NUM_TYPE| match num_ty {
+                            NumberDataType::NUM_TYPE => {
+                                self.visit_type::<SimpleRowsAsc<NumberType<NUM_TYPE>>, SimpleRowConverter<NumberType<NUM_TYPE>>>()
+                            }
+                        }),
+                        _ => self.visit_type::<CommonRows, CommonConverter>()
+                    }
+                }
+            }
+            _ => self.visit_type::<CommonRows, CommonConverter>(),
+        }
+    }
+
+    fn visit_type<R, C>(&mut self) -> Result<Box<dyn Processor>>
+    where
+        R: Rows + 'static,
+        C: RowConverter<R> + Send + 'static,
+        R::Type: ArgType + Send + Sync,
+        <R::Type as AccessType>::Scalar: Ord + Send + Sync,
+    {
+        match self.typ {
+            BuilderType::AddOrderColumn => self.build_add_order_column::<R, C>(),
+            BuilderType::ReclusterSample => self.build_recluster_sample::<R::Type>(),
+            BuilderType::RangePartitionIndexer => self.build_range_partition_indexer::<R::Type>(),
+        }
+    }
+
+    fn build_add_order_column<R, C>(&mut self) -> Result<Box<dyn Processor>>
+    where
+        R: Rows + 'static,
+        C: RowConverter<R> + Send + 'static,
+    {
+        let inner = TransformAddOrderColumn::<R, C>::try_new(
+            self.base.sort_desc.clone(),
+            self.base.schema.clone(),
+        )?;
+        Ok(Transformer::create(
+            self.input.clone(),
+            self.output.clone(),
+            inner,
+        ))
+    }
+
+    fn build_range_partition_indexer<T>(&mut self) -> Result<Box<dyn Processor>>
+    where
+        T: ArgType + Send + Sync,
+        T::Scalar: Ord + Send + Sync,
+    {
+        Ok(TransformRangePartitionIndexer::<T>::create(
+            self.input.clone(),
+            self.output.clone(),
+            self.state.clone().unwrap(),
+        ))
+    }
+
+    fn build_recluster_sample<T>(&mut self) -> Result<Box<dyn Processor>>
+    where
+        T: ArgType + Send + Sync,
+        T::Scalar: Ord + Send + Sync,
+    {
+        let offset = self.base.schema.fields().len();
+        Ok(AccumulatingTransformer::create(
+            self.input.clone(),
+            self.output.clone(),
+            TransformReclusterCollect::<T>::new(offset, self.base.sample_rate, self.base.seed),
+        ))
+    }
+}

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

@@ -12,9 +12,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod builder;
+mod range_bound_sampler;
 mod recluster_partition_exchange;
 mod recluster_partition_strategy;
+mod recluster_sample_state;
+mod transform_add_order_column;
+mod transform_range_partition_indexer;
+mod transform_recluster_collect;
 
+pub use range_bound_sampler::RangeBoundSampler;
 pub use recluster_partition_exchange::ReclusterPartitionExchange;
 pub use recluster_partition_strategy::CompactPartitionStrategy;
 pub use recluster_partition_strategy::ReclusterPartitionStrategy;
+pub(crate) use recluster_sample_state::SampleState;
+pub use transform_range_partition_indexer::TransformRangePartitionIndexer;
+pub(crate) use transform_recluster_collect::ReclusterSampleMeta;
+pub use transform_recluster_collect::TransformReclusterCollect;

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

@@ -0,0 +1,86 @@
+// 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::marker::PhantomData;
+
+use databend_common_expression::types::ArgType;
+use databend_common_expression::types::ValueType;
+use databend_common_expression::DataBlock;
+use databend_common_expression::Scalar;
+use rand::prelude::SliceRandom;
+use rand::prelude::SmallRng;
+use rand::SeedableRng;
+
+pub struct RangeBoundSampler<T>
+where T: ValueType
+{
+    offset: usize,
+    sample_rate: f64,
+    rng: SmallRng,
+
+    values: Vec<(u64, Vec<Scalar>)>,
+    _t: PhantomData<T>,
+}
+
+impl<T> RangeBoundSampler<T>
+where T: ValueType
+{
+    pub fn new(offset: usize, sample_rate: f64, seed: u64) -> Self {
+        let rng = SmallRng::seed_from_u64(seed);
+        Self {
+            offset,
+            sample_rate,
+            rng,
+            values: vec![],
+            _t: PhantomData,
+        }
+    }
+}
+
+impl<T> RangeBoundSampler<T>
+where
+    T: ArgType,
+    T::Scalar: Ord + Send,
+{
+    pub fn add_block(&mut self, data: &DataBlock) {
+        let rows = data.num_rows();
+        assert!(rows > 0);
+        let column = data.get_by_offset(self.offset).to_column(rows);
+
+        let sample_size = std::cmp::max((self.sample_rate * rows as f64).ceil() as usize, 100);
+        let mut indices = (0..rows).collect::<Vec<_>>();
+
+        let sampled_indices = if rows > sample_size {
+            indices.shuffle(&mut self.rng);
+            &indices[..sample_size]
+        } else {
+            &indices
+        };
+
+        let column = T::try_downcast_column(&column).unwrap();
+        let sample_values = sampled_indices
+            .iter()
+            .map(|i| {
+                T::upcast_scalar(T::to_owned_scalar(unsafe {
+                    T::index_column_unchecked(&column, *i)
+                }))
+            })
+            .collect::<Vec<_>>();
+        self.values.push((rows as u64, sample_values));
+    }
+
+    pub fn sample_values(&mut self) -> Vec<(u64, Vec<Scalar>)> {
+        std::mem::take(&mut self.values)
+    }
+}

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

@@ -33,8 +33,7 @@ impl ReclusterPartitionExchange {
 
 impl Exchange for ReclusterPartitionExchange {
     const NAME: &'static str = "Recluster";
-    fn partition(&self, data_block: DataBlock, n: usize) -> Result<Vec<DataBlock>> {
-        let mut data_block = data_block;
+    fn partition(&self, mut data_block: DataBlock, n: usize) -> Result<Vec<DataBlock>> {
         let range_ids = data_block
             .get_last_column()
             .as_number()