comphead commented on code in PR #6725: URL: https://github.com/apache/datafusion-comet/pull/6725#discussion_r4265104608
########## spark/src/test/resources/sql-tests/iceberg/nested_schema_pruning.sql: ########## @@ -0,0 +1,135 @@ +-- 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. + +-- The native Iceberg scan reads only the nested fields that Spark's nested schema pruning keeps. +-- These queries check the pruned reads against Spark. CometIcebergNativeSuite checks that they +-- scan fewer bytes, and covers deletes, time travel, and files without field ids. + +-- No Iceberg spark-runtime is published for Spark 4.2 yet, and the 4.0 runtime the build reuses +-- is binary-incompatible with it. See https://github.com/apache/datafusion-comet/issues/4969. +-- MaxSparkVersion: 4.1 + +-- Config: spark.sql.catalog.test_cat=org.apache.iceberg.spark.SparkCatalog +-- Config: spark.sql.catalog.test_cat.type=hadoop +-- Config: spark.sql.catalog.test_cat.warehouse=/tmp/comet-iceberg-sql-test +-- Config: spark.comet.scan.icebergNative.enabled=true + +statement +DROP TABLE IF EXISTS test_cat.db.nested_pruning + +statement +CREATE TABLE test_cat.db.nested_pruning ( + id INT, + s STRUCT<a: INT, pad: STRING, inner: STRUCT<b: INT, pad: STRING>>, + items ARRAY<STRUCT<x: INT, pad: STRING>>, + m MAP<STRING, STRUCT<v: INT, pad: STRING>> +) USING iceberg + +-- NULL structs, lists, maps, list elements, map values, and fields, so that the pruned read has +-- to rebuild the validity of every level from the leaves it keeps. +statement +INSERT INTO test_cat.db.nested_pruning +SELECT + CAST(id AS INT), + IF(id % 7 = 0, NULL, named_struct( + 'a', IF(id % 5 = 0, NULL, CAST(id AS INT)), + 'pad', repeat('p', 20), + 'inner', IF(id % 3 = 0, NULL, named_struct('b', CAST(id * 2 AS INT), 'pad', repeat('q', 20))))), + IF(id % 7 = 0, NULL, array( + named_struct('x', CAST(id AS INT), 'pad', repeat('r', 20)), + IF(id % 5 = 0, NULL, named_struct('x', CAST(-id AS INT), 'pad', repeat('s', 20))))), + IF(id % 7 = 0, NULL, map('k', IF(id % 3 = 0, NULL, + named_struct('v', CAST(id AS INT), 'pad', repeat('t', 20))))) Review Comment: Added in 5395e50ba3 as you suggested: every 11th row now has an empty list and an empty map. ########## spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala: ########## @@ -2284,6 +2324,373 @@ class CometIcebergNativeSuite } } + // Spark's nested schema pruning reaches the native scan through the scan schema, so only the + // nested fields a query uses are read and the wide `pad` strings beside them are skipped. + // sql-tests/iceberg/nested_schema_pruning.sql checks the results. One data file holds every + // row, so each `pad` column chunk is larger than iceberg-rust's 1 MiB read coalescing, which + // would otherwise merge the reads of the kept chunks across the skipped ones. + test("nested schema pruning reads only the nested fields the query uses") { + assume(icebergAvailable, "Iceberg not available in classpath") + + withTempIcebergDir { warehouseDir => + withSQLConf( + "spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.test_cat.type" -> "hadoop", + "spark.sql.catalog.test_cat.warehouse" -> warehouseDir.getAbsolutePath, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + + val table = "test_cat.db.nested_pruning" + spark.sql(s""" + CREATE TABLE $table ( + id INT, + s STRUCT<a: INT, pad: STRING>, + items ARRAY<STRUCT<x: INT, pad: STRING>>, + m MAP<STRING, STRUCT<v: INT, pad: STRING>> + ) USING iceberg + """) + spark.sql(s""" + INSERT INTO $table + SELECT + CAST(id AS INT), + named_struct('a', CAST(id AS INT), 'pad', pad), + array(named_struct('x', CAST(id AS INT), 'pad', pad)), + map('k', named_struct('v', CAST(id AS INT), 'pad', pad)) + FROM ( + SELECT id, concat_ws('', transform(array('a', 'b', 'c', 'd'), + salt -> sha2(concat(CAST(id AS STRING), salt), 256))) AS pad + FROM range(0, 20000, 1, 1)) + """) + + def bytesScanned(pruneNestedFields: Boolean): Long = { + var bytes = 0L + withSQLConf( + CometConf.COMET_ICEBERG_NESTED_SCHEMA_PRUNING_ENABLED.key -> + pruneNestedFields.toString) { + val df = spark.sql(s"SELECT sum(s.a), count(items.x), count(m['k'].v) FROM $table") + df.collect() + val scans = collectIcebergNativeScans(df.queryExecution.executedPlan) + assert(scans.length == 1, s"expected one native scan, got ${scans.length}") + bytes = scans.head.metrics("bytes_scanned").value + } + bytes + } + val prunedBytes = bytesScanned(pruneNestedFields = true) + val fullBytes = bytesScanned(pruneNestedFields = false) + assert( + prunedBytes * 4 < fullBytes, + s"pruned read should skip the pad fields: pruned=$prunedBytes, full=$fullBytes") + + spark.sql(s"DROP TABLE $table") + } + } + } + + // A pruned task schema still needs the columns iceberg-rust uses beyond the projection: the + // partition source and the equality-delete key when the query projects neither. Tasks with + // deletes read with the pruned schema too, and every task of a scan reads with the same one. + // VERSION AS OF reads of columns dropped, renamed, or swapped after the snapshot use the + // snapshot's names. + test("nested schema pruning with deletes, partitions, and time travel") { + assume(icebergAvailable, "Iceberg not available in classpath") + + withTempIcebergDir { warehouseDir => + withSQLConf( + "spark.sql.catalog.test_cat" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.test_cat.type" -> "hadoop", + "spark.sql.catalog.test_cat.warehouse" -> warehouseDir.getAbsolutePath, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") { + + // Also checks that every task, tasks with deletes included, reads with one schema that + // leaves out the pruned `pad` fields. + def checkPrunedNativeScan(query: String): Unit = { Review Comment: Done in 5395e50ba3. `checkPrunedNativeScan` and the equality-delete naming test now also run each query with pruning off, so the dropped, renamed, and swapped columns are read through the full schema as well. -- 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]
