fix: check for table scan before expanding

src/query/src/dist_plan/analyzer.rs

@@ -14,7 +14,6 @@
 
 use std::sync::Arc;
 
-use common_telemetry::info;
 use datafusion::datasource::DefaultTableSource;
 use datafusion::error::Result as DfResult;
 use datafusion_common::config::ConfigOptions;
@@ -44,13 +43,51 @@ impl AnalyzerRule for DistPlannerAnalyzer {
         plan: LogicalPlan,
         _config: &ConfigOptions,
     ) -> datafusion_common::Result<LogicalPlan> {
+        if !Self::is_query(&plan) {
+            return Ok(plan);
+        }
+
         let plan = plan.transform(&Self::inspect_plan_with_subquery)?;
         let mut rewriter = PlanRewriter::default();
-        plan.rewrite(&mut rewriter)
+        let result = plan.rewrite(&mut rewriter)?;
+
+        Ok(result)
     }
 }
 
 impl DistPlannerAnalyzer {
+    fn is_query(plan: &LogicalPlan) -> bool {
+        match plan {
+            LogicalPlan::Projection(_)
+            | LogicalPlan::Filter(_)
+            | LogicalPlan::Window(_)
+            | LogicalPlan::Aggregate(_)
+            | LogicalPlan::Sort(_)
+            | LogicalPlan::Join(_)
+            | LogicalPlan::CrossJoin(_)
+            | LogicalPlan::Repartition(_)
+            | LogicalPlan::Union(_)
+            | LogicalPlan::TableScan(_)
+            | LogicalPlan::Subquery(_)
+            | LogicalPlan::SubqueryAlias(_)
+            | LogicalPlan::Limit(_)
+            | LogicalPlan::Explain(_)
+            | LogicalPlan::Analyze(_)
+            | LogicalPlan::Extension(_)
+            | LogicalPlan::Distinct(_)
+            | LogicalPlan::Unnest(_) => true,
+
+            // empty relation and plain values are also counted as non-query here
+            LogicalPlan::EmptyRelation(_)
+            | LogicalPlan::Values(_)
+            | LogicalPlan::Statement(_)
+            | LogicalPlan::Prepare(_)
+            | LogicalPlan::Dml(_)
+            | LogicalPlan::Ddl(_)
+            | LogicalPlan::DescribeTable(_) => false,
+        }
+    }
+
     fn inspect_plan_with_subquery(plan: LogicalPlan) -> DfResult<Transformed<LogicalPlan>> {
         let exprs = plan
             .expressions()
@@ -138,10 +175,6 @@ impl PlanRewriter {
     /// Return true if should stop and expand. The input plan is the parent node of current node
     fn should_expand(&mut self, plan: &LogicalPlan) -> bool {
         if DFLogicalSubstraitConvertor.encode(plan).is_err() {
-            info!(
-                "substrait error: {:?}",
-                DFLogicalSubstraitConvertor.encode(plan)
-            );
             return true;
         }
 
@@ -251,6 +284,13 @@ impl TreeNodeRewriter for PlanRewriter {
             return Ok(node);
         }
 
+        // only expand when the leaf is table scan
+        if node.inputs().is_empty() && !matches!(node, LogicalPlan::TableScan(_)) {
+            self.set_expanded();
+            self.pop_stack();
+            return Ok(node);
+        }
+
         self.maybe_set_partitions(&node);
 
         let Some(parent) = self.get_parent() else {

tests-integration/tests/sql.rs

@@ -135,12 +135,12 @@ pub async fn test_mysql_crud(store_type: StorageType) {
         .await
         .unwrap();
 
-    assert!(sqlx::query(
+    sqlx::query(
         "create table demo(i bigint, ts timestamp time index, d date, dt datetime, b blob)",
     )
     .execute(&pool)
     .await
-    .is_ok());
+    .unwrap();
     for i in 0..10 {
         let dt: DateTime<Utc> = DateTime::from_naive_utc_and_offset(
             NaiveDateTime::from_timestamp_opt(60, i).unwrap(),
@@ -149,15 +149,15 @@ pub async fn test_mysql_crud(store_type: StorageType) {
         let d = NaiveDate::from_yo_opt(2015, 100).unwrap();
         let hello = format!("hello{i}");
         let bytes = hello.as_bytes();
-        assert!(sqlx::query("insert into demo values(?, ?, ?, ?, ?)")
+        sqlx::query("insert into demo values(?, ?, ?, ?, ?)")
             .bind(i)
             .bind(i)
             .bind(d)
             .bind(dt)
             .bind(bytes)
             .execute(&pool)
             .await
-            .is_ok());
+            .unwrap();
     }
 
     let rows = sqlx::query("select i, d, dt, b from demo")

tests/cases/distributed/explain/subqueries.result

@@ -113,12 +113,11 @@ EXPLAIN INSERT INTO other SELECT i, 2 FROM integers WHERE i=(SELECT MAX(i) FROM
 | logical_plan | Dml: op=[Insert] table=[other]                                    |
 |              |   Projection: integers.i AS i, TimestampMillisecond(2, None) AS j |
 |              |     Inner Join: integers.i = __scalar_sq_1.MAX(integers.i)        |
-|              |       Projection: integers.i                                      |
-|              |         MergeScan [is_placeholder=false]                          |
+|              |       TableScan: integers projection=[i]                          |
 |              |       SubqueryAlias: __scalar_sq_1                                |
-|              |         Aggregate: groupBy=[[]], aggr=[[MAX(integers.i)]]         |
-|              |           Projection: integers.i                                  |
-|              |             MergeScan [is_placeholder=false]                      |
+|              |         Projection: MAX(integers.i)                               |
+|              |           Aggregate: groupBy=[[]], aggr=[[MAX(integers.i)]]       |
+|              |             TableScan: integers projection=[i]                    |
 +--------------+-------------------------------------------------------------------+
 
 drop table other;