Skip to content

Commit 34c8e07

Browse files
robert3005robert3005
committed
fixes
Signed-off-by: Robert Kruszewski <github@robertk.io>
1 parent 1662bba commit 34c8e07

8 files changed

Lines changed: 167 additions & 25 deletions

File tree

‎java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,12 @@ private static Optional<Expression> comparison(
143143
return Optional.empty();
144144
}
145145
Optional<Expression> literal = literal(value, column.get().dataType());
146-
return literal.map(expression -> Expression.binary(op, column.get().expression(), expression));
146+
return literal.map(expression -> {
147+
Expression field = column.get().expression();
148+
Expression comparison = Expression.binary(op, field, expression);
149+
// Null-safe equality must return false for a null field, including beneath NOT.
150+
return nullSafe ? Expression.and(Expression.isNotNull(field), comparison) : comparison;
151+
});
147152
}
148153

149154
private static Optional<Expression> convertIn(In in, StructType dataSchema) {

‎java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java‎

Lines changed: 24 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,12 @@
2323
import java.io.UncheckedIOException;
2424
import java.util.ArrayList;
2525
import java.util.Arrays;
26-
import java.util.HashSet;
26+
import java.util.HashMap;
2727
import java.util.List;
2828
import java.util.Locale;
29+
import java.util.Map;
2930
import java.util.Optional;
30-
import java.util.Set;
31+
import java.util.TreeMap;
3132
import org.apache.spark.sql.connector.read.PartitionReader;
3233
import org.apache.spark.sql.execution.datasources.PartitionedFile;
3334
import org.apache.spark.sql.execution.vectorized.ConstantColumnVector;
@@ -69,7 +70,8 @@ public VortexPartitionReader(
6970
StructType readPartitionSchema,
7071
VortexIo io,
7172
VortexOptions formatOptions,
72-
Filter[] pushedFilters) {
73+
Filter[] pushedFilters,
74+
boolean caseSensitive) {
7375
this.file = file;
7476
this.readDataSchema = readDataSchema;
7577
this.readPartitionSchema = readPartitionSchema;
@@ -79,14 +81,15 @@ public VortexPartitionReader(
7981
readable = io.openReadable(new VortexFile(file.toPath().toString(), file.fileSize()));
8082
dataSource = DataSource.open(session, List.of(readable), io.readConcurrency());
8183

82-
Set<String> fileFields = fieldNames(dataSource);
84+
Map<String, String> fileFields = fieldNames(dataSource, caseSensitive);
8385
List<String> projection = new ArrayList<>(readDataSchema.length());
8486
this.vectorSlots = new int[readDataSchema.length()];
8587
StructField[] fields = readDataSchema.fields();
8688
for (int i = 0; i < fields.length; i++) {
87-
if (fileFields.contains(fields[i].name())) {
89+
String fileName = fileFields.get(fields[i].name());
90+
if (fileName != null) {
8891
vectorSlots[i] = projection.size();
89-
projection.add(fields[i].name());
92+
projection.add(fileName);
9093
} else if (fields[i].nullable()) {
9194
vectorSlots[i] = -1;
9295
} else {
@@ -113,17 +116,27 @@ public VortexPartitionReader(
113116
}
114117
}
115118

116-
private Set<String> fieldNames(DataSource source) {
117-
Set<String> names = new HashSet<>();
119+
private Map<String, String> fieldNames(DataSource source, boolean caseSensitive) {
120+
Map<String, String> names = caseSensitive ? new HashMap<>() : new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
118121
for (Field field : source.arrowSchema(allocator).getFields()) {
119-
names.add(field.getName());
122+
String name = field.getName();
123+
String previous = names.putIfAbsent(name, name);
124+
if (previous != null) {
125+
throw new IllegalArgumentException(String.format(
126+
Locale.ROOT,
127+
"Ambiguous columns %s and %s in %s under Spark's case sensitivity setting",
128+
previous,
129+
name,
130+
file.toPath()));
131+
}
120132
}
121133
return names;
122134
}
123135

124-
private static StructType restrictTo(StructType schema, Set<String> names) {
136+
private static StructType restrictTo(StructType schema, Map<String, String> names) {
137+
// Native filters use physical field names. Filters with different casing stay with Spark.
125138
StructField[] present = Arrays.stream(schema.fields())
126-
.filter(field -> names.contains(field.name()))
139+
.filter(field -> field.name().equals(names.get(field.name())))
127140
.toArray(StructField[]::new);
128141
return present.length == schema.length() ? schema : new StructType(present);
129142
}

‎java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ public final class VortexPartitionReaderFactory extends FilePartitionReaderFacto
2727
private final StructType readDataSchema;
2828
private final StructType readPartitionSchema;
2929
private final Filter[] pushedFilters;
30+
private final boolean caseSensitive;
3031

3132
public VortexPartitionReaderFactory(
3233
FileSourceOptions fileOptions,
@@ -35,14 +36,16 @@ public VortexPartitionReaderFactory(
3536
StructType dataSchema,
3637
StructType readDataSchema,
3738
StructType readPartitionSchema,
38-
Filter[] pushedFilters) {
39+
Filter[] pushedFilters,
40+
boolean caseSensitive) {
3941
this.fileOptions = fileOptions;
4042
this.io = io;
4143
this.formatOptions = formatOptions;
4244
this.dataSchema = dataSchema;
4345
this.readDataSchema = readDataSchema;
4446
this.readPartitionSchema = readPartitionSchema;
4547
this.pushedFilters = pushedFilters == null ? new Filter[0] : pushedFilters.clone();
48+
this.caseSensitive = caseSensitive;
4649
}
4750

4851
@Override
@@ -58,7 +61,7 @@ public PartitionReader<InternalRow> buildReader(PartitionedFile file) {
5861
@Override
5962
public PartitionReader<ColumnarBatch> buildColumnarReader(PartitionedFile file) {
6063
return new VortexPartitionReader(
61-
file, dataSchema, readDataSchema, readPartitionSchema, io, formatOptions, pushedFilters);
64+
file, dataSchema, readDataSchema, readPartitionSchema, io, formatOptions, pushedFilters, caseSensitive);
6265
}
6366

6467
@Override

‎java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexFileFormat.scala‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,11 @@ import org.apache.spark.sql.execution.datasources.{
2121
import org.apache.spark.sql.internal.SQLConf
2222
import org.apache.spark.sql.sources.{DataSourceRegister, Filter}
2323
import org.apache.spark.sql.types.{DataType, StructType}
24-
import org.apache.spark.sql.vectorized.ColumnarBatch
24+
import org.apache.spark.sql.vectorized.{ColumnVector, ColumnarBatch}
2525

2626
import dev.vortex.spark.VortexOptions
2727
import dev.vortex.spark.io.VortexIo
28-
import dev.vortex.spark.read.{
29-
VortexArrowColumnVector,
30-
VortexFooterReader,
31-
VortexPartitionReader
32-
}
28+
import dev.vortex.spark.read.{VortexFooterReader, VortexPartitionReader}
3329
import dev.vortex.spark.write.{SparkToArrowSchema, VortexOutputWriterFactory}
3430

3531
/** V1 fallback used for catalog tables and when Vortex is listed in
@@ -82,7 +78,7 @@ final class VortexFileFormat
8278
): Option[Seq[String]] =
8379
Some(
8480
Seq.fill(requiredSchema.length)(
85-
classOf[VortexArrowColumnVector].getName
81+
classOf[ColumnVector].getName
8682
) ++
8783
Seq.fill(partitionSchema.length)(classOf[ConstantColumnVector].getName)
8884
)
@@ -106,6 +102,7 @@ final class VortexFileFormat
106102
): PartitionedFile => Iterator[InternalRow] = {
107103
val optionMap = VortexOptions.of(options.asJava)
108104
val io = VortexIo.create(optionMap, hadoopConf)
105+
val caseSensitive = spark.sessionState.conf.caseSensitiveAnalysis
109106
(file: PartitionedFile) => {
110107
val reader = new VortexPartitionReader(
111108
file,
@@ -114,7 +111,8 @@ final class VortexFileFormat
114111
new StructType(),
115112
io,
116113
optionMap,
117-
filters.toArray
114+
filters.toArray,
115+
caseSensitive
118116
)
119117
Option(TaskContext.get())
120118
.foreach(_.addTaskCompletionListener[Unit](_ => reader.close()))
@@ -176,6 +174,7 @@ final class VortexFileFormat
176174

177175
val optionMap = VortexOptions.of(options.asJava)
178176
val io = VortexIo.create(optionMap, hadoopConf)
177+
val caseSensitive = spark.sessionState.conf.caseSensitiveAnalysis
179178
(file: PartitionedFile) => {
180179
val reader = new VortexPartitionReader(
181180
file,
@@ -184,7 +183,8 @@ final class VortexFileFormat
184183
partitionSchema,
185184
io,
186185
optionMap,
187-
filters.toArray
186+
filters.toArray,
187+
caseSensitive
188188
)
189189
Option(TaskContext.get()).foreach(
190190
_.addTaskCompletionListener[Unit](_ => reader.close())

‎java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ final class VortexScan(
4949
private val hadoopConf =
5050
sparkSession.sessionState.newHadoopConfWithOptions(caseSensitiveOptions)
5151
private val io = VortexIo.create(options, hadoopConf)
52+
private val caseSensitive =
53+
sparkSession.sessionState.conf.caseSensitiveAnalysis
5254

5355
override def isSplitable(path: Path): Boolean = false
5456

@@ -75,7 +77,8 @@ final class VortexScan(
7577
dataSchema,
7678
readDataSchema,
7779
readPartitionSchema,
78-
pushedFilters
80+
pushedFilters,
81+
caseSensitive
7982
)
8083
}
8184
}
@@ -117,6 +120,7 @@ final class VortexScan(
117120
super.equals(
118121
scan
119122
) && dataSchema == scan.dataSchema && options == scan.options &&
123+
caseSensitive == scan.caseSensitive &&
120124
equivalentFilters(pushedFilters, scan.pushedFilters) && aggregationsEqual
121125
case _ => false
122126
}

‎java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexFilterPushdownTest.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@
3434
import org.junit.jupiter.api.Test;
3535
import org.junit.jupiter.api.TestInstance;
3636
import org.junit.jupiter.api.io.TempDir;
37+
import org.junit.jupiter.params.ParameterizedTest;
38+
import org.junit.jupiter.params.provider.ValueSource;
3739

3840
/**
3941
* Tests that Spark predicate pushdown into the Vortex datasource produces correct results.
@@ -70,6 +72,31 @@ public void tearDown() {
7072
}
7173
}
7274

75+
@ParameterizedTest
76+
@ValueSource(booleans = {false, true})
77+
void nullSafeEqualityPreservesNullsUnderNegation(boolean v1) {
78+
SparkSession session = spark.newSession();
79+
SparkSession.setActiveSession(session);
80+
session.conf().set("spark.sql.sources.useV1SourceList", v1 ? "vortex" : "");
81+
Dataset<Row> source = session.sql(
82+
"SELECT 0 AS id, CAST(NULL AS BIGINT) AS value UNION ALL SELECT 1, 1L UNION ALL SELECT 2, 2L");
83+
Path output = tempDir.resolve("null_safe_" + v1);
84+
source.coalesce(1).write().format("vortex").save(output.toString());
85+
Dataset<Row> data = session.read().format("vortex").load(output.toString());
86+
87+
for (String predicate : List.of(
88+
"NOT (value <=> 1L)",
89+
"value <=> 1L",
90+
"value <=> NULL",
91+
"NOT (value <=> NULL)",
92+
"NOT ((value <=> 1L) OR id = 2)")) {
93+
assertEquals(
94+
source.where(predicate).orderBy("id").collectAsList(),
95+
data.where(predicate).orderBy("id").collectAsList(),
96+
predicate);
97+
}
98+
}
99+
73100
@Test
74101
@DisplayName("Equality, comparison, IS NULL, IN, AND/OR/NOT all return correct rows after pushdown")
75102
public void testFilterPushdownCorrectness() throws IOException {
@@ -574,6 +601,7 @@ private static List<Integer> idsOf(Dataset<Row> df) {
574601

575602
@AfterEach
576603
public void cleanupTempFiles() throws IOException {
604+
SparkSession.setActiveSession(spark);
577605
if (tempDir != null && Files.exists(tempDir)) {
578606
try (Stream<Path> paths = Files.walk(tempDir)) {
579607
paths.sorted(Comparator.reverseOrder()).forEach(path -> {

‎java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexProjectionTest.java‎

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,22 +5,29 @@
55

66
import static org.junit.jupiter.api.Assertions.assertEquals;
77
import static org.junit.jupiter.api.Assertions.assertFalse;
8+
import static org.junit.jupiter.api.Assertions.assertThrows;
89
import static org.junit.jupiter.api.Assertions.assertTrue;
910

1011
import java.nio.file.Path;
1112
import java.util.List;
1213
import org.apache.spark.sql.Dataset;
1314
import org.apache.spark.sql.Row;
15+
import org.apache.spark.sql.RowFactory;
1416
import org.apache.spark.sql.SaveMode;
1517
import org.apache.spark.sql.SparkSession;
18+
import org.apache.spark.sql.types.DataTypes;
19+
import org.apache.spark.sql.types.StructType;
1620
import org.junit.jupiter.api.AfterAll;
21+
import org.junit.jupiter.api.AfterEach;
1722
import org.junit.jupiter.api.BeforeAll;
1823
import org.junit.jupiter.api.Test;
1924
import org.junit.jupiter.api.TestInstance;
2025
import org.junit.jupiter.api.io.TempDir;
26+
import org.junit.jupiter.params.ParameterizedTest;
27+
import org.junit.jupiter.params.provider.CsvSource;
2128

2229
/**
23-
* Reads that need no data column at all.
30+
* Column pruning and resolution of user-specified read schemas.
2431
*
2532
* <p>A query over partition columns alone, and a count Spark declines to push down, both leave the read data schema
2633
* empty. The scan must still push an empty projection, or it pulls every column off storage to answer them.
@@ -50,6 +57,59 @@ public void tearDown() {
5057
}
5158
}
5259

60+
@AfterEach
61+
void restoreActiveSession() {
62+
SparkSession.setActiveSession(spark);
63+
}
64+
65+
@ParameterizedTest
66+
@CsvSource({"v2,false", "v2,true", "v1,false", "v1,true", "v1_rows,false", "v1_rows,true"})
67+
void explicitSchemaRespectsCaseSensitivity(String mode, boolean caseSensitive) {
68+
SparkSession session = newActiveSession();
69+
session.conf().set("spark.sql.sources.useV1SourceList", mode.startsWith("v1") ? "vortex" : "");
70+
session.conf().set("spark.sql.codegen.wholeStage", !mode.equals("v1_rows"));
71+
session.conf().set("spark.sql.caseSensitive", caseSensitive);
72+
Path output = tempDir.resolve("schema_case_" + mode + "_" + caseSensitive);
73+
session.range(1, 3).coalesce(1).write().format("vortex").save(output.toString());
74+
75+
Dataset<Row> data = session.read()
76+
.format("vortex")
77+
.schema(new StructType().add("ID", DataTypes.LongType, true))
78+
.load(output.toString());
79+
List<Row> expected = caseSensitive
80+
? List.of(RowFactory.create((Object) null), RowFactory.create((Object) null))
81+
: List.of(RowFactory.create(1L), RowFactory.create(2L));
82+
assertEquals(expected, data.orderBy("ID").collectAsList());
83+
assertEquals(
84+
caseSensitive ? List.of() : List.of(RowFactory.create(2L)),
85+
data.where("ID > 1").collectAsList());
86+
}
87+
88+
@Test
89+
void explicitSchemaRejectsAmbiguousColumnNames() {
90+
SparkSession session = newActiveSession();
91+
session.conf().set("spark.sql.sources.useV1SourceList", "");
92+
session.conf().set("spark.sql.caseSensitive", true);
93+
Path output = tempDir.resolve("ambiguous_schema");
94+
session.range(1)
95+
.selectExpr("id", "id + 1 AS ID")
96+
.coalesce(1)
97+
.write()
98+
.format("vortex")
99+
.save(output.toString());
100+
session.conf().set("spark.sql.caseSensitive", false);
101+
Dataset<Row> data = session.read()
102+
.format("vortex")
103+
.schema(new StructType().add("id", DataTypes.LongType, true))
104+
.load(output.toString());
105+
106+
Throwable failure = assertThrows(Exception.class, data::collectAsList);
107+
while (failure.getCause() != null) {
108+
failure = failure.getCause();
109+
}
110+
assertTrue(failure.getMessage().contains("Ambiguous columns id and ID"), failure.getMessage());
111+
}
112+
53113
@Test
54114
void selectingOnlyAPartitionColumnReadsNoDataColumn() {
55115
Dataset<Row> data = writeAndRead("partition_only", true).select("group");
@@ -89,6 +149,13 @@ private String plan(Dataset<Row> data) {
89149
return data.queryExecution().executedPlan().toString();
90150
}
91151

152+
private SparkSession newActiveSession() {
153+
SparkSession session = spark.newSession();
154+
// Spark 3.5's FileDataSourceV2 resolves its session through SparkSession.active.
155+
SparkSession.setActiveSession(session);
156+
return session;
157+
}
158+
92159
private Dataset<Row> writeAndRead(String name, boolean partitioned) {
93160
Path output = tempDir.resolve(name);
94161
write(output, partitioned);

0 commit comments

Comments
 (0)