[SQL] Preserve SELECT-list order when extracting multiple generators

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -3284,8 +3284,8 @@ class Analyzer(
    *
    * This rule will throw [[AnalysisException]] for following cases:
    * 1. [[Generator]] is nested in expressions, e.g. `SELECT explode(list) + 1 FROM tbl`
-   * 2. more than one [[Generator]] is found in projectList,
-   *    e.g. `SELECT explode(list), explode(list) FROM tbl`
+   * 2. more than one [[Generator]] is found in an aggregate result list,
+   *    e.g. `SELECT explode(list), explode(list), count(*) FROM tbl`
    * 3. [[Generator]] is found in other operators that are not [[Project]] or [[Generate]],
    *    e.g. `SELECT * FROM tbl SORT BY explode(list)`
    */
@@ -3346,6 +3346,62 @@ class Analyzer(
       }
     }
 
+    private def isGeneratorOrUnresolvedGenerator(expr: Expression): Boolean = {
+      expr.containsPattern(GENERATOR) || (
+        expr.containsPattern(UNRESOLVED_FUNCTION) &&
+        expr.exists {
+          case u: UnresolvedFunction =>
+            u.nameParts.lastOption.exists(_.equalsIgnoreCase("json_tuple")) ||
+            functionResolution
+              .lookupBuiltinOrTempFunction(u.nameParts, Some(u))
+              .exists(_.getGroup == "generator_funcs")
+          case _ => false
+        }
+      )
+    }
+
+    /** Allows bypassing earlier generators only when they depend on a generator to their right. */
+    private def canExtractGeneratorInProjectListOrder(
+        precedingExpressions: Seq[NamedExpression],
+        remainingExpressions: Seq[NamedExpression]): Boolean = {
+      def generatorOutputNames(expression: NamedExpression): Seq[String] = expression match {
+        case AliasedGenerator(_, names, _) =>
+          names
+        case Alias(child, name) if isGeneratorOrUnresolvedGenerator(child) =>
+          name :: Nil
+        case MultiAlias(child, names) if isGeneratorOrUnresolvedGenerator(child) =>
+          names
+        case _ =>
+          Nil
+      }
+
+      val neededGeneratorOutputNames = ArrayBuffer.empty[String]
+      remainingExpressions.foreach { expression =>
+        generatorOutputNames(expression).foreach { name =>
+          neededGeneratorOutputNames += name
+        }
+      }
+
+      precedingExpressions.reverseIterator.forall { expression =>
+        if (!isGeneratorOrUnresolvedGenerator(expression)) {
+          true
+        } else {
+          val dependsOnGeneratorToRight = expression.exists {
+            case unresolvedAttribute: UnresolvedAttribute =>
+              neededGeneratorOutputNames.exists(
+                conf.resolver(_, unresolvedAttribute.nameParts.head))
+            case _ => false
+          }
+          if (dependsOnGeneratorToRight) {
+            generatorOutputNames(expression).foreach { name =>
+              neededGeneratorOutputNames += name
+            }
+          }
+          dependsOnGeneratorToRight
+        }
+      }
+    }
+
     def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsUpWithPruning(
       _.containsPattern(GENERATOR), ruleId) {
       case p @ Project(Seq(UnresolvedStarWithColumns(_, _, _)), _) =>
@@ -3415,13 +3471,16 @@ class Analyzer(
 
       // The star will be expanded differently if we insert `Generate` under `Project` too early.
       case p @ Project(projectList, child) if !projectList.exists(_.exists(_.isInstanceOf[Star])) =>
-        val (resolvedGenerator, newProjectList) = projectList
-          .map(trimNonTopLevelAliases)
-          .foldLeft((None: Option[Generate], Nil: Seq[NamedExpression])) { (res, e) =>
-            e match {
-              // If there are more than one generator, we only rewrite the first one and wait for
-              // the next analyzer iteration to rewrite the next one.
-              case AliasedGenerator(generator, names, outer) if res._1.isEmpty &&
+        val trimmedProjectList = projectList.map(trimNonTopLevelAliases)
+        val (resolvedGenerator, newProjectList) = trimmedProjectList.zipWithIndex
+          .foldLeft((None: Option[Generate], Nil: Seq[NamedExpression])) {
+            case (res, (AliasedGenerator(generator, names, outer), index))
+                if res._1.isEmpty &&
+                  (!conf.getConf(SQLConf.GENERATOR_PRESERVE_SELECT_LIST_ORDER) ||
+                    canExtractGeneratorInProjectListOrder(
+                      precedingExpressions = res._2,
+                      remainingExpressions = trimmedProjectList.drop(index)
+                    )) &&
                   generator.childrenResolved =>
                 val g = Generate(
                   generator,
@@ -3431,9 +3490,8 @@ class Analyzer(
                   generatorOutput = GeneratorResolution.makeGeneratorOutput(generator, names),
                   child)
                 (Some(g), res._2 ++ g.nullableOutput)
-              case other =>
-                (res._1, res._2 :+ other)
-            }
+            case (res, (other, _)) =>
+              (res._1, res._2 :+ other)
           }
 
         if (resolvedGenerator.isDefined) {

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -584,6 +584,17 @@ object SQLConf {
       .booleanConf
       .createWithDefault(true)
 
+  val GENERATOR_PRESERVE_SELECT_LIST_ORDER =
+    buildConf("spark.sql.generator.preserveSelectListOrder")
+      .internal()
+      .doc("When true, independent generator expressions in a SELECT list are extracted in " +
+        "SELECT-list order. When false, Spark extracts the first generator whose children are " +
+        "resolved, which can reorder generator evaluation.")
+      .version("4.4.0")
+      .withBindingPolicy(ConfigBindingPolicy.SESSION)
+      .booleanConf
+      .createWithDefault(true)
+
   val MULTI_COMMUTATIVE_OP_OPT_THRESHOLD =
     buildConf("spark.sql.analyzer.canonicalization.multiCommutativeOpMemoryOptThreshold")
       .internal()

sql/core/src/test/resources/sql-tests/analyzer-results/generators-resolution-edge-cases.sql.out

@@ -38,8 +38,8 @@ Project [col#x, col#x]
 SELECT explode(array(sin(0), 1, 2)), explode(array(10, 20))
 -- !query analysis
 Project [col#x, col#x]
-+- Generate explode(array(SIN(cast(0 as double)), cast(1 as double), cast(2 as double))), false, [col#x]
-   +- Generate explode(array(10, 20)), false, [col#x]
++- Generate explode(array(10, 20)), false, [col#x]
+   +- Generate explode(array(SIN(cast(0 as double)), cast(1 as double), cast(2 as double))), false, [col#x]
       +- OneRowRelation
 
 
@@ -480,6 +480,18 @@ Project [col#x, arr#x]
       +- OneRowRelation
 
 
+-- !query
+SELECT explode(b) AS a,
+       explode(c) AS b,
+       explode(array(array(array(1)), array(array(2)))) AS c
+-- !query analysis
+Project [a#x, b#x, c#x]
++- Generate explode(b#x), false, [a#x]
+   +- Generate explode(c#x), false, [b#x]
+      +- Generate explode(array(array(array(1)), array(array(2)))), false, [c#x]
+         +- OneRowRelation
+
+
 -- !query
 SELECT col + 1 as col2, explode(array(1, 2, 3)) as col
 -- !query analysis

sql/core/src/test/resources/sql-tests/analyzer-results/udf/generator-ordering-legacy.sql.out

@@ -0,0 +1,32 @@
+-- Automatically generated by SQLQueryTestSuite
+-- !query
+SELECT
+  explode(array(sin(0) + udf(monotonically_increasing_id()))) AS left_value,
+  explode(array(10, 20)) AS right_value
+-- !query analysis
+[Analyzer test output redacted due to nondeterminism]
+
+
+-- !query
+SELECT
+  explode(array(udf(sin(0)), udf(1))) AS left_value,
+  explode(array(monotonically_increasing_id())) AS right_value
+-- !query analysis
+[Analyzer test output redacted due to nondeterminism]
+
+
+-- !query
+SELECT
+  explode(array(g2)) AS g1,
+  explode(array(g3)) AS g2,
+  explode(array(udf(base + monotonically_increasing_id()) + sin(0))) AS g3,
+  explode(array(10L, 20L)) AS independent,
+  gid
+FROM (
+  SELECT coalesce(id, 0L) AS base, grouping_id(id) AS gid
+  FROM VALUES (1L) AS t(id)
+  GROUP BY GROUPING SETS ((id), ())
+) grouped
+WHERE gid = 1
+-- !query analysis
+[Analyzer test output redacted due to nondeterminism]

sql/core/src/test/resources/sql-tests/analyzer-results/udf/generator-ordering.sql.out

@@ -0,0 +1,32 @@
+-- Automatically generated by SQLQueryTestSuite
+-- !query
+SELECT
+  explode(array(sin(0) + udf(monotonically_increasing_id()))) AS left_value,
+  explode(array(10, 20)) AS right_value
+-- !query analysis
+[Analyzer test output redacted due to nondeterminism]
+
+
+-- !query
+SELECT
+  explode(array(udf(sin(0)), udf(1))) AS left_value,
+  explode(array(monotonically_increasing_id())) AS right_value
+-- !query analysis
+[Analyzer test output redacted due to nondeterminism]
+
+
+-- !query
+SELECT
+  explode(array(g2)) AS g1,
+  explode(array(g3)) AS g2,
+  explode(array(udf(base + monotonically_increasing_id()) + sin(0))) AS g3,
+  explode(array(10L, 20L)) AS independent,
+  gid
+FROM (
+  SELECT coalesce(id, 0L) AS base, grouping_id(id) AS gid
+  FROM VALUES (1L) AS t(id)
+  GROUP BY GROUPING SETS ((id), ())
+) grouped
+WHERE gid = 1
+-- !query analysis
+[Analyzer test output redacted due to nondeterminism]

sql/core/src/test/resources/sql-tests/inputs/generators-resolution-edge-cases.sql

@@ -7,7 +7,7 @@ SELECT 1 + explode(array(1, 2, 3));
 -- multiple generators should work
 SELECT explode(array(0, 1, 2)), explode(array(10, 20));
 
--- multiple generators' order is not fixed and depends on rule ordering
+-- multiple generators preserve SELECT-list order when an earlier generator resolves later
 SELECT explode(array(sin(0), 1, 2)), explode(array(10, 20));
 
 -- multiple generators in aggregate should fail
@@ -134,6 +134,11 @@ SELECT explode(array(array(0), array(1), array(2))) as arr, explode(arr) as col;
 -- generator LCA right-to-left should work
 SELECT explode(arr) as col, explode(array(array(0), array(1), array(2))) as arr;
 
+-- transitive generator LCA right-to-left should work
+SELECT explode(b) AS a,
+       explode(c) AS b,
+       explode(array(array(array(1)), array(array(2)))) AS c;
+
 -- generator output LCA right-to-left should fail (reference before definition)
 SELECT col + 1 as col2, explode(array(1, 2, 3)) as col;
 

sql/core/src/test/resources/sql-tests/inputs/udf/generator-ordering-legacy.sql

@@ -0,0 +1,25 @@
+--SET spark.sql.generator.preserveSelectListOrder=false
+
+-- A nondeterministic expression in the first generator is evaluated after later generators.
+SELECT
+  explode(array(sin(0) + udf(monotonically_increasing_id()))) AS left_value,
+  explode(array(10, 20)) AS right_value;
+
+-- A nondeterministic expression in the second generator is evaluated before the first generator.
+SELECT
+  explode(array(udf(sin(0)), udf(1))) AS left_value,
+  explode(array(monotonically_increasing_id())) AS right_value;
+
+-- UDF, nondeterministic expression, grouping analytics, and transitive rightward generator LCA.
+SELECT
+  explode(array(g2)) AS g1,
+  explode(array(g3)) AS g2,
+  explode(array(udf(base + monotonically_increasing_id()) + sin(0))) AS g3,
+  explode(array(10L, 20L)) AS independent,
+  gid
+FROM (
+  SELECT coalesce(id, 0L) AS base, grouping_id(id) AS gid
+  FROM VALUES (1L) AS t(id)
+  GROUP BY GROUPING SETS ((id), ())
+) grouped
+WHERE gid = 1;

sql/core/src/test/resources/sql-tests/inputs/udf/generator-ordering.sql

@@ -0,0 +1,25 @@
+--SET spark.sql.generator.preserveSelectListOrder=true
+
+-- A nondeterministic expression in the first generator is evaluated before later generators.
+SELECT
+  explode(array(sin(0) + udf(monotonically_increasing_id()))) AS left_value,
+  explode(array(10, 20)) AS right_value;
+
+-- A nondeterministic expression in the second generator is evaluated per first-generator row.
+SELECT
+  explode(array(udf(sin(0)), udf(1))) AS left_value,
+  explode(array(monotonically_increasing_id())) AS right_value;
+
+-- UDF, nondeterministic expression, grouping analytics, and transitive rightward generator LCA.
+SELECT
+  explode(array(g2)) AS g1,
+  explode(array(g3)) AS g2,
+  explode(array(udf(base + monotonically_increasing_id()) + sin(0))) AS g3,
+  explode(array(10L, 20L)) AS independent,
+  gid
+FROM (
+  SELECT coalesce(id, 0L) AS base, grouping_id(id) AS gid
+  FROM VALUES (1L) AS t(id)
+  GROUP BY GROUPING SETS ((id), ())
+) grouped
+WHERE gid = 1;

sql/core/src/test/resources/sql-tests/results/generators-resolution-edge-cases.sql.out

@@ -482,6 +482,17 @@ struct<col:int,arr:array<int>>
 2	[2]
 
 
+-- !query
+SELECT explode(b) AS a,
+       explode(c) AS b,
+       explode(array(array(array(1)), array(array(2)))) AS c
+-- !query schema
+struct<a:int,b:array<int>,c:array<array<int>>>
+-- !query output
+1	[1]	[[1]]
+2	[2]	[[2]]
+
+
 -- !query
 SELECT col + 1 as col2, explode(array(1, 2, 3)) as col
 -- !query schema

sql/core/src/test/resources/sql-tests/results/udf/generator-ordering-legacy.sql.out

@@ -0,0 +1,41 @@
+-- Automatically generated by SQLQueryTestSuite
+-- !query
+SELECT
+  explode(array(sin(0) + udf(monotonically_increasing_id()))) AS left_value,
+  explode(array(10, 20)) AS right_value
+-- !query schema
+struct<left_value:double,right_value:int>
+-- !query output
+0.0	10
+1.0	20
+
+
+-- !query
+SELECT
+  explode(array(udf(sin(0)), udf(1))) AS left_value,
+  explode(array(monotonically_increasing_id())) AS right_value
+-- !query schema
+struct<left_value:double,right_value:bigint>
+-- !query output
+0.0	0
+1.0	0
+
+
+-- !query
+SELECT
+  explode(array(g2)) AS g1,
+  explode(array(g3)) AS g2,
+  explode(array(udf(base + monotonically_increasing_id()) + sin(0))) AS g3,
+  explode(array(10L, 20L)) AS independent,
+  gid
+FROM (
+  SELECT coalesce(id, 0L) AS base, grouping_id(id) AS gid
+  FROM VALUES (1L) AS t(id)
+  GROUP BY GROUPING SETS ((id), ())
+) grouped
+WHERE gid = 1
+-- !query schema
+struct<g1:double,g2:double,g3:double,independent:bigint,gid:bigint>
+-- !query output
+0.0	0.0	0.0	10	1
+1.0	1.0	1.0	20	1

sql/core/src/test/resources/sql-tests/results/udf/generator-ordering.sql.out

@@ -0,0 +1,41 @@
+-- Automatically generated by SQLQueryTestSuite
+-- !query
+SELECT
+  explode(array(sin(0) + udf(monotonically_increasing_id()))) AS left_value,
+  explode(array(10, 20)) AS right_value
+-- !query schema
+struct<left_value:double,right_value:int>
+-- !query output
+0.0	10
+0.0	20
+
+
+-- !query
+SELECT
+  explode(array(udf(sin(0)), udf(1))) AS left_value,
+  explode(array(monotonically_increasing_id())) AS right_value
+-- !query schema
+struct<left_value:double,right_value:bigint>
+-- !query output
+0.0	0
+1.0	1
+
+
+-- !query
+SELECT
+  explode(array(g2)) AS g1,
+  explode(array(g3)) AS g2,
+  explode(array(udf(base + monotonically_increasing_id()) + sin(0))) AS g3,
+  explode(array(10L, 20L)) AS independent,
+  gid
+FROM (
+  SELECT coalesce(id, 0L) AS base, grouping_id(id) AS gid
+  FROM VALUES (1L) AS t(id)
+  GROUP BY GROUPING SETS ((id), ())
+) grouped
+WHERE gid = 1
+-- !query schema
+struct<g1:double,g2:double,g3:double,independent:bigint,gid:bigint>
+-- !query output
+0.0	0.0	0.0	10	1
+0.0	0.0	0.0	20	1