fix replace into panic

Author: dantengsky

Cargo.lock

@@ -56,3 +56,6 @@ tokio-rayon = "2.1.0"
 tracing = "0.1.36"
 typetag = "0.2.3"
 uuid = { version = "1.1.2", features = ["serde", "v4"] }
+
+[dev-dependencies]
+rand = "0.8.5"

src/query/storages/fuse/Cargo.toml

@@ -25,6 +25,7 @@ use common_pipeline_core::pipe::PipeItem;
 use common_pipeline_core::processors::processor::ProcessorPtr;
 use common_pipeline_transforms::processors::transforms::create_dummy_item;
 use common_pipeline_transforms::processors::transforms::AsyncAccumulatingTransformer;
+use storages_common_table_meta::meta::Location;
 use storages_common_table_meta::meta::Statistics;
 use storages_common_table_meta::meta::TableSnapshot;
 use uuid::Uuid;
@@ -37,6 +38,7 @@ use crate::operations::merge_into::CommitSink;
 use crate::operations::merge_into::MergeIntoOperationAggregator;
 use crate::operations::merge_into::OnConflictField;
 use crate::operations::merge_into::TableMutationAggregator;
+use crate::operations::mutation::base_mutator::SegmentIndex;
 use crate::operations::replace_into::processor_replace_into::ReplaceIntoProcessor;
 use crate::pipelines::Pipeline;
 use crate::FuseTable;
@@ -259,28 +261,7 @@ impl FuseTable {
         on_conflicts: Vec<OnConflictField>,
         table_snapshot: &TableSnapshot,
     ) -> Result<Vec<PipeItem>> {
-        let segments = table_snapshot.segments.as_slice();
-        let chunk_size = segments.len() / num_partition;
-        // caller site should guarantee this
-        assert!(chunk_size >= 1);
-
-        let mut chunks = vec![];
-        let mut chunk = vec![];
-        for (segment_idx, segment_location) in segments.iter().enumerate() {
-            chunk.push((segment_idx, segment_location.clone()));
-            if (segment_idx + 1) % chunk_size == 0 {
-                chunks.push(std::mem::take(&mut chunk))
-            }
-        }
-
-        if !chunk.is_empty() {
-            if chunks.len() == num_partition {
-                chunks.last_mut().unwrap().append(&mut chunk);
-            } else {
-                chunks.push(std::mem::take(&mut chunk))
-            }
-        }
-
+        let chunks = Self::partition_segments(&table_snapshot.segments, num_partition);
         let read_settings = ReadSettings::from_ctx(&ctx)?;
         let mut items = vec![];
         for chunk_of_segment_locations in chunks {
@@ -299,6 +280,28 @@ impl FuseTable {
         Ok(items)
     }
 
+    fn partition_segments(
+        segments: &[Location],
+        num_partition: usize,
+    ) -> Vec<Vec<(SegmentIndex, Location)>> {
+        let chunk_size = segments.len() / num_partition;
+        // caller site guarantees this
+        assert!(chunk_size >= 1);
+
+        let mut chunks = vec![];
+        for (chunk_idx, chunk) in segments.chunks(chunk_size).enumerate() {
+            let mut segment_chunk = (chunk_idx * chunk_size..)
+                .zip(chunk.to_vec())
+                .collect::<Vec<_>>();
+            if chunks.len() < num_partition {
+                chunks.push(segment_chunk);
+            } else {
+                chunks.last_mut().unwrap().append(&mut segment_chunk);
+            }
+        }
+        chunks
+    }
+
     fn create_append_transform(&self, ctx: Arc<dyn TableContext>) -> AppendTransform {
         AppendTransform::try_create(
             ctx,
@@ -362,3 +365,42 @@ impl FuseTable {
         )
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_partition() -> Result<()> {
+        use rand::Rng;
+
+        let mut rng = rand::thread_rng();
+
+        for _ in 0..100 {
+            let number_segment: usize = rng.gen_range(1..100);
+
+            // do not matter, arbitrarily picked
+            let format_version = 2;
+
+            let segments = (0..number_segment)
+                .into_iter()
+                .map(|idx| (format!("{idx}"), format_version))
+                .collect::<Vec<_>>();
+
+            for _ in 0..100 {
+                let num_partition: usize = if number_segment == 1 {
+                    1
+                } else {
+                    rng.gen_range(1..number_segment)
+                };
+
+                let chunks = FuseTable::partition_segments(&segments, num_partition);
+                assert_eq!(chunks.len(), num_partition);
+                for (idx, (segment_idx, _)) in chunks.clone().into_iter().flatten().enumerate() {
+                    assert_eq!(idx, segment_idx)
+                }
+            }
+        }
+        Ok(())
+    }
+}

src/query/storages/fuse/src/operations/replace.rs

@@ -323,6 +323,44 @@ DROP TABLE test;
 statement ok
 drop stage test_stage
 
+###############################################################
+# issue https://github.com/datafuselabs/databend/issues/10572 #
+###############################################################
+
+# with 5 segments and max_threads set to 3, `replace into` leads to
+# ERROR HY000 (1105): Code: 1068, displayText = Cannot join handle from context's runtime
+
+statement ok
+CREATE TABLE test(a int);
+
+statement ok
+insert into test values(1);
+
+statement ok
+insert into test values(2);
+
+statement ok
+insert into test values(3);
+
+statement ok
+insert into test values(4);
+
+statement ok
+insert into test values(5);
+
+statement ok
+set max_threads = 3;
+
+# if not fixed: ERROR HY000 (1105): Code: 1068, displayText = Cannot join handle from context's runtime
+statement ok
+replace into test on(a) values(6);
+
+statement ok
+drop table test
+
+
+#####################
+
 
 statement ok
 DROP DATABASE db_09_0023