Skip to content

Commit ff79104

Browse files
gitRasheedgitRasheed
authored
fix(spark): return the requested field when Spark prunes nested fields (#10724)
## Summary Fixes #10723 Reading one field of a nested column returned a different field, because `readSchema()` described pruned structs while the native scan returns each top-level column whole. ## Changes `pruneColumns` now keeps each required column with its full table type, in table order, so `readSchema()` matches what the reader returns. A query that uses one field of a nested column still reads the whole column. `VortexNestedSchemaPruningTest` selects a non-first field from a struct, an array of structs, a map of structs and a struct inside a struct. All four fail without the fix. `VortexScanBuilderTest.pruneColumnsKeepsTableTypes` checks the read schema directly. I used AI tools to help write and review this change. --------- Signed-off-by: gitRasheed <99349713+gitRasheed@users.noreply.github.com>
1 parent 4c8425f commit ff79104

3 files changed

Lines changed: 144 additions & 4 deletions

File tree

‎java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import java.util.List;
1616
import java.util.Map;
1717
import java.util.Set;
18-
import org.apache.spark.sql.connector.catalog.CatalogV2Util;
1918
import org.apache.spark.sql.connector.catalog.Column;
2019
import org.apache.spark.sql.connector.expressions.NamedReference;
2120
import org.apache.spark.sql.connector.expressions.Transform;
@@ -131,15 +130,20 @@ public Scan build() {
131130
/**
132131
* Prunes the columns to only include those specified in the required schema.
133132
*
134-
* <p>This method clears the current column list and replaces it with columns derived from the required schema.
135-
* Currently only supports top-level schema pruning - deeply nested schema pruning is not yet implemented.
133+
* <p>Required columns keep their table type because the native scan returns whole top-level columns, even when
134+
* Spark prunes nested fields.
136135
*
137136
* @param requiredSchema the schema specifying which columns are required
138137
*/
139138
@Override
140139
public void pruneColumns(StructType requiredSchema) {
140+
Set<String> required = Set.copyOf(Arrays.asList(requiredSchema.fieldNames()));
141141
readColumns.clear();
142-
readColumns.addAll(Arrays.asList(CatalogV2Util.structTypeToV2Columns(requiredSchema)));
142+
for (Column column : tableColumns) {
143+
if (required.contains(column.name())) {
144+
readColumns.add(column);
145+
}
146+
}
143147
}
144148

145149
/**

‎java/vortex-spark/src/test/java/dev/vortex/spark/VortexNestedSchemaPruningTest.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
package dev.vortex.spark;
5+
6+
import static org.junit.jupiter.api.Assertions.assertEquals;
7+
8+
import java.nio.file.Path;
9+
import java.util.List;
10+
import java.util.stream.Collectors;
11+
import org.apache.spark.sql.Column;
12+
import org.apache.spark.sql.Dataset;
13+
import org.apache.spark.sql.Row;
14+
import org.apache.spark.sql.SaveMode;
15+
import org.apache.spark.sql.SparkSession;
16+
import org.apache.spark.sql.functions;
17+
import org.junit.jupiter.api.AfterAll;
18+
import org.junit.jupiter.api.BeforeAll;
19+
import org.junit.jupiter.api.DisplayName;
20+
import org.junit.jupiter.api.Test;
21+
import org.junit.jupiter.api.TestInstance;
22+
import org.junit.jupiter.api.io.TempDir;
23+
24+
/**
25+
* Tests that reading one field of a nested column returns that field.
26+
*
27+
* <p>With {@code spark.sql.optimizer.nestedSchemaPruning.enabled}, Spark asks the scan for only the nested fields a
28+
* query uses. Each row of the table holds {@code a = id + 100}, {@code b = id + 200} and {@code c = id + 300}, so a
29+
* field read from the wrong position shows up as the wrong hundred.
30+
*/
31+
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
32+
public final class VortexNestedSchemaPruningTest {
33+
34+
private static final List<Long> C_VALUES = List.of(300L, 301L, 302L, 303L, 304L);
35+
36+
private SparkSession spark;
37+
private Dataset<Row> table;
38+
39+
@BeforeAll
40+
public void setUp(@TempDir Path tempDir) {
41+
spark = SparkSession.builder()
42+
.appName("VortexNestedSchemaPruningTest")
43+
.master("local[2]")
44+
.config("spark.driver.host", "127.0.0.1")
45+
.config("spark.sql.shuffle.partitions", "2")
46+
.config("spark.sql.adaptive.enabled", "false")
47+
.config("spark.sql.optimizer.nestedSchemaPruning.enabled", "true")
48+
.config("spark.ui.enabled", "false")
49+
.getOrCreate();
50+
51+
Column abc = functions.struct(
52+
functions.col("id").plus(100).alias("a"),
53+
functions.col("id").plus(200).alias("b"),
54+
functions.col("id").plus(300).alias("c"));
55+
Dataset<Row> df = spark.range(5)
56+
.select(
57+
functions.col("id"),
58+
abc.alias("s"),
59+
functions.array(abc).alias("arr"),
60+
functions.map(functions.lit("k"), abc).alias("m"),
61+
functions
62+
.struct(functions.lit("x").alias("p"), abc.alias("inner"))
63+
.alias("outer"));
64+
65+
String path = tempDir.resolve("nested").toUri().toString();
66+
df.write()
67+
.format("vortex")
68+
.option("path", path)
69+
.mode(SaveMode.Overwrite)
70+
.save();
71+
table = spark.read().format("vortex").option("path", path).load();
72+
}
73+
74+
@AfterAll
75+
public void tearDown() {
76+
if (spark != null) {
77+
spark.stop();
78+
}
79+
}
80+
81+
@Test
82+
@DisplayName("Selecting a struct field other than the first returns that field")
83+
public void structField() {
84+
assertEquals(C_VALUES, longsOf(functions.col("s.c")));
85+
}
86+
87+
@Test
88+
@DisplayName("Selecting a field of the structs in an array returns that field")
89+
public void arrayElementField() {
90+
assertEquals(C_VALUES, longsOf(functions.element_at(functions.col("arr.c"), 1)));
91+
}
92+
93+
@Test
94+
@DisplayName("Selecting a field of the structs in a map returns that field")
95+
public void mapValueField() {
96+
assertEquals(C_VALUES, longsOf(functions.col("m").getItem("k").getField("c")));
97+
}
98+
99+
@Test
100+
@DisplayName("Selecting a field of a struct inside a struct returns that field")
101+
public void structInStructField() {
102+
assertEquals(C_VALUES, longsOf(functions.col("outer.inner.c")));
103+
}
104+
105+
private List<Long> longsOf(Column column) {
106+
return table.orderBy("id").select(column).collectAsList().stream()
107+
.map(row -> row.getLong(0))
108+
.collect(Collectors.toList());
109+
}
110+
}

‎java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.apache.spark.sql.connector.expressions.Transform;
1717
import org.apache.spark.sql.connector.expressions.filter.Predicate;
1818
import org.apache.spark.sql.types.DataTypes;
19+
import org.apache.spark.sql.types.StructType;
1920
import org.apache.spark.unsafe.types.UTF8String;
2021
import org.junit.jupiter.api.DisplayName;
2122
import org.junit.jupiter.api.Test;
@@ -160,6 +161,31 @@ void pushedPredicatesReturnsCopy() {
160161
assertArrayEquals(new Predicate[] {predicate}, builder.pushedPredicates());
161162
}
162163

164+
// --- pruneColumns ---
165+
166+
@Test
167+
@DisplayName("Required columns keep their full table type, so readSchema matches what the reader returns")
168+
void pruneColumnsKeepsTableTypes() {
169+
StructType abc = new StructType()
170+
.add("a", DataTypes.LongType)
171+
.add("b", DataTypes.LongType)
172+
.add("c", DataTypes.LongType);
173+
VortexScanBuilder builder = builderWithColumns(
174+
new Transform[0],
175+
Column.create("id", DataTypes.IntegerType),
176+
Column.create("s", abc),
177+
Column.create("name", DataTypes.StringType));
178+
179+
// With nested schema pruning, Spark asks for s with only the field the query uses.
180+
builder.pruneColumns(new StructType()
181+
.add("s", new StructType().add("c", DataTypes.LongType))
182+
.add("id", DataTypes.IntegerType));
183+
184+
assertEquals(
185+
new StructType().add("id", DataTypes.IntegerType).add("s", abc),
186+
builder.build().readSchema());
187+
}
188+
163189
// --- build ---
164190

165191
@Test

0 commit comments

Comments
 (0)