Merge pull request #8452 from RinChanNOWWW/improve_sort

feat(query): improve sort.
datafuselabs · Nov 4, 2022 · 93bb4ba · 93bb4ba · vercel · Nov 4, 2022
2 parents 05e7db2 + bbbf55e
commit 93bb4ba
Show file tree

Hide file tree

Showing 15 changed files with 617 additions and 70 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/common/arrow/Cargo.toml b/src/common/arrow/Cargo.toml
@@ -34,7 +34,7 @@ simd = ["arrow/simd"]
 # Workspace dependencies
 
 # Crates.io dependencies
-arrow = { package = "arrow2", version = "0.14.0", default-features = false, features = [
+arrow = { package = "arrow2", git = "https://github.com/RinChanNOWWW/arrow2", rev = "8bd6417", default-features = false, features = [
     "io_parquet",
     "io_parquet_compression",
 ] }

diff --git a/src/query/datablocks/src/kernels/data_block_sort.rs b/src/query/datablocks/src/kernels/data_block_sort.rs
@@ -133,7 +133,7 @@ impl DataBlock {
             .map(|f| {
                 let left = lhs.try_column_by_name(f.name())?;
                 let right = rhs.try_column_by_name(f.name())?;
-                Self::take_columns_by_slices_limit(
+                Self::take_column_by_slices_limit(
                     f.data_type(),
                     &[left.clone(), right.clone()],
                     &slices,
@@ -210,7 +210,6 @@ impl DataBlock {
 fn compare_variant(left: &dyn Array, right: &dyn Array) -> ArrowResult<DynComparator> {
     let left = VariantColumn::from_arrow_array(left);
     let right = VariantColumn::from_arrow_array(right);
-
     Ok(Box::new(move |i, j| {
         left.get_data(i).cmp(right.get_data(j))
     }))

diff --git a/src/query/datablocks/src/kernels/data_block_take.rs b/src/query/datablocks/src/kernels/data_block_take.rs
@@ -61,7 +61,29 @@ impl DataBlock {
         Ok(DataBlock::create(schema.clone(), result_columns))
     }
 
-    pub fn take_columns_by_slices_limit(
+    pub fn block_take_by_slices_limit(
+        raw: &DataBlock,
+        slice: (usize, usize),
+        limit: Option<usize>,
+    ) -> Result<DataBlock> {
+        let fields = raw.schema().fields();
+        let columns = fields
+            .iter()
+            .map(|f| {
+                let column = raw.try_column_by_name(f.name())?.clone();
+                Self::take_column_by_slices_limit(
+                    f.data_type(),
+                    &[column],
+                    &[(0, slice.0, slice.1)],
+                    limit,
+                )
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let data = DataBlock::create(raw.schema().clone(), columns);
+        Ok(data)
+    }
+
+    pub fn take_column_by_slices_limit(
         data_type: &DataTypeImpl,
         columns: &[ColumnRef],
         slices: &[MergeSlice],

diff --git a/src/query/pipeline/core/src/pipeline_display.rs b/src/query/pipeline/core/src/pipeline_display.rs
@@ -71,21 +71,34 @@ impl<'a> Display for PipelineIndentDisplayWrapper<'a> {
                     ..
                 } => {
                     let prev_name = Self::pipe_name(&pipes[pipes.len() - index - 2]);
-                    let post_name = Self::pipe_name(&pipes[pipes.len() - index]);
-
-                    write!(
-                        f,
-                        "Merge ({} × {} {}) to ({} × {})",
-                        prev_name,
-                        inputs_port.len(),
-                        if inputs_port.len() == 1 {
-                            "processor"
-                        } else {
-                            "processors"
-                        },
-                        post_name,
-                        outputs_port.len(),
-                    )?;
+                    if index > 0 {
+                        let post_name = Self::pipe_name(&pipes[pipes.len() - index]);
+                        write!(
+                            f,
+                            "Merge ({} × {} {}) to ({} × {})",
+                            prev_name,
+                            inputs_port.len(),
+                            if inputs_port.len() == 1 {
+                                "processor"
+                            } else {
+                                "processors"
+                            },
+                            post_name,
+                            outputs_port.len(),
+                        )?;
+                    } else {
+                        write!(
+                            f,
+                            "Merge ({} × {} {})",
+                            prev_name,
+                            inputs_port.len(),
+                            if inputs_port.len() == 1 {
+                                "processor"
+                            } else {
+                                "processors"
+                            },
+                        )?;
+                    }
                 }
             }
         }

diff --git a/src/query/pipeline/transforms/Cargo.toml b/src/query/pipeline/transforms/Cargo.toml
@@ -9,7 +9,9 @@ doctest = false
 test = false
 
 [dependencies]
+common-arrow = { path = "../../../common/arrow" }
 common-datablocks = { path = "../../datablocks" }
+common-datavalues = { path = "../../datavalues" }
 common-exception = { path = "../../../common/exception" }
 common-pipeline-core = { path = "../core" }
 

diff --git a/src/query/pipeline/transforms/src/processors/transforms/mod.rs b/src/query/pipeline/transforms/src/processors/transforms/mod.rs
@@ -16,12 +16,14 @@ pub mod transform;
 pub mod transform_block_compact;
 pub mod transform_compact;
 pub mod transform_limit;
+pub mod transform_multi_sort_merge;
 pub mod transform_sort_merge;
 pub mod transform_sort_partial;
 
 pub use transform::*;
 pub use transform_block_compact::*;
 pub use transform_compact::*;
 pub use transform_limit::*;
+pub use transform_multi_sort_merge::*;
 pub use transform_sort_merge::*;
 pub use transform_sort_partial::*;