Merge pull request #10593 from dantengsky/fix-replace-into-panic

fix: replace into panic

Author: BohuTANG

src/query/service/tests/it/storages/fuse/operations/mod.rs

@@ -25,5 +25,6 @@ mod optimize;
 mod purge_drop;
 mod purge_truncate;
 mod read_plan;
+mod replace_into;
 mod table_analyze;
 mod truncate;

src/query/service/tests/it/storages/fuse/operations/replace_into.rs

@@ -0,0 +1,47 @@
+// Copyright 2023 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 common_exception::Result;
+use common_storages_fuse::FuseTable;
+
+#[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)
+            .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

@@ -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)
     }
 
+    pub 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,

tests/sqllogictests/suites/base/09_fuse_engine/09_0023_replace_into

@@ -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