comphead commented on code in PR #5889:
URL: https://github.com/apache/datafusion-comet/pull/5889#discussion_r4265558667
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -2566,6 +2566,325 @@ class CometExecSuite extends CometTestBase {
}
}
+ test("scalar subqueries merged into a struct") {
+ val numRows = 1024
+ withTempPath { dir =>
+ // Keep the consuming input in one file: local[5] writes the old
five-row input as five
+ // one-row files, hiding a scalar field incorrectly returned as a
one-element array.
+ (0 until numRows)
+ .map(i => (i, i + 10))
+ .toDF("_1", "_2")
+ .coalesce(1)
+ .write
+ .parquet(dir.getCanonicalPath)
+ Seq(false, true).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
+ SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true",
+ CometConf.COMET_BATCH_SIZE.key -> numRows.toString,
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
+ withParquetTable(dir.getCanonicalPath, "tbl") {
+ assert(spark.table("tbl").inputFiles.length == 1)
+ // MergeScalarSubqueries combines the aggregate results into one
struct and reads its
+ // fields at each original scalar-subquery site. There is no
explicit struct in the SQL.
+ val df = sql("""
+ |SELECT _1,
+ | (SELECT max(_1) AS maximum FROM tbl) AS maximum,
+ | (SELECT sum(_2) AS total FROM tbl) AS total,
+ | (SELECT avg(_2) AS mean FROM tbl) AS mean
+ |FROM tbl
+ |""".stripMargin)
+ val mergedSubqueries = df.queryExecution.optimizedPlan.collect {
case p =>
+ p.expressions.flatMap(_.collect {
+ case GetStructField(s: LogicalScalarSubquery, _, _)
+ if s.dataType.isInstanceOf[StructType] =>
+ s
+ })
+ }.flatten
+ assert(
+ mergedSubqueries.nonEmpty,
+ s"Expected merged struct scalar
subqueries:\n${df.queryExecution.optimizedPlan}")
+
assert(mergedSubqueries.exists(_.dataType.asInstanceOf[StructType].length == 3))
+
+ val (_, cometPlan) =
+ checkSparkAnswerAndOperator(df, Seq(classOf[CometProjectExec]))
+ val nativeProjections = stripAQEPlan(cometPlan)
+ .collect { case p: CometProjectExec =>
+ val fields = p.projectList.flatMap(_.collect {
+ case g @ GetStructField(s: ScalarSubquery, _, _)
+ if s.dataType.isInstanceOf[StructType] =>
+ g
+ })
+ (p, fields)
+ }
+ .filter(_._2.nonEmpty)
+ assert(
+ nativeProjections.nonEmpty,
+ s"Expected CometProjectExec to consume a struct scalar
subquery:\n$cometPlan")
+ nativeProjections.foreach { case (projection, fields) =>
+ assert(fields.map(_.ordinal).toSet == Set(0, 1, 2))
+
assert(fields.forall(_.child.dataType.asInstanceOf[StructType].length == 3))
+ // Projection preserves its input batch length. Read sizes
inside each task while
+ // the iterator owns the batches, and bring only integers back
to the driver.
+ val batchSizes = projection
+ .executeColumnar()
+ .mapPartitions(batches => batches.map(_.numRows()))
+ .collect()
+ assert(batchSizes.sum == numRows)
+ assert(
+ batchSizes.exists(_ > 1),
+ s"Expected a multi-row consuming batch, got
${batchSizes.mkString(", ")}")
+ }
+ }
+ }
+ }
+ }
+ }
+
+ Seq("hour", "minute", "second", "unix_timestamp", "lpad", "rpad").foreach {
function =>
+ test(s"native functions consume scalar subqueries: $function") {
Review Comment:
Most of this could live in the SQL file tests. That framework already
disables `ConstantFolding`, and the literal cases this PR re-enables in
`hour.sql`, `minute.sql`, `second.sql` and `unix_timestamp.sql` already pin the
scalar kernel input. `scalar_subquery_struct.sql` also runs under both session
zones already. Could the plain and merged subquery shapes move there as `query
expect_native(hour)` queries? For example, `SELECT id, hour((SELECT
max(payload.instant) FROM test_struct_subq)), hour((SELECT min(payload.instant)
FROM test_struct_subq)) FROM test_struct_subq`, which I expect
`MergeScalarSubqueries` to merge (not run). The multi-row batch check is
already covered by `scalar subqueries merged into a struct` above, so this test
mostly repeats that setup for six functions.
##########
native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs:
##########
@@ -48,6 +48,84 @@ fn spark_read_side_padding2(
is_left_pad: bool,
) -> Result<ColumnarValue, DataFusionError> {
match args {
+ [ColumnarValue::Scalar(string @ (ScalarValue::Utf8(_) |
ScalarValue::LargeUtf8(_))), rest @ ..]
Review Comment:
Comet already keeps a scalar `str` away from this kernel. `CometStringLPad`
and `CometStringRPad` return `Unsupported(literalStrReason)` ("Scalar values
are not supported for the `str` argument") when `expr.str` is a `Literal`, and
the dispatcher takes over (#2630). A plain or merged subquery has the same
scalar shape. andygrove's concern was that these queries work today through
fallback and fail with this PR, and falling back also answers that. Could the
check be widened to `expr.str.references.isEmpty` instead? The dispatcher
rejects `ExecSubqueryExpression`, so these queries would fall back to Spark the
way they do in 1.1. Then this arm, `spark_pad_scalar_string`, most of the new
unit tests and `tests/padding_scalar_alloc.rs` could go. The two null-length
arms below would still be needed, because a runtime NULL `len` from a subquery
still reaches the column path. If you would rather keep this native, the
one-row wrap you used for `hour` and `unix_timestamp` covers the scalar-l
ength case in a few lines. Only the per-row-length case would then need the
serde check.
##########
spark/src/test/resources/sql-tests/expressions/misc/scalar_subquery_struct.sql:
##########
@@ -0,0 +1,150 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you 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.
+
+-- ConfigMatrix: spark.sql.session.timeZone=UTC,America/Los_Angeles
+
+statement
+CREATE TABLE test_struct_subq(
+ id int,
+ payload
struct<flag:boolean,tiny:tinyint,small:smallint,number:int,large:bigint,
+
single:float,dbl:double,amount:decimal(38,10),compact:decimal(10,2),
+
text:string,bytes:binary,day:date,instant:timestamp,local_time:timestamp_ntz,
+ nested:struct<last:string,first:bigint>>) USING parquet
+
+statement
+INSERT INTO test_struct_subq VALUES
+ (1, named_struct(
+ 'flag', true, 'tiny', -128, 'small', -32768, 'number', -2147483648,
+ 'large', -9223372036854775808, 'single', 1.25, 'dbl', -2.5,
+ 'amount', 1234567890123456789012345678.1234567890, 'compact', -12345.67,
+ 'text', '中文-é', 'bytes', X'00FF41', 'day', DATE '1969-12-31',
+ 'instant', TIMESTAMP '1969-12-31 23:59:59.123456',
+ 'local_time', TIMESTAMP_NTZ '2024-02-29 12:34:56.654321',
+ 'nested', named_struct('last', 'tail', 'first', 9876543210))),
+ (2, named_struct(
+ 'flag', NULL, 'tiny', NULL, 'small', NULL, 'number', NULL, 'large', NULL,
+ 'single', NULL, 'dbl', NULL, 'amount', NULL, 'compact', NULL, 'text', NULL,
+ 'bytes', NULL, 'day', NULL, 'instant', NULL, 'local_time', NULL, 'nested',
NULL)),
+ (3, NULL),
+ (4, named_struct(
+ 'flag', false, 'tiny', 127, 'small', 32767, 'number', 2147483647,
+ 'large', 9223372036854775807, 'single', 0.0, 'dbl', 3.5,
Review Comment:
None of the `single` and `dbl` fields hold `-0.0`, `NaN` or `Infinity`.
Could you add a row with them and a query that consumes them, such as `s.dbl =
0.0`, `isnan(s.dbl)` and `CAST(s.single AS STRING)`? Write them as `CAST('-0.0'
AS DOUBLE)` and so on, because a bare `-0.0` literal is a decimal and loses the
sign. These values matched in the earlier manual runs on this thread, but
nothing in the PR pins them.
--
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]
