feat(query): Vector type support recluster using the K-Means algorithm

src/query/service/src/physical_plans/physical_column_mutation.rs

@@ -18,6 +18,7 @@ use std::collections::HashMap;
 use databend_common_catalog::table::Table;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_expression::DataSchema;
 use databend_common_expression::DataSchemaRef;
 use databend_common_expression::RemoteExpr;
 use databend_common_functions::BUILTIN_FUNCTIONS;
@@ -189,7 +190,8 @@ impl IPhysicalPlan for ColumnMutation {
 
         let block_thresholds = table.get_block_thresholds();
         let cluster_stats_gen = if matches!(self.mutation_kind, MutationKind::Delete) {
-            table.get_cluster_stats_gen(builder.ctx.clone(), 0, block_thresholds, None)?
+            let input_schema = DataSchema::from(table.schema_with_stream()).into();
+            table.get_cluster_stats_gen(builder.ctx.clone(), 0, block_thresholds, input_schema)?
         } else {
             table.cluster_gen_for_append(
                 builder.ctx.clone(),

src/query/service/src/physical_plans/physical_multi_table_insert.rs

@@ -691,7 +691,7 @@ impl IPhysicalPlan for ChunkAppendData {
                 builder.ctx.clone(),
                 0,
                 block_thresholds,
-                Some(schema),
+                schema,
             )?;
             let operators = cluster_stats_gen.operators.clone();
             if !operators.is_empty() {

src/query/service/src/physical_plans/physical_mutation.rs

@@ -165,8 +165,9 @@ impl IPhysicalPlan for Mutation {
         let table = FuseTable::try_from_table(tbl.as_ref())?;
         let block_thresholds = table.get_block_thresholds();
 
+        let input_schema = DataSchema::from(table.schema_with_stream()).into();
         let cluster_stats_gen =
-            table.get_cluster_stats_gen(builder.ctx.clone(), 0, block_thresholds, None)?;
+            table.get_cluster_stats_gen(builder.ctx.clone(), 0, block_thresholds, input_schema)?;
 
         let max_threads = builder.settings.get_max_threads()? as usize;
         let io_request_semaphore = Arc::new(Semaphore::new(max_threads));

src/query/service/src/physical_plans/physical_recluster.rs

@@ -24,9 +24,9 @@ use databend_common_catalog::table::Table;
 use databend_common_config::GlobalConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_expression::DataSchema;
 use databend_common_expression::DataSchemaRef;
 use databend_common_expression::DataSchemaRefExt;
-use databend_common_expression::SortColumnDescription;
 use databend_common_io::constants::DEFAULT_BLOCK_BUFFER_SIZE;
 use databend_common_meta_app::schema::TableInfo;
 use databend_common_metrics::storage::metrics_inc_recluster_block_bytes_to_read;
@@ -44,6 +44,7 @@ use databend_common_sql::executor::physical_plans::MutationKind;
 use databend_common_storages_fuse::FUSE_OPT_KEY_BLOCK_IN_MEM_SIZE_THRESHOLD;
 use databend_common_storages_fuse::FuseTable;
 use databend_common_storages_fuse::operations::TransformSerializeBlock;
+use databend_common_storages_fuse::operations::TransformVectorClusterKmeans;
 use databend_common_storages_fuse::statistics::ClusterStatsGenerator;
 use databend_storages_common_cache::TempDirManager;
 use databend_storages_common_table_meta::meta::TableMetaTimestamps;
@@ -185,11 +186,12 @@ impl IPhysicalPlan for Recluster {
                         .add_transformer(|| TransformAddStreamColumns::new(stream_ctx.clone()));
                 }
 
+                let input_schema = DataSchema::from(table.schema_with_stream()).into();
                 let cluster_stats_gen = table.get_cluster_stats_gen(
                     builder.ctx.clone(),
                     task.level + 1,
                     block_thresholds,
-                    None,
+                    input_schema,
                 )?;
                 let operators = cluster_stats_gen.operators.clone();
                 if !operators.is_empty() {
@@ -203,27 +205,33 @@ impl IPhysicalPlan for Recluster {
                     });
                 }
 
-                // construct output fields
-                let output_fields = cluster_stats_gen.out_fields.clone();
-                let schema = DataSchemaRefExt::create(output_fields);
-                let sort_descs: Vec<_> = cluster_stats_gen
-                    .cluster_key_index
-                    .iter()
-                    .map(|offset| SortColumnDescription {
-                        offset: *offset,
-                        asc: true,
-                        nulls_first: false,
-                    })
-                    .collect();
+                let settings = builder.ctx.get_settings();
+                let max_threads = settings.get_max_threads()? as usize;
 
-                // merge sort
                 let (rows_per_block, bytes_per_block) = block_thresholds.calc_rows_for_recluster(
                     task.total_rows,
                     task.total_bytes,
                     task.total_compressed,
                 );
 
-                let settings = builder.ctx.get_settings();
+                if let Some(vector_operator) = cluster_stats_gen.vector_operator.clone() {
+                    builder.main_pipeline.try_resize(1)?;
+                    builder.main_pipeline.add_accumulating_transformer(move || {
+                        TransformVectorClusterKmeans::new(
+                            vector_operator.vector_column_input_offset,
+                            vector_operator.info.distance_type,
+                            rows_per_block,
+                        )
+                    });
+                    builder.main_pipeline.try_resize(max_threads)?;
+                }
+
+                // construct output fields
+                let output_fields = cluster_stats_gen.out_fields.clone();
+                let schema = DataSchemaRefExt::create(output_fields);
+                let sort_descs = cluster_stats_gen.sort_descs();
+
+                // merge sort
                 let sort_pipeline_builder = SortPipelineBuilder::create(
                     builder.ctx.clone(),
                     schema,
@@ -241,7 +249,6 @@ impl IPhysicalPlan for Recluster {
                 let compact_thresholds = block_thresholds
                     .set_rows_per_block(rows_per_block)
                     .set_bytes_per_block(bytes_per_block);
-                let max_threads = settings.get_max_threads()? as usize;
                 build_ordered_compact_pipeline(
                     &mut builder.main_pipeline,
                     compact_thresholds,

src/query/service/src/physical_plans/physical_replace_into.rs

@@ -114,12 +114,8 @@ impl IPhysicalPlan for ReplaceInto {
             .build_table_by_table_info(&self.table_info, None)?;
         let table = FuseTable::try_from_table(table.as_ref())?;
         let schema = DataSchema::from(table.schema()).into();
-        let cluster_stats_gen = table.get_cluster_stats_gen(
-            builder.ctx.clone(),
-            0,
-            self.block_thresholds,
-            Some(schema),
-        )?;
+        let cluster_stats_gen =
+            table.get_cluster_stats_gen(builder.ctx.clone(), 0, self.block_thresholds, schema)?;
 
         // connect to broadcast processor and append transform
         let serialize_block_transform = TransformSerializeBlock::try_create(

src/query/service/tests/it/storages/fuse/bloom_index_meta_size.rs

@@ -341,6 +341,7 @@ fn build_test_segment_info(
         spatial_index_size: None,
         spatial_index_location: None,
         spatial_stats: None,
+        vector_stats: None,
         virtual_block_meta: None,
         compression: Compression::Lz4,
         create_on: Some(Utc::now()),