diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java index 62c8085aa0f..97da8d42302 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java @@ -15,7 +15,6 @@ import java.util.List; import java.util.Map; import java.util.Set; -import org.apache.spark.sql.connector.catalog.CatalogV2Util; import org.apache.spark.sql.connector.catalog.Column; import org.apache.spark.sql.connector.expressions.NamedReference; import org.apache.spark.sql.connector.expressions.Transform; @@ -131,15 +130,20 @@ public Scan build() { /** * Prunes the columns to only include those specified in the required schema. * - *

This method clears the current column list and replaces it with columns derived from the required schema. - * Currently only supports top-level schema pruning - deeply nested schema pruning is not yet implemented. + *

Required columns keep their table type because the native scan returns whole top-level columns, even when + * Spark prunes nested fields. * * @param requiredSchema the schema specifying which columns are required */ @Override public void pruneColumns(StructType requiredSchema) { + Set required = Set.copyOf(Arrays.asList(requiredSchema.fieldNames())); readColumns.clear(); - readColumns.addAll(Arrays.asList(CatalogV2Util.structTypeToV2Columns(requiredSchema))); + for (Column column : tableColumns) { + if (required.contains(column.name())) { + readColumns.add(column); + } + } } /** diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexNestedSchemaPruningTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexNestedSchemaPruningTest.java new file mode 100644 index 00000000000..16bdc59a0f8 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexNestedSchemaPruningTest.java @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.nio.file.Path; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.spark.sql.Column; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SaveMode; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.functions; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.io.TempDir; + +/** + * Tests that reading one field of a nested column returns that field. + * + *

With {@code spark.sql.optimizer.nestedSchemaPruning.enabled}, Spark asks the scan for only the nested fields a + * query uses. Each row of the table holds {@code a = id + 100}, {@code b = id + 200} and {@code c = id + 300}, so a + * field read from the wrong position shows up as the wrong hundred. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public final class VortexNestedSchemaPruningTest { + + private static final List C_VALUES = List.of(300L, 301L, 302L, 303L, 304L); + + private SparkSession spark; + private Dataset table; + + @BeforeAll + public void setUp(@TempDir Path tempDir) { + spark = SparkSession.builder() + .appName("VortexNestedSchemaPruningTest") + .master("local[2]") + .config("spark.driver.host", "127.0.0.1") + .config("spark.sql.shuffle.partitions", "2") + .config("spark.sql.adaptive.enabled", "false") + .config("spark.sql.optimizer.nestedSchemaPruning.enabled", "true") + .config("spark.ui.enabled", "false") + .getOrCreate(); + + Column abc = functions.struct( + functions.col("id").plus(100).alias("a"), + functions.col("id").plus(200).alias("b"), + functions.col("id").plus(300).alias("c")); + Dataset df = spark.range(5) + .select( + functions.col("id"), + abc.alias("s"), + functions.array(abc).alias("arr"), + functions.map(functions.lit("k"), abc).alias("m"), + functions + .struct(functions.lit("x").alias("p"), abc.alias("inner")) + .alias("outer")); + + String path = tempDir.resolve("nested").toUri().toString(); + df.write() + .format("vortex") + .option("path", path) + .mode(SaveMode.Overwrite) + .save(); + table = spark.read().format("vortex").option("path", path).load(); + } + + @AfterAll + public void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + @Test + @DisplayName("Selecting a struct field other than the first returns that field") + public void structField() { + assertEquals(C_VALUES, longsOf(functions.col("s.c"))); + } + + @Test + @DisplayName("Selecting a field of the structs in an array returns that field") + public void arrayElementField() { + assertEquals(C_VALUES, longsOf(functions.element_at(functions.col("arr.c"), 1))); + } + + @Test + @DisplayName("Selecting a field of the structs in a map returns that field") + public void mapValueField() { + assertEquals(C_VALUES, longsOf(functions.col("m").getItem("k").getField("c"))); + } + + @Test + @DisplayName("Selecting a field of a struct inside a struct returns that field") + public void structInStructField() { + assertEquals(C_VALUES, longsOf(functions.col("outer.inner.c"))); + } + + private List longsOf(Column column) { + return table.orderBy("id").select(column).collectAsList().stream() + .map(row -> row.getLong(0)) + .collect(Collectors.toList()); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java index 85d8046833f..485c093daff 100644 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java @@ -16,6 +16,7 @@ import org.apache.spark.sql.connector.expressions.Transform; import org.apache.spark.sql.connector.expressions.filter.Predicate; import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; import org.apache.spark.unsafe.types.UTF8String; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; @@ -160,6 +161,31 @@ void pushedPredicatesReturnsCopy() { assertArrayEquals(new Predicate[] {predicate}, builder.pushedPredicates()); } + // --- pruneColumns --- + + @Test + @DisplayName("Required columns keep their full table type, so readSchema matches what the reader returns") + void pruneColumnsKeepsTableTypes() { + StructType abc = new StructType() + .add("a", DataTypes.LongType) + .add("b", DataTypes.LongType) + .add("c", DataTypes.LongType); + VortexScanBuilder builder = builderWithColumns( + new Transform[0], + Column.create("id", DataTypes.IntegerType), + Column.create("s", abc), + Column.create("name", DataTypes.StringType)); + + // With nested schema pruning, Spark asks for s with only the field the query uses. + builder.pruneColumns(new StructType() + .add("s", new StructType().add("c", DataTypes.LongType)) + .add("id", DataTypes.IntegerType)); + + assertEquals( + new StructType().add("id", DataTypes.IntegerType).add("s", abc), + builder.build().readSchema()); + } + // --- build --- @Test