feat: make RegionScanner aware of PartitionRange (#4170)

* define PartitionRange

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* add optimizer rule

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* implement interfaces

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* impl aggr stream

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* add fallback method

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix tests

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix tests

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix typo

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* update sqlness result

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* add document and rename struct

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* add more comments

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix typo

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

src/query/src/lib.rs

@@ -15,6 +15,7 @@
 #![feature(let_chains)]
 #![feature(int_roundings)]
 #![feature(option_get_or_insert_default)]
+#![feature(trait_upcasting)]
 
 mod analyze;
 pub mod dataframe;

src/query/src/optimizer.rs

@@ -14,6 +14,7 @@
 
 pub mod count_wildcard;
 pub mod order_hint;
+pub mod parallelize_scan;
 pub mod remove_duplicate;
 pub mod string_normalization;
 #[cfg(test)]

src/query/src/optimizer/parallelize_scan.rs

@@ -0,0 +1,97 @@
+// Copyright 2023 Greptime Team
+//
+// 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 common_telemetry::debug;
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::{DataFusionError, Result};
+use store_api::region_engine::PartitionRange;
+use table::table::scan::RegionScanExec;
+
+pub struct ParallelizeScan;
+
+impl PhysicalOptimizerRule for ParallelizeScan {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Self::do_optimize(plan, config)
+    }
+
+    fn name(&self) -> &str {
+        "parallelize_scan"
+    }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
+}
+
+impl ParallelizeScan {
+    fn do_optimize(
+        plan: Arc<dyn ExecutionPlan>,
+        config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let result = plan
+            .transform_down(|plan| {
+                if let Some(region_scan_exec) = plan.as_any().downcast_ref::<RegionScanExec>() {
+                    let ranges = region_scan_exec.get_partition_ranges();
+                    let total_range_num = ranges.len();
+                    let expected_partition_num = config.execution.target_partitions;
+
+                    // assign ranges to each partition
+                    let partition_ranges =
+                        Self::assign_partition_range(ranges, expected_partition_num);
+                    debug!(
+                        "Assign {total_range_num} ranges to {expected_partition_num} partitions"
+                    );
+
+                    // update the partition ranges
+                    region_scan_exec
+                        .set_partitions(partition_ranges)
+                        .map_err(|e| DataFusionError::External(e.into_inner()))?;
+                }
+
+                // The plan might be modified, but it's modified in-place so we always return
+                // Transformed::no(plan) to indicate there is no "new child"
+                Ok(Transformed::no(plan))
+            })?
+            .data;
+
+        Ok(result)
+    }
+
+    /// Distribute [`PartitionRange`]s to each partition.
+    ///
+    /// Currently we use a simple round-robin strategy to assign ranges to partitions.
+    fn assign_partition_range(
+        ranges: Vec<PartitionRange>,
+        expected_partition_num: usize,
+    ) -> Vec<Vec<PartitionRange>> {
+        let mut partition_ranges = vec![vec![]; expected_partition_num];
+
+        // round-robin assignment
+        for (i, range) in ranges.into_iter().enumerate() {
+            let partition_idx = i % expected_partition_num;
+            partition_ranges[partition_idx].push(range);
+        }
+
+        partition_ranges
+    }
+}

src/query/src/query_engine/state.rs

@@ -43,6 +43,7 @@ use table::TableRef;
 use crate::dist_plan::{DistExtensionPlanner, DistPlannerAnalyzer};
 use crate::optimizer::count_wildcard::CountWildcardToTimeIndexRule;
 use crate::optimizer::order_hint::OrderHintRule;
+use crate::optimizer::parallelize_scan::ParallelizeScan;
 use crate::optimizer::remove_duplicate::RemoveDuplicate;
 use crate::optimizer::string_normalization::StringNormalizationRule;
 use crate::optimizer::type_conversion::TypeConversionRule;
@@ -112,6 +113,11 @@ impl QueryEngineState {
 
         // add physical optimizer
         let mut physical_optimizer = PhysicalOptimizer::new();
+        // Change TableScan's partition at first
+        physical_optimizer
+            .rules
+            .insert(0, Arc::new(ParallelizeScan));
+        // Add rule to remove duplicate nodes generated by other rules. Run this in the last.
         physical_optimizer.rules.push(Arc::new(RemoveDuplicate));
 
         let session_state = SessionState::new_with_config_rt(session_config, runtime_env)