ganeshashree commented on code in PR #57630:
URL: https://github.com/apache/spark/pull/57630#discussion_r3687854202
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/GeneratorExpressionSuite.scala:
##########
@@ -114,4 +115,106 @@ class GeneratorExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
)
)
}
+
+ test("unnest - eval is lazy and only reads elements as rows are pulled") {
+ // Backing array that records which ordinals were read, to prove that
pulling the first N rows
+ // touches only the first N elements rather than materializing the whole
expansion up front.
+ val readOrdinals = scala.collection.mutable.ArrayBuffer.empty[Int]
+ val tracking = new GenericArrayData(Array[Any](10, 20, 30, 40, 50)) {
+ override def get(ordinal: Int, elementType: DataType): AnyRef = {
+ readOrdinals += ordinal
+ super.get(ordinal, elementType)
+ }
+ }
+ val result = Unnest(Seq(Literal(tracking, ArrayType(IntegerType))),
withOrdinality = true)
+ .eval(null)
+ // The returned value is a lazy Iterator, not an eagerly materialized
collection, and building
+ // it must not read any element.
+ assert(result.isInstanceOf[Iterator[_]])
+ // Literal's constructor validates its value by reading element 0; ignore
reads made before the
+ // iterator is created and observe only what pulling rows drives.
+ readOrdinals.clear()
+
+ val it = result.iterator
+ assert(readOrdinals.isEmpty, "no element should be read before the
iterator is advanced")
+ assert(it.next() === create_row(10, 1L))
+ assert(it.next() === create_row(20, 2L))
+ // Only the two consumed rows' elements were read; rows 2..4 remain
untouched.
+ assert(readOrdinals.toSeq === Seq(0, 1))
+ }
+
+ test("unnest - single array") {
+ checkTuple(Unnest(Seq(empty_array), withOrdinality = false), Seq.empty)
+ checkTuple(
+ Unnest(Seq(int_array), withOrdinality = false),
+ Seq(create_row(1), create_row(2), create_row(3)))
+ // A null array is treated as empty and contributes no rows.
+ checkTuple(
+ Unnest(Seq(Literal.create(null, ArrayType(IntegerType))), withOrdinality
= false),
+ Seq.empty)
+ }
+
+ test("unnest - single column naming and ordinality") {
+ // With a single array the output column keeps the default name `col`.
+ assert(Unnest(Seq(int_array), withOrdinality = false).elementSchema ===
+ new StructType().add("col", IntegerType, nullable = false))
+ // WITH ORDINALITY appends a 1-based, non-nullable bigint column.
+ assert(Unnest(Seq(int_array), withOrdinality = true).elementSchema ===
+ new StructType()
+ .add("col", IntegerType, nullable = false)
+ .add("ordinality", LongType, nullable = false))
+ checkTuple(
+ Unnest(Seq(str_array), withOrdinality = true),
+ Seq(create_row("a", 1L), create_row("b", 2L), create_row("c", 3L)))
+ }
+
+ test("unnest - multiple arrays are zipped and padded with nulls") {
+ val short_array = CreateArray(Seq(10, 20).map(Literal(_)))
+ // With several arrays the columns are named positionally and padded
columns are nullable.
+ assert(Unnest(Seq(int_array, short_array), withOrdinality =
false).elementSchema ===
+ new StructType()
+ .add("col0", IntegerType, nullable = true)
+ .add("col1", IntegerType, nullable = true))
+ checkTuple(
+ Unnest(Seq(int_array, short_array), withOrdinality = false),
+ Seq(create_row(1, 10), create_row(2, 20), create_row(3, null)))
+ // WITH ORDINALITY spans the full (longest) length.
+ checkTuple(
+ Unnest(Seq(int_array, short_array), withOrdinality = true),
+ Seq(create_row(1, 10, 1L), create_row(2, 20, 2L), create_row(3, null,
3L)))
+ }
+
+ test("unnest - type checks") {
+ assert(Unnest(Seq(int_array), withOrdinality =
false).checkInputDataTypes().isSuccess)
+
+ // No arguments is rejected.
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala:
##########
@@ -701,6 +701,126 @@ object InlineOuterGeneratorBuilder extends
InlineGeneratorBuilderBase {
override def isOuter: Boolean = true
}
+/**
+ * Expands one or more arrays into a table, one row per element, implementing
the ANSI SQL
+ * `UNNEST` collection derived table used in the FROM clause.
+ *
+ * When several arrays are supplied they are expanded in parallel: the number
of output rows is
+ * the length of the longest array, and shorter arrays are padded with NULLs.
A NULL array is
+ * treated as an empty array (contributes no elements). This matches the
multi-array semantics of
+ * PostgreSQL and Trino.
+ *
+ * Each array contributes exactly one output column holding its element as-is;
unlike `inline`,
+ * arrays of structs are not expanded into one column per field. When
`withOrdinality` is set, a
+ * trailing 1-based `BIGINT` ordinality column is appended, matching `WITH
ORDINALITY` in
+ * PostgreSQL and Trino (BigQuery's 0-based `WITH OFFSET` is intentionally not
adopted).
+ *
+ * {{{
+ * SELECT * FROM UNNEST(array(10, 20), array(30)) WITH ORDINALITY ->
+ * 10 30 1
+ * 20 NULL 2
+ * }}}
+ *
+ * This generator uses interpreted evaluation ([[CodegenFallback]]);
[[GenerateExec]] therefore
+ * disables whole-stage codegen for the enclosing `Generate`. The
per-array/per-ordinality zip with
+ * NULL padding does not map onto the existing [[CollectionGenerator]] codegen
path (which emits a
+ * single `ArrayData`/`MapData`), and its interpreted `eval` is already lazy
(one row built per
+ * pull). Interpreted generation is the same choice made by other
non-`CollectionGenerator`
+ * generators such as [[ReplicateRows]]. A dedicated codegen path (analogous
to `arrays_zip`) is
+ * possible future work if UNNEST becomes hot in whole-stage-codegen pipelines.
+ */
+case class Unnest(children: Seq[Expression], withOrdinality: Boolean)
+ extends Generator with CodegenFallback {
+
+ private lazy val arrayElementTypes: Seq[ArrayType] =
+ children.map(_.dataType.asInstanceOf[ArrayType])
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ if (children.isEmpty) {
+ throw QueryCompilationErrors.wrongNumArgsError(
+ toSQLId(prettyName), Seq("> 0"), children.length)
+ }
+ val nonArray = children.zipWithIndex.collectFirst {
+ case (e, idx) if !e.dataType.isInstanceOf[ArrayType] => (e, idx)
+ }
+ nonArray match {
+ case Some((e, idx)) =>
+ DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_INPUT_TYPE",
+ messageParameters = Map(
+ "paramIndex" -> ordinalNumber(idx),
+ "requiredType" -> toSQLType(ArrayType),
+ "inputSql" -> toSQLExpr(e),
+ "inputType" -> toSQLType(e.dataType)))
+ case None =>
+ TypeCheckResult.TypeCheckSuccess
+ }
+ }
+
+ override def elementSchema: StructType = {
+ // With a single array keep the `explode`-compatible default name `col`;
with several arrays
+ // use positional names `col0`, `col1`, .... A shorter array yields NULLs
in the padded rows,
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala:
##########
@@ -701,6 +701,126 @@ object InlineOuterGeneratorBuilder extends
InlineGeneratorBuilderBase {
override def isOuter: Boolean = true
}
+/**
+ * Expands one or more arrays into a table, one row per element, implementing
the ANSI SQL
+ * `UNNEST` collection derived table used in the FROM clause.
+ *
+ * When several arrays are supplied they are expanded in parallel: the number
of output rows is
+ * the length of the longest array, and shorter arrays are padded with NULLs.
A NULL array is
+ * treated as an empty array (contributes no elements). This matches the
multi-array semantics of
+ * PostgreSQL and Trino.
+ *
+ * Each array contributes exactly one output column holding its element as-is;
unlike `inline`,
+ * arrays of structs are not expanded into one column per field. When
`withOrdinality` is set, a
+ * trailing 1-based `BIGINT` ordinality column is appended, matching `WITH
ORDINALITY` in
+ * PostgreSQL and Trino (BigQuery's 0-based `WITH OFFSET` is intentionally not
adopted).
+ *
+ * {{{
+ * SELECT * FROM UNNEST(array(10, 20), array(30)) WITH ORDINALITY ->
+ * 10 30 1
+ * 20 NULL 2
+ * }}}
+ *
+ * This generator uses interpreted evaluation ([[CodegenFallback]]);
[[GenerateExec]] therefore
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
