Merge pull request #7595 from sandflee/slow2

feat(hive): delay start the worker for simple select hive query

Author: BohuTANG

src/query/storages/hive/src/hive_table.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashSet;
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -32,6 +33,7 @@ use common_legacy_planners::Extras;
 use common_legacy_planners::Partitions;
 use common_legacy_planners::Projection;
 use common_legacy_planners::ReadDataSourcePlan;
+use common_legacy_planners::RequireColumnsVisitor;
 use common_legacy_planners::Statistics;
 use common_legacy_planners::TruncateTablePlan;
 use common_meta_app::schema::TableInfo;
@@ -109,19 +111,69 @@ impl HiveTable {
         let max_threads = std::cmp::min(parts_len, max_threads);
 
         let mut source_builder = SourcePipeBuilder::create();
+        let delay_timer = if self.is_simple_select_query(plan) {
+            // 0, 0, 200, 200, 400,400
+            |x: usize| (x / 2).min(10) * 200
+        } else {
+            |_| 0
+        };
 
-        for _index in 0..std::cmp::max(1, max_threads) {
+        for index in 0..std::cmp::max(1, max_threads) {
             let output = OutputPort::create();
             source_builder.add_source(
                 output.clone(),
-                HiveTableSource::create(ctx.clone(), output, block_reader.clone())?,
+                HiveTableSource::create(
+                    ctx.clone(),
+                    output,
+                    block_reader.clone(),
+                    delay_timer(index),
+                )?,
             );
         }
 
         pipeline.add_pipe(source_builder.finalize());
         Ok(())
     }
 
+    // simple select query is the sql likes `select * from xx limit 10` or
+    // `select * from xx where p_date = '20220201' limit 10` where p_date is a partition column;
+    // we just need to read few datas from table
+    fn is_simple_select_query(&self, plan: &ReadDataSourcePlan) -> bool {
+        // couldn't get groupby order by info
+        if let Some(Extras {
+            filters: f,
+            limit: Some(lm),
+            ..
+        }) = &plan.push_downs
+        {
+            if *lm > 100000 {
+                return false;
+            }
+
+            // filter out the partition column related expressions
+            let partition_keys = self.get_partition_key_sets();
+            let columns = Self::get_columns_from_expressions(f);
+            if columns.difference(&partition_keys).count() == 0 {
+                return true;
+            }
+        }
+        false
+    }
+
+    fn get_partition_key_sets(&self) -> HashSet<String> {
+        match &self.table_options.partition_keys {
+            Some(v) => v.iter().cloned().collect::<HashSet<_>>(),
+            None => HashSet::new(),
+        }
+    }
+
+    fn get_columns_from_expressions(expressions: &[Expression]) -> HashSet<String> {
+        expressions
+            .iter()
+            .flat_map(|e| RequireColumnsVisitor::collect_columns_from_expr(e).unwrap())
+            .collect::<HashSet<_>>()
+    }
+
     fn create_block_reader(
         &self,
         ctx: &Arc<dyn TableContext>,

src/query/storages/hive/src/hive_table_source.rs

@@ -15,6 +15,8 @@
 use std::any::Any;
 use std::sync::Arc;
 
+use common_base::base::tokio::time::sleep;
+use common_base::base::tokio::time::Duration;
 use common_base::base::Progress;
 use common_base::base::ProgressValues;
 use common_catalog::table_context::TableContext;
@@ -57,13 +59,15 @@ pub struct HiveTableSource {
     scan_progress: Arc<Progress>,
     block_reader: Arc<HiveParquetBlockReader>,
     output: Arc<OutputPort>,
+    delay: usize,
 }
 
 impl HiveTableSource {
     pub fn create(
         ctx: Arc<dyn TableContext>,
         output: Arc<OutputPort>,
         block_reader: Arc<HiveParquetBlockReader>,
+        delay: usize,
     ) -> Result<ProcessorPtr> {
         let scan_progress = ctx.get_scan_progress();
         let mut partitions = ctx.try_get_partitions(1)?;
@@ -74,13 +78,15 @@ impl HiveTableSource {
                 block_reader,
                 scan_progress,
                 state: State::Finish,
+                delay,
             }))),
             false => Ok(ProcessorPtr::create(Box::new(HiveTableSource {
                 ctx,
                 output,
                 block_reader,
                 scan_progress,
                 state: State::ReadMeta(partitions.remove(0)),
+                delay,
             }))),
         }
     }
@@ -179,6 +185,11 @@ impl Processor for HiveTableSource {
     async fn async_process(&mut self) -> Result<()> {
         match std::mem::replace(&mut self.state, State::Finish) {
             State::ReadMeta(part) => {
+                if self.delay > 0 {
+                    sleep(Duration::from_millis(self.delay as u64)).await;
+                    tracing::debug!("sleep for {}ms", self.delay);
+                    self.delay = 0;
+                }
                 let part = HivePartInfo::from_part(&part)?;
                 let file_meta = self
                     .block_reader