().all {
dependsOn("makeTestFiles")
}
diff --git a/java/vortex-jni/src/main/java/dev/vortex/io/Closeables.java b/java/vortex-jni/src/main/java/dev/vortex/io/Closeables.java
new file mode 100644
index 00000000000..cdd6afdded2
--- /dev/null
+++ b/java/vortex-jni/src/main/java/dev/vortex/io/Closeables.java
@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.io;
+
+import java.io.Closeable;
+import java.io.IOException;
+
+/** Helpers for closing groups of resources. */
+public final class Closeables {
+ private Closeables() {}
+
+ /**
+ * Closes every item, always attempting all of them; the first failure is thrown once everything has been visited,
+ * with any further failures suppressed.
+ */
+ public static void closeAll(Iterable closeables) throws IOException {
+ IOException failure = null;
+ for (Closeable closeable : closeables) {
+ try {
+ closeable.close();
+ } catch (IOException e) {
+ if (failure == null) {
+ failure = e;
+ } else {
+ failure.addSuppressed(e);
+ }
+ }
+ }
+ if (failure != null) {
+ throw failure;
+ }
+ }
+}
diff --git a/java/vortex-jni/src/main/java/dev/vortex/io/PooledReadable.java b/java/vortex-jni/src/main/java/dev/vortex/io/PooledReadable.java
new file mode 100644
index 00000000000..e3ab92f0709
--- /dev/null
+++ b/java/vortex-jni/src/main/java/dev/vortex/io/PooledReadable.java
@@ -0,0 +1,171 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.io;
+
+import com.google.common.base.Preconditions;
+import java.io.Closeable;
+import java.io.EOFException;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Locale;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+
+/**
+ * A {@link NativeReadable} that serves each of Vortex's concurrent reads from its own stream.
+ *
+ * Streams are opened on demand and reused until this readable is closed. Subclasses supply only
+ * {@link #openStream()}.
+ */
+public abstract class PooledReadable implements NativeReadable {
+ /** Largest scratch array retained for reuse on a pooled stream; bigger reads get a one-shot array. */
+ private static final int SCRATCH_RETAIN_LIMIT = 1 << 20;
+
+ private final String name;
+ private final long length;
+ private final Queue pool = new ConcurrentLinkedQueue<>();
+ private volatile boolean closed = false;
+
+ /**
+ * @param name stable unique name for this source, as required by {@link NativeReadable#name()}
+ * @param length total size of the source in bytes
+ */
+ protected PooledReadable(String name, long length) {
+ Preconditions.checkArgument(name != null && !name.isEmpty(), "name is required");
+ Preconditions.checkArgument(length >= 0, "length must not be negative: %s", length);
+ this.name = name;
+ this.length = length;
+ }
+
+ /** Opens one more handle on this source. Called whenever a read finds no free stream in the pool. */
+ protected abstract PositionalStream openStream() throws IOException;
+
+ @Override
+ public final String name() {
+ return name;
+ }
+
+ @Override
+ public final long length() {
+ return length;
+ }
+
+ @Override
+ public final void readFully(long position, ByteBuffer buffer) throws IOException {
+ Preconditions.checkState(!closed, "Cannot read %s: already closed", name);
+ int requested = buffer.remaining();
+ if (position < 0 || position + requested > length) {
+ throw new EOFException(String.format(
+ Locale.ROOT,
+ "Cannot read %d bytes at position %d: %s is %d bytes long",
+ requested,
+ position,
+ name,
+ length));
+ }
+
+ PooledStream pooled = pool.poll();
+ if (pooled == null) {
+ pooled = new PooledStream(openStream());
+ }
+
+ try {
+ pooled.stream.readFully(position, buffer, pooled);
+ } catch (IOException | RuntimeException e) {
+ try {
+ pooled.close();
+ } catch (IOException suppressed) {
+ e.addSuppressed(suppressed);
+ }
+ throw e;
+ }
+
+ if (buffer.hasRemaining()) {
+ EOFException failure = new EOFException(String.format(
+ Locale.ROOT,
+ "Read of %d bytes at position %d in %s left %d bytes unfilled",
+ requested,
+ position,
+ name,
+ buffer.remaining()));
+ try {
+ pooled.close();
+ } catch (IOException e) {
+ failure.addSuppressed(e);
+ }
+ throw failure;
+ }
+
+ pool.add(pooled);
+ if (closed) {
+ // Racing with close(): make sure nothing this read returned to the pool leaks. The read
+ // itself succeeded, so a failure to close an idle stream must not fail it retroactively.
+ try {
+ closeAllPooled();
+ } catch (IOException ignored) {
+ // Best-effort: the streams were drained from the pool and close was attempted on each.
+ }
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
+ closed = true;
+ closeAllPooled();
+ }
+
+ private void closeAllPooled() throws IOException {
+ List drained = new ArrayList<>();
+ PooledStream pooled;
+ while ((pooled = pool.poll()) != null) {
+ drained.add(pooled);
+ }
+ Closeables.closeAll(drained);
+ }
+
+ /** One handle on the source, used by at most one read at a time. */
+ public interface PositionalStream extends Closeable {
+ /**
+ * Fills {@code buffer} completely from absolute {@code position}, in as few requests to storage as possible:
+ * Vortex has already coalesced neighbouring reads into this one.
+ *
+ * Implementations limited to {@code byte[]} APIs should stage through {@code scratch} rather than allocate.
+ */
+ void readFully(long position, ByteBuffer buffer, ScratchBytes scratch) throws IOException;
+ }
+
+ /** A staging array, reused across the reads served by one stream. */
+ public interface ScratchBytes {
+ /** Returns an array of at least {@code length} bytes, whose contents are undefined. */
+ byte[] bytes(int length);
+ }
+
+ private static final class PooledStream implements ScratchBytes, Closeable {
+ private final PositionalStream stream;
+ private byte[] scratch = new byte[0];
+
+ private PooledStream(PositionalStream stream) {
+ this.stream = stream;
+ }
+
+ @Override
+ public byte[] bytes(int length) {
+ if (scratch.length >= length) {
+ return scratch;
+ } else if (length > SCRATCH_RETAIN_LIMIT) {
+ return new byte[length];
+ }
+
+ scratch = new byte[length];
+ return scratch;
+ }
+
+ @Override
+ public void close() throws IOException {
+ stream.close();
+ }
+ }
+}
diff --git a/java/vortex-spark/README.md b/java/vortex-spark/README.md
index ea27b747ebe..07e1d593ec6 100644
--- a/java/vortex-spark/README.md
+++ b/java/vortex-spark/README.md
@@ -1,115 +1,90 @@
# vortex-spark
-A Spark DataSource V2 connector for reading and writing [Vortex](https://vortex.dev) files.
-It registers itself under the format name `vortex` and supports both the DataFrame API and
-Spark SQL.
+A Spark file data source for reading and writing [Vortex](https://vortex.dev) files. It
+registers itself as `vortex` and supports the DataFrame API and Spark SQL without session
+extensions or custom catalogs.
-Two flavors are published to Maven Central:
+Choose the artifact matching both Spark and Scala:
-| Artifact | Spark | Scala |
-|-------------------------------|-----------|-------|
-| `dev.vortex:vortex-spark_2.13` | Spark 4.x | 2.13 |
-| `dev.vortex:vortex-spark_2.12` | Spark 3.5.x | 2.12 |
+| Artifact | Spark | Scala |
+|------------------------------------|---------------|-------|
+| `dev.vortex:vortex-spark-3.5_2.12` | 3.5.x | 2.12 |
+| `dev.vortex:vortex-spark-3.5_2.13` | 3.5.x | 2.13 |
+| `dev.vortex:vortex-spark-4.0_2.13` | 4.0.x and 4.1.x | 2.13 |
-Use the `all` classifier JAR (e.g. `vortex-spark_2.13-0.83.0-all.jar`). It is self-contained:
-it bundles the Vortex JNI bindings, native libraries for Linux (x86_64 and aarch64) and macOS
-(aarch64), and relocates its Arrow, Guava, and Jackson dependencies to avoid classpath
-conflicts with Spark. The thin (unclassified) JAR does not work on its own because it
-references relocated classes that only ship in the `all` JAR.
+Use the `all` classifier JAR. It bundles the Vortex JNI bindings, native libraries, and
+relocated Arrow, Guava, and Jackson dependencies. The unclassified thin JAR is not usable by
+itself.
-## Getting Vortex into Spark
-
-Pass the `all` JAR to `spark-shell`, `spark-submit`, or `pyspark` with `--jars`. Spark accepts
-either a local path or a URL, so you can point directly at Maven Central:
+For example, with version `VERSION`:
```shell
-spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark_2.13/0.83.0/vortex-spark_2.13-0.83.0-all.jar
+spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark-4.0_2.13/VERSION/vortex-spark-4.0_2.13-VERSION-all.jar
```
-Or configure it on the session builder, e.g. in PySpark:
-
-```python
-spark = (
- SparkSession.builder
- .config("spark.jars", "/path/to/vortex-spark_2.13-0.83.0-all.jar")
- .getOrCreate()
-)
-```
-
-Note that `--packages dev.vortex:vortex-spark_2.13:0.83.0` does not work: `--packages` cannot
-select the `all` classifier and resolves the thin JAR, which fails at runtime with
-`NoClassDefFoundError: dev/vortex/relocated/...`.
-
-To depend on the connector from a JVM project instead, add the `all` classifier to the
-dependency:
+Or add the classified artifact to a JVM project:
```kotlin
-implementation("dev.vortex:vortex-spark_2.13:0.83.0:all")
+implementation("dev.vortex:vortex-spark-4.0_2.13:VERSION:all")
```
-## Usage
-
-Paths may be local filesystem paths (`/path/to/data`) or URLs (`file:///path/to/data`,
-`s3://bucket/path/to/data`).
+Spark's `--packages` flag cannot select the `all` classifier, so pass the classified JAR with
+`--jars` instead.
-### DataFrame API
+## Usage
```java
-// Write
df.write()
.format("vortex")
- .option("path", "/path/to/output")
.mode(SaveMode.Overwrite)
- .save();
+ .save("/path/to/output");
-// Read a single file or a directory of .vortex files
-Dataset df = spark.read()
+Dataset result = spark.read()
.format("vortex")
- .option("path", "/path/to/output")
- .load();
+ .load("/path/to/output");
```
-### Spark SQL
+Spark discovers files, Hive-style partitions, and storage credentials through its Hadoop
+configuration. This also enables standard file-source options such as `pathGlobFilter` and
+`recursiveFileLookup`. All file content is read and written through Hadoop streams.
-```sql
--- Query existing Vortex files through a temporary view
-CREATE TEMPORARY VIEW people
-USING vortex
-OPTIONS (path '/path/to/data');
-
-SELECT name, age FROM people WHERE age > 30;
-
--- Create a table and write to it. With a LOCATION clause the table is external,
--- backed by the files at that path; without one it is managed by Spark.
-CREATE TABLE student (id INT, name STRING, age INT)
-USING vortex;
-
-INSERT INTO student VALUES (1, 'Alice', 20), (2, 'Bob', 21);
+Every file in a Vortex dataset must end with `.vortex`. Writes produce that extension, and a
+dataset holding any other file is rejected rather than read in part.
+```sql
+CREATE TABLE student (id INT, name STRING) USING vortex;
+INSERT INTO student VALUES (1, 'Alice'), (2, 'Bob');
SELECT * FROM student;
+
+SELECT * FROM vortex.`/path/to/output`;
```
-On Spark 3.5, `CREATE TABLE ... USING vortex` additionally requires replacing the session
-catalog with `spark.sql.catalog.spark_catalog=dev.vortex.spark.VortexSessionCatalog`,
-because Spark 3.5's built-in catalog cannot read tables backed by a DataSource V2-only
-connector. The extension delegates everything else to the built-in session catalog and
-leaves tables of other providers untouched; it is not needed on Spark 4.
+`CREATE TABLE ... USING vortex` and direct path queries work on Spark 3.5 and 4.x without
+catalog configuration. Catalog tables read and write through the V1 file format;
+`spark.sql.sources.useV1SourceList=vortex` puts every path on it, which is what the `_metadata`
+column and dynamic partition pruning need.
-### Direct file queries
+## Benchmarks
-Spark's built-in ``SELECT * FROM format.`path` `` syntax only works for built-in file
-formats, so the connector ships a path-based catalog that provides the equivalent for
-Vortex. Register it under the name `vortex` with this session config:
+JMH benchmarks live in `common/src/jmh/java`. They measure the per-file open cost of the
+framework's file-splitting model and the footer reads behind statistics and `COUNT(*)` pushdown.
+Run them for one variant and pass JMH arguments through `-PjmhArgs`:
-```
-spark.sql.catalog.vortex=dev.vortex.spark.VortexCatalog
+```bash
+cd java
+./gradlew :vortex-spark-4.0_2.13:jmh -PjmhArgs="SparkScanBenchmark -p fileCount=1,128"
+./gradlew :vortex-spark-4.0_2.13:jmh -PjmhArgs="FooterReadBenchmark -f 1 -wi 2 -i 5"
```
-Then query (or insert into) Vortex files directly by path:
+## Migrating from the old connector
-```sql
-SELECT * FROM vortex.`/path/to/data`;
-```
+- Replace `vortex-spark_2.12` or `vortex-spark_2.13` with the versioned artifact above.
+- Remove `spark.sql.catalog.spark_catalog=dev.vortex.spark.VortexSessionCatalog` and
+ `spark.sql.catalog.vortex=dev.vortex.spark.VortexCatalog`; those classes no longer exist.
+- Configure remote filesystems through Hadoop and use their Hadoop schemes, such as `s3a://`
+ and `abfs://`. Spark's file index now owns listing and skips hidden `_`-prefixed files.
+- All reads and writes use Hadoop streams. The `vortex.io` option and Vortex's native storage
+ clients (with their `aws_*`/`azure_*` options) are no longer available from Spark.
-See the [Spark user guide](https://docs.vortex.dev/user-guide/spark.html) for the full
-documentation, including supported types, write options, and S3 configuration.
+See the [Spark user guide](https://docs.vortex.dev/user-guide/spark.html) for supported types,
+options, partitioning, and remote storage details.
diff --git a/java/vortex-spark/build.gradle.kts b/java/vortex-spark/build.gradle.kts
deleted file mode 100644
index 9dc46f4c6d9..00000000000
--- a/java/vortex-spark/build.gradle.kts
+++ /dev/null
@@ -1,142 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
-
-plugins {
- `java-library`
- `jvm-test-suite`
- id("com.gradleup.shadow") version "9.4.2"
-}
-
-// Derive Scala and Spark versions from the Gradle project name (vortex-spark_2.12 or vortex-spark_2.13)
-val scalaVersion: String = project.name.substringAfterLast("_")
-val sparkVersion: String =
- when (scalaVersion) {
- "2.12" -> {
- libs.versions.spark3.get()
- }
-
- "2.13" -> {
- libs.versions.spark4.get()
- }
-
- else -> {
- throw GradleException(
- "Unsupported Scala version: $scalaVersion (project name must end with _2.12 or _2.13)",
- )
- }
- }
-
-// Both vortex-spark_2.12 and vortex-spark_2.13 share this projectDir.
-// Give each its own build directory to avoid output conflicts.
-layout.buildDirectory = layout.projectDirectory.dir("build/${project.name}")
-
-dependencies {
- compileOnly("org.apache.spark:spark-catalyst_$scalaVersion:$sparkVersion")
- compileOnly("org.apache.spark:spark-sql_$scalaVersion:$sparkVersion")
- api(project(":vortex-jni", configuration = "shadow"))
-
- compileOnly(libs.immutables.value)
- annotationProcessor(libs.immutables.value)
-
- implementation(libs.guava)
- implementation(libs.slf4j.api)
-}
-
-testing {
- suites {
- val test by getting(JvmTestSuite::class) {
- useJUnitJupiter()
-
- dependencies {
- implementation(libs.junit.jupiter)
- implementation("org.apache.spark:spark-core_$scalaVersion:$sparkVersion")
- implementation("org.apache.spark:spark-sql_$scalaVersion:$sparkVersion")
- implementation(libs.s3mock.testcontainers)
- implementation(libs.testcontainers.juputer)
- runtimeOnly(libs.slf4j.simple)
- if (scalaVersion == "2.12") {
- // Spark 3.5 marks javax.servlet-api as provided; needed at test runtime for MetricsServlet
- runtimeOnly("javax.servlet:javax.servlet-api:4.0.1")
- }
- }
- }
- }
-}
-
-mavenPublishing {
- coordinates(groupId = "dev.vortex", artifactId = "vortex-spark_$scalaVersion", version = "${rootProject.version}")
-
- publishToMavenCentral()
-
- if (!project.hasProperty("skip.signing")) {
- signAllPublications()
- }
-
- pom {
- name = "vortex-spark_$scalaVersion"
- description = project.description
- url = "https://vortex.dev"
- inceptionYear = "2025"
-
- licenses {
- license {
- name = "Apache-2.0"
- url = "https://spdx.org/licenses/Apache-2.0.html"
- }
- }
- developers {
- developer {
- id = "spiraldb"
- name = "Vortex Authors"
- }
- }
- scm {
- connection = "scm:git:https://github.com/spiraldb/vortex.git"
- developerConnection = "scm:git:ssh://github.com/spiraldb/vortex.git"
- url = "https://github.com/spiraldb/vortex"
- }
- }
-
- repositories {
- mavenCentral()
- mavenLocal()
- }
-}
-
-// shade guava and arrow dependencies
-tasks.withType {
- relocate("com.google.common", "dev.vortex.relocated.com.google.common")
- relocate("org.apache.arrow", "dev.vortex.relocated.org.apache.arrow") {
- // exclude C Data Interface since JNI cannot be relocated
- exclude("org.apache.arrow.c.jni.JniWrapper")
- exclude("org.apache.arrow.c.jni.PrivateData")
- exclude("org.apache.arrow.c.jni.CDataJniException")
- // Also used by JNI: https://github.com/apache/arrow/blob/apache-arrow-11.0.0/java/c/src/main/cpp/jni_wrapper.cc#L341
- // Note this class is not used by us, but required when loading the native lib
- exclude("org.apache.arrow.c.ArrayStreamExporter\$ExportedArrayStreamPrivateData")
- }
- relocate("com.fasterxml.jackson", "dev.vortex.relocated.com.fasterxml.jackson")
-}
-
-tasks.withType().all {
- classpath +=
- project(":vortex-jni")
- .tasks
- .named("shadowJar")
- .get()
- .outputs.files
- jvmArgs(
- "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED",
- "--add-opens=java.base/java.nio=ALL-UNNAMED",
- "--add-opens=java.base/sun.util.calendar=ALL-UNNAMED",
- "--add-opens=java.base/sun.security.action=ALL-UNNAMED",
- )
-}
-
-tasks.build {
- dependsOn("shadowJar")
-}
-
-description = "Apache Spark bindings for reading Vortex file datasets"
diff --git a/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/BenchmarkSparkSession.java b/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/BenchmarkSparkSession.java
new file mode 100644
index 00000000000..d78b7d21efe
--- /dev/null
+++ b/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/BenchmarkSparkSession.java
@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.bench;
+
+import java.nio.file.Path;
+import org.apache.spark.sql.SaveMode;
+import org.apache.spark.sql.SparkSession;
+
+/** Local Spark session and dataset writer shared by the benchmarks. */
+final class BenchmarkSparkSession {
+ private BenchmarkSparkSession() {}
+
+ static SparkSession create(String appName) {
+ return SparkSession.builder()
+ .appName(appName)
+ .master("local[4]")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.ui.enabled", "false")
+ .config("spark.sql.shuffle.partitions", "4")
+ .getOrCreate();
+ }
+
+ /**
+ * Writes {@code rows} rows spread over {@code fileCount} Vortex files. Each row carries an integer id, a string,
+ * and a double so scans touch several encodings.
+ */
+ static void writeDataset(SparkSession spark, Path output, long rows, int fileCount) {
+ spark.range(0, rows)
+ .selectExpr(
+ "cast(id as int) as id",
+ "concat('value_', cast(id % 1000 as string)) as value",
+ "cast(id as double) / 7.0 as measure")
+ .repartition(fileCount)
+ .write()
+ .format("vortex")
+ .mode(SaveMode.Overwrite)
+ .save(output.toString());
+ }
+}
diff --git a/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/FooterReadBenchmark.java b/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/FooterReadBenchmark.java
new file mode 100644
index 00000000000..124340ce9b0
--- /dev/null
+++ b/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/FooterReadBenchmark.java
@@ -0,0 +1,98 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.bench;
+
+import dev.vortex.spark.VortexOptions;
+import dev.vortex.spark.io.VortexFile;
+import dev.vortex.spark.io.VortexIo;
+import dev.vortex.spark.read.VortexFooterReader;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.OptionalLong;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Stream;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.spark.sql.SparkSession;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+
+/**
+ * Cost of opening a Vortex file just to read its footer.
+ *
+ * Footer reads drive schema inference, row-count statistics, and COUNT(*) pushdown, so each of those pays this price
+ * once per file. The {@code ioMode} axis compares Hadoop stream upcalls with native file reads.
+ */
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@Warmup(iterations = 3)
+@Measurement(iterations = 5)
+@Fork(1)
+public class FooterReadBenchmark {
+ private static final long ROWS = 500_000L;
+ private static final int FILE_COUNT = 32;
+
+ @Param({"1", "8"})
+ private int statsParallelism;
+
+ private SparkSession spark;
+ private Path dataDir;
+ private VortexIo io;
+ private VortexOptions options;
+ private List files;
+
+ @Setup(Level.Trial)
+ public void setUp() throws IOException {
+ spark = BenchmarkSparkSession.create("FooterReadBenchmark");
+ dataDir = Files.createTempDirectory("vortex-footer-bench");
+ Path data = dataDir.resolve("data");
+ BenchmarkSparkSession.writeDataset(spark, data, ROWS, FILE_COUNT);
+
+ options = VortexOptions.of(
+ Map.of(VortexFooterReader.FOOTER_PARALLELISM_OPTION, Integer.toString(statsParallelism)));
+ io = VortexIo.create(options, new Configuration());
+
+ files = new ArrayList<>();
+ try (Stream paths = Files.list(data)) {
+ for (Path path : paths.toList()) {
+ if (path.getFileName().toString().endsWith(".vortex")) {
+ files.add(new VortexFile(path.toUri().toString(), Files.size(path)));
+ }
+ }
+ }
+ }
+
+ @TearDown(Level.Trial)
+ public void tearDown() throws IOException {
+ spark.stop();
+ SparkScanBenchmark.deleteRecursively(dataDir);
+ }
+
+ /** One footer open and row-count read, the unit of work behind every statistic. */
+ @Benchmark
+ public OptionalLong singleFooterRowCount() {
+ return VortexFooterReader.estimatedRowCount(files.get(0), io, options);
+ }
+
+ /** Footer row counts summed over all files through the bounded pool the scan uses for statistics. */
+ @Benchmark
+ public OptionalLong sumRowCountsAcrossFiles() {
+ return VortexFooterReader.sumRowCounts(files, io, options);
+ }
+}
diff --git a/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/SparkScanBenchmark.java b/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/SparkScanBenchmark.java
new file mode 100644
index 00000000000..fcb26b52a31
--- /dev/null
+++ b/java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/SparkScanBenchmark.java
@@ -0,0 +1,101 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.bench;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Comparator;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Stream;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSession;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+
+/**
+ * End-to-end Spark scans over a fixed number of rows split into a varying number of files.
+ *
+ * The connector opens one Vortex reader per {@code PartitionedFile}, so the {@code fileCount} axis measures the
+ * per-file open overhead that bin-packing adds on top of the pure decode cost.
+ */
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Warmup(iterations = 3)
+@Measurement(iterations = 5)
+@Fork(1)
+public class SparkScanBenchmark {
+ private static final long ROWS = 2_000_000L;
+
+ @Param({"1", "16", "128"})
+ private int fileCount;
+
+ private SparkSession spark;
+ private Path dataDir;
+
+ @Setup(Level.Trial)
+ public void setUp() throws IOException {
+ spark = BenchmarkSparkSession.create("SparkScanBenchmark");
+ dataDir = Files.createTempDirectory("vortex-scan-bench");
+ BenchmarkSparkSession.writeDataset(spark, dataDir.resolve("data"), ROWS, fileCount);
+ }
+
+ @TearDown(Level.Trial)
+ public void tearDown() throws IOException {
+ spark.stop();
+ deleteRecursively(dataDir);
+ }
+
+ /** COUNT(*) is pushed down to footer row counts, so this measures footer reads plus planning. */
+ @Benchmark
+ public long countStar() {
+ return read().count();
+ }
+
+ /** Reads every column of every file and aggregates, measuring the full decode path. */
+ @Benchmark
+ public Row fullScanAggregate() {
+ return read().selectExpr("sum(id)", "count(value)", "sum(measure)").first();
+ }
+
+ /** Pushes a predicate into each per-file reader and counts the survivors. */
+ @Benchmark
+ public long filteredScan() {
+ return read().filter("id % 97 = 0").count();
+ }
+
+ /** Projects a single column to isolate per-file open cost from decode volume. */
+ @Benchmark
+ public Row singleColumnScan() {
+ return read().selectExpr("sum(id)").first();
+ }
+
+ private Dataset read() {
+ return spark.read().format("vortex").load(dataDir.resolve("data").toString());
+ }
+
+ static void deleteRecursively(Path root) throws IOException {
+ if (!Files.exists(root)) {
+ return;
+ }
+ try (Stream paths = Files.walk(root)) {
+ for (Path path : paths.sorted(Comparator.reverseOrder()).toList()) {
+ Files.delete(path);
+ }
+ }
+ }
+}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/ArrowUtils.java
similarity index 100%
rename from java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java
rename to java/vortex-spark/common/src/main/java/dev/vortex/spark/ArrowUtils.java
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexOptions.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexOptions.java
new file mode 100644
index 00000000000..92d6af05750
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexOptions.java
@@ -0,0 +1,122 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark;
+
+import com.google.common.collect.ImmutableMap;
+import java.io.Serializable;
+import java.util.HashMap;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+
+/**
+ * The format options of one Spark read or write, resolved without regard to key case.
+ *
+ * Spark treats data source options as case-insensitive, but the connector is reached through maps that disagree: a
+ * V2 scan carries the keys as the user typed them, while the V1 file format lower-cases them first. Every
+ * {@code vortex.*} lookup goes through here so both paths answer the same.
+ *
+ *
The keys as given are kept too. Hadoop configuration keys are case-sensitive, so {@link #asCaseSensitiveMap()} is
+ * what feeds a Hadoop {@code Configuration}.
+ *
+ *
Built on the driver and shipped to executors inside the reader and writer factories.
+ */
+public final class VortexOptions implements Serializable {
+ private static final long serialVersionUID = 1L;
+
+ private final ImmutableMap original;
+ private final ImmutableMap byLowerCasedKey;
+
+ /** Captures {@code options} as given. A null map is read as no options at all. */
+ public static VortexOptions of(Map options) {
+ ImmutableMap original = options == null ? ImmutableMap.of() : ImmutableMap.copyOf(options);
+
+ // A map arriving from Spark may already hold keys that differ only in case. Later entries win, as they
+ // do in Spark's own CaseInsensitiveStringMap, and `original` still carries both.
+ Map lowerCased = new HashMap<>(original.size());
+ original.forEach((key, value) -> lowerCased.put(key.toLowerCase(Locale.ROOT), value));
+ return new VortexOptions(original, ImmutableMap.copyOf(lowerCased));
+ }
+
+ /** No options at all, for callers with no Spark read or write to draw them from. */
+ public static VortexOptions empty() {
+ return of(ImmutableMap.of());
+ }
+
+ private VortexOptions(ImmutableMap original, ImmutableMap byLowerCasedKey) {
+ this.original = original;
+ this.byLowerCasedKey = byLowerCasedKey;
+ }
+
+ /** The options with the keys as given, for the case-sensitive world of Hadoop configuration. */
+ public Map asCaseSensitiveMap() {
+ return original;
+ }
+
+ /** The value set for {@code key} under any casing. */
+ public Optional get(String key) {
+ Objects.requireNonNull(key, "key");
+ return Optional.ofNullable(byLowerCasedKey.get(key.toLowerCase(Locale.ROOT)));
+ }
+
+ /** The value set for {@code key} under any casing, or {@code fallback} when it is not set. */
+ public String get(String key, String fallback) {
+ return get(key).orElse(fallback);
+ }
+
+ /**
+ * The value set for {@code key} as an integer, or {@code fallback} when it is not set.
+ *
+ * @throws IllegalArgumentException if the value is not an integer
+ */
+ public int getInt(String key, int fallback) {
+ Optional value = get(key);
+ if (value.isEmpty()) {
+ return fallback;
+ }
+ try {
+ return Integer.parseInt(value.get().trim());
+ } catch (NumberFormatException e) {
+ throw new IllegalArgumentException(
+ String.format(Locale.ROOT, "%s must be an integer, got '%s'", key, value.get()), e);
+ }
+ }
+
+ /**
+ * The value set for {@code key} as a boolean, or {@code fallback} when it is not set.
+ *
+ * @throws IllegalArgumentException if the value is neither {@code true} nor {@code false}
+ */
+ public boolean getBoolean(String key, boolean fallback) {
+ Optional value = get(key);
+ if (value.isEmpty()) {
+ return fallback;
+ }
+ String trimmed = value.get().trim();
+ if (trimmed.equalsIgnoreCase("true")) {
+ return true;
+ }
+ if (trimmed.equalsIgnoreCase("false")) {
+ return false;
+ }
+ throw new IllegalArgumentException(
+ String.format(Locale.ROOT, "%s must be true or false, got '%s'", key, value.get()));
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ return other instanceof VortexOptions options && original.equals(options.original);
+ }
+
+ @Override
+ public int hashCode() {
+ return original.hashCode();
+ }
+
+ @Override
+ public String toString() {
+ return original.toString();
+ }
+}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionProvider.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexSessionProvider.java
similarity index 100%
rename from java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionProvider.java
rename to java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexSessionProvider.java
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSparkSession.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexSparkSession.java
similarity index 69%
rename from java/vortex-spark/src/main/java/dev/vortex/spark/VortexSparkSession.java
rename to java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexSparkSession.java
index 1d3d07d3b25..773a25b7a82 100644
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSparkSession.java
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexSparkSession.java
@@ -4,9 +4,11 @@
package dev.vortex.spark;
import dev.vortex.api.Session;
-import java.util.Map;
+import dev.vortex.jni.NativeRuntime;
import java.util.Objects;
+import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
/**
* JVM-wide holder for one or more Vortex {@link Session}s used by Spark readers and writers. The Rust side multiplexes
@@ -33,7 +35,13 @@ public final class VortexSparkSession {
/** Options key used to select a {@link VortexSessionProvider} by class name. */
public static final String PROVIDER_OPTION = "vortex.session.provider";
+ /** Options key sizing the JVM-wide pool of background threads that drive Vortex futures. */
+ public static final String WORKER_THREADS_OPTION = "vortex.workerThreads";
+
+ private static final int DEFAULT_WORKER_THREADS = 4;
+
private static final ConcurrentHashMap providerCache = new ConcurrentHashMap<>();
+ private static final AtomicBoolean runtimeConfigured = new AtomicBoolean();
private static volatile Session defaultSession;
private VortexSparkSession() {}
@@ -56,12 +64,32 @@ public static Session get() {
* Resolve the session to use for a given set of Spark format options. Honours the {@value #PROVIDER_OPTION} key;
* falls back to {@link #get()} otherwise.
*/
- public static Session get(Map options) {
- String providerClass = options == null ? null : options.get(PROVIDER_OPTION);
- if (providerClass == null || providerClass.isEmpty()) {
+ public static Session get(VortexOptions options) {
+ configureRuntime(options);
+ Optional providerClass = options.get(PROVIDER_OPTION).filter(name -> !name.isEmpty());
+ if (providerClass.isEmpty()) {
return get();
}
- return providerCache.computeIfAbsent(providerClass, VortexSparkSession::loadProvider);
+ return providerCache.computeIfAbsent(providerClass.get(), VortexSparkSession::loadProvider);
+ }
+
+ /**
+ * Sizes the shared worker pool from {@value #WORKER_THREADS_OPTION}, once per JVM.
+ *
+ * The pool is JVM-wide and resizing it starts or stops threads, so every read and write path resolves its
+ * session through here and the first caller wins. Repeating the call per file, or per task, would have concurrent
+ * tasks resize the pool underneath each other.
+ */
+ private static void configureRuntime(VortexOptions options) {
+ if (!runtimeConfigured.compareAndSet(false, true)) {
+ return;
+ }
+ int workers = options.getInt(WORKER_THREADS_OPTION, DEFAULT_WORKER_THREADS);
+ if (workers < 0) {
+ runtimeConfigured.set(false);
+ throw new IllegalArgumentException(WORKER_THREADS_OPTION + " must be >= 0, got " + workers);
+ }
+ NativeRuntime.setWorkerThreads(workers);
}
/**
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java
new file mode 100644
index 00000000000..26da27450b3
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java
@@ -0,0 +1,82 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.io;
+
+import dev.vortex.io.PooledReadable;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.ByteBuffer;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataInputStream;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.StreamCapabilities;
+
+/** A {@link dev.vortex.io.NativeReadable} that serves Vortex's reads from a Hadoop {@link FileSystem}. */
+public final class HadoopReadable extends PooledReadable {
+ private final Path path;
+ private final FileSystem fs;
+
+ /** Opens a readable over {@code path}, stating the file for its length. */
+ public static HadoopReadable open(Configuration conf, String path) {
+ return open(conf, path, -1);
+ }
+
+ /**
+ * Opens a readable over {@code path}. A non-negative {@code length} is taken as authoritative — the driver already
+ * stats every file when it plans the scan, so executors need not stat them again.
+ */
+ public static HadoopReadable open(Configuration conf, String path, long length) {
+ Path file = new Path(path);
+ try {
+ FileSystem fs = file.getFileSystem(conf);
+ long resolved = length >= 0 ? length : fs.getFileStatus(file).getLen();
+ return new HadoopReadable(path, file, fs, resolved);
+ } catch (IOException e) {
+ throw new UncheckedIOException("Failed to open " + path, e);
+ }
+ }
+
+ private HadoopReadable(String name, Path path, FileSystem fs, long length) {
+ super(name, length);
+ this.path = path;
+ this.fs = fs;
+ }
+
+ @Override
+ protected PositionalStream openStream() throws IOException {
+ return new HadoopStream(fs.open(path));
+ }
+
+ private static final class HadoopStream implements PositionalStream {
+ private final FSDataInputStream stream;
+ // Hadoop mandates this capability probe: streams may implement ByteBufferPositionedReadable
+ // yet throw UnsupportedOperationException when their inner stream cannot serve it.
+ private final boolean byteBufferPositionedRead;
+
+ private HadoopStream(FSDataInputStream stream) {
+ this.stream = stream;
+ this.byteBufferPositionedRead = stream.hasCapability(StreamCapabilities.PREADBYTEBUFFER);
+ }
+
+ @Override
+ public void readFully(long position, ByteBuffer buffer, ScratchBytes scratch) throws IOException {
+ if (byteBufferPositionedRead) {
+ // HDFS and friends fill Vortex's own memory, with no staging array in between.
+ stream.readFully(position, buffer);
+ return;
+ }
+
+ int length = buffer.remaining();
+ byte[] target = scratch.bytes(length);
+ stream.readFully(position, target, 0, length);
+ buffer.put(target, 0, length);
+ }
+
+ @Override
+ public void close() throws IOException {
+ stream.close();
+ }
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopWritable.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopWritable.java
new file mode 100644
index 00000000000..e68aa3fc5e1
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopWritable.java
@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.io;
+
+import dev.vortex.io.NativeWritable;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+
+/**
+ * A {@link NativeWritable} that streams a Vortex file into a Hadoop {@link FileSystem}.
+ *
+ *
Vortex writes and flushes but never closes; the file is finalized when the owner closes this.
+ */
+public final class HadoopWritable implements NativeWritable {
+ private final FSDataOutputStream stream;
+
+ /** Creates (or overwrites) the file at {@code path}, along with any missing parent directories. */
+ public static HadoopWritable create(Configuration conf, String path) {
+ Path file = new Path(path);
+ try {
+ FileSystem fs = file.getFileSystem(conf);
+ return new HadoopWritable(fs.create(file, true));
+ } catch (IOException e) {
+ throw new UncheckedIOException("Failed to create " + path, e);
+ }
+ }
+
+ private HadoopWritable(FSDataOutputStream stream) {
+ this.stream = stream;
+ }
+
+ @Override
+ public void write(byte[] buffer, int offset, int length) throws IOException {
+ stream.write(buffer, offset, length);
+ }
+
+ @Override
+ public void flush() throws IOException {
+ stream.flush();
+ }
+
+ @Override
+ public void close() throws IOException {
+ stream.close();
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/SerializableHadoopConf.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/SerializableHadoopConf.java
new file mode 100644
index 00000000000..c6b9f64fb4f
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/SerializableHadoopConf.java
@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.io;
+
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import org.apache.hadoop.conf.Configuration;
+
+/**
+ * Java-serializable wrapper around a Hadoop {@link Configuration}.
+ *
+ *
{@code org.apache.spark.util.SerializableConfiguration} does the same job but is {@code private[spark]}, so this
+ * class keeps the connector off Spark internals.
+ */
+final class SerializableHadoopConf implements Serializable {
+ private static final long serialVersionUID = 1L;
+
+ private transient Configuration conf;
+
+ SerializableHadoopConf(Configuration conf) {
+ this.conf = conf;
+ }
+
+ Configuration value() {
+ return conf;
+ }
+
+ private void writeObject(ObjectOutputStream out) throws IOException {
+ out.defaultWriteObject();
+ conf.write(out);
+ }
+
+ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+ in.defaultReadObject();
+ conf = new Configuration(false);
+ conf.readFields(in);
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexFile.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexFile.java
new file mode 100644
index 00000000000..9d3b5619010
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexFile.java
@@ -0,0 +1,70 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.io;
+
+import java.io.Serializable;
+import java.util.Locale;
+
+/** A Vortex file and, when a listing already reported it, its size on storage. */
+public final class VortexFile implements Serializable {
+ private static final long serialVersionUID = 1L;
+
+ /** The extension every Vortex data file must carry. */
+ public static final String EXTENSION = ".vortex";
+
+ /** Size of a file nothing has stat'ed yet. Whoever opens it pays for the stat, if it needs one at all. */
+ public static final long UNKNOWN_LENGTH = -1;
+
+ private final String path;
+ private final long length;
+
+ /**
+ * @param path the path as the caller spelled it, not as Hadoop qualifies it
+ * @param length size in bytes, or {@link #UNKNOWN_LENGTH} for a path named directly rather than listed
+ */
+ public VortexFile(String path, long length) {
+ this.path = path;
+ this.length = length;
+ }
+
+ public String path() {
+ return path;
+ }
+
+ public long length() {
+ return length;
+ }
+
+ /** A file whose size is not known yet. */
+ public static VortexFile unsized(String path) {
+ return new VortexFile(path, UNKNOWN_LENGTH);
+ }
+
+ /**
+ * Returns whether {@code path} names a Vortex data file.
+ *
+ *
Spark's file index hides names that begin with {@code _} or {@code .}, but it keeps {@code _metadata} and
+ * {@code _common_metadata}, and it keeps every other extension. The connector reads only files that carry
+ * {@link #EXTENSION}, so this is the one test that decides what belongs to a Vortex dataset.
+ */
+ public static boolean hasVortexExtension(String path) {
+ return path.toLowerCase(Locale.ROOT).endsWith(EXTENSION);
+ }
+
+ /**
+ * Fails unless {@code path} names a Vortex data file.
+ *
+ * @throws IllegalArgumentException if the path does not carry {@link #EXTENSION}
+ */
+ public static void requireVortexExtension(String path) {
+ if (!hasVortexExtension(path)) {
+ throw new IllegalArgumentException(String.format(
+ Locale.ROOT,
+ "%s is not a Vortex file: every file in a Vortex dataset must end with %s. "
+ + "Remove the file, or restrict the listing with the pathGlobFilter option.",
+ path,
+ EXTENSION));
+ }
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java
new file mode 100644
index 00000000000..e1f42c21354
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java
@@ -0,0 +1,84 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.io;
+
+import dev.vortex.io.NativeReadable;
+import dev.vortex.spark.VortexOptions;
+import java.io.Serializable;
+import java.util.Locale;
+import java.util.Objects;
+import org.apache.hadoop.conf.Configuration;
+
+/**
+ * The Hadoop configuration this Spark job reaches storage with, and the read settings that go with it.
+ *
+ *
File contents are read and written through Hadoop streams, so the connector sees the same schemes and credential
+ * providers as Spark's file index and commit protocol. Reads open a {@link HadoopReadable}; writes go through
+ * {@link HadoopWritable} on the task path Spark's commit protocol assigns.
+ *
+ *
Built on the driver and shipped to executors, so the configuration travels with it.
+ */
+public final class VortexIo implements Serializable {
+ private static final long serialVersionUID = 1L;
+
+ /**
+ * Option bounding how many concurrent read upcalls the native reader issues against one file. Each in-flight upcall
+ * leases one pooled Hadoop input stream. Zero keeps the native default.
+ */
+ public static final String READ_CONCURRENCY_OPTION = "vortex.readConcurrency";
+
+ private static final int DEFAULT_READ_CONCURRENCY = 0;
+
+ private final SerializableHadoopConf conf;
+ private final int readConcurrency;
+
+ /** Captures {@code hadoopConf} and the read settings found in the format options. */
+ public static VortexIo create(VortexOptions options, Configuration hadoopConf) {
+ Objects.requireNonNull(options, "options");
+ Objects.requireNonNull(hadoopConf, "hadoopConf");
+ return new VortexIo(new SerializableHadoopConf(hadoopConf), parseReadConcurrency(options));
+ }
+
+ /** Hadoop I/O over a default configuration, for callers that have no Spark session to draw one from. */
+ public static VortexIo defaults() {
+ return new VortexIo(new SerializableHadoopConf(new Configuration()), DEFAULT_READ_CONCURRENCY);
+ }
+
+ private VortexIo(SerializableHadoopConf conf, int readConcurrency) {
+ this.conf = conf;
+ this.readConcurrency = readConcurrency;
+ }
+
+ private static int parseReadConcurrency(VortexOptions options) {
+ int parsed = options.getInt(READ_CONCURRENCY_OPTION, DEFAULT_READ_CONCURRENCY);
+ if (parsed < 0) {
+ throw new IllegalArgumentException(
+ String.format(Locale.ROOT, "%s must be >= 0, got %d", READ_CONCURRENCY_OPTION, parsed));
+ }
+ return parsed;
+ }
+
+ public Configuration hadoopConf() {
+ return conf.value();
+ }
+
+ /** Bound on concurrent read upcalls per file; zero keeps the native default. */
+ public int readConcurrency() {
+ return readConcurrency;
+ }
+
+ /**
+ * Opens a byte source for {@code file}, stating it only if the listing that produced it did not report a size. The
+ * caller owns the result and must close it once the scan built on it is done.
+ *
+ *
Every read path in the connector goes through here, so this is where the {@code *.vortex} requirement is
+ * enforced.
+ *
+ * @throws IllegalArgumentException if the path does not carry {@link VortexFile#EXTENSION}
+ */
+ public NativeReadable openReadable(VortexFile file) {
+ VortexFile.requireVortexExtension(file.path());
+ return HadoopReadable.open(hadoopConf(), file.path(), file.length());
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/PartitionColumnVectors.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/PartitionColumnVectors.java
new file mode 100644
index 00000000000..4178e322242
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/PartitionColumnVectors.java
@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read;
+
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.execution.vectorized.ConstantColumnVector;
+import org.apache.spark.sql.types.BinaryType;
+import org.apache.spark.sql.types.BooleanType;
+import org.apache.spark.sql.types.ByteType;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.DateType;
+import org.apache.spark.sql.types.DayTimeIntervalType;
+import org.apache.spark.sql.types.DecimalType;
+import org.apache.spark.sql.types.DoubleType;
+import org.apache.spark.sql.types.FloatType;
+import org.apache.spark.sql.types.IntegerType;
+import org.apache.spark.sql.types.LongType;
+import org.apache.spark.sql.types.ShortType;
+import org.apache.spark.sql.types.StringType;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.types.TimestampNTZType;
+import org.apache.spark.sql.types.TimestampType;
+import org.apache.spark.sql.types.YearMonthIntervalType;
+import org.apache.spark.sql.vectorized.ColumnVector;
+
+/** Creates constant column vectors from Spark's typed partition values. */
+final class PartitionColumnVectors {
+ private PartitionColumnVectors() {}
+
+ static ColumnVector[] create(int rowCount, StructType schema, InternalRow values) {
+ StructField[] fields = schema.fields();
+ ColumnVector[] vectors = new ColumnVector[fields.length];
+ for (int i = 0; i < fields.length; i++) {
+ vectors[i] = create(rowCount, fields[i], values, i);
+ }
+ return vectors;
+ }
+
+ static ConstantColumnVector create(int rowCount, StructField field, InternalRow values, int ordinal) {
+ ConstantColumnVector vector = new ConstantColumnVector(rowCount, field.dataType());
+ if (values.isNullAt(ordinal)) {
+ vector.setNull();
+ return vector;
+ }
+ vector.setNotNull();
+ DataType type = field.dataType();
+ if (type instanceof BooleanType) {
+ vector.setBoolean(values.getBoolean(ordinal));
+ } else if (type instanceof ByteType) {
+ vector.setByte(values.getByte(ordinal));
+ } else if (type instanceof ShortType) {
+ vector.setShort(values.getShort(ordinal));
+ } else if (type instanceof IntegerType || type instanceof DateType || type instanceof YearMonthIntervalType) {
+ vector.setInt(values.getInt(ordinal));
+ } else if (type instanceof LongType
+ || type instanceof TimestampType
+ || type instanceof TimestampNTZType
+ || type instanceof DayTimeIntervalType) {
+ vector.setLong(values.getLong(ordinal));
+ } else if (type instanceof FloatType) {
+ vector.setFloat(values.getFloat(ordinal));
+ } else if (type instanceof DoubleType) {
+ vector.setDouble(values.getDouble(ordinal));
+ } else if (type instanceof StringType) {
+ vector.setUtf8String(values.getUTF8String(ordinal));
+ } else if (type instanceof BinaryType) {
+ vector.setBinary(values.getBinary(ordinal));
+ } else if (type instanceof DecimalType decimalType) {
+ vector.setDecimal(
+ values.getDecimal(ordinal, decimalType.precision(), decimalType.scale()), decimalType.precision());
+ } else {
+ vector.close();
+ throw new UnsupportedOperationException("Unsupported partition column type: " + type);
+ }
+ return vector;
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java
new file mode 100644
index 00000000000..4f994b92fc5
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java
@@ -0,0 +1,401 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read;
+
+import dev.vortex.api.Expression;
+import dev.vortex.api.Expression.BinaryOp;
+import dev.vortex.api.Expression.TimeUnit;
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.sql.Date;
+import java.sql.Timestamp;
+import java.time.Instant;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.OffsetDateTime;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import org.apache.spark.sql.sources.AlwaysFalse;
+import org.apache.spark.sql.sources.AlwaysTrue;
+import org.apache.spark.sql.sources.And;
+import org.apache.spark.sql.sources.EqualNullSafe;
+import org.apache.spark.sql.sources.EqualTo;
+import org.apache.spark.sql.sources.Filter;
+import org.apache.spark.sql.sources.GreaterThan;
+import org.apache.spark.sql.sources.GreaterThanOrEqual;
+import org.apache.spark.sql.sources.In;
+import org.apache.spark.sql.sources.IsNotNull;
+import org.apache.spark.sql.sources.IsNull;
+import org.apache.spark.sql.sources.LessThan;
+import org.apache.spark.sql.sources.LessThanOrEqual;
+import org.apache.spark.sql.sources.Not;
+import org.apache.spark.sql.sources.Or;
+import org.apache.spark.sql.sources.StringContains;
+import org.apache.spark.sql.sources.StringEndsWith;
+import org.apache.spark.sql.sources.StringStartsWith;
+import org.apache.spark.sql.types.BinaryType;
+import org.apache.spark.sql.types.BooleanType;
+import org.apache.spark.sql.types.ByteType;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.DateType;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.types.DecimalType;
+import org.apache.spark.sql.types.DoubleType;
+import org.apache.spark.sql.types.FloatType;
+import org.apache.spark.sql.types.IntegerType;
+import org.apache.spark.sql.types.LongType;
+import org.apache.spark.sql.types.ShortType;
+import org.apache.spark.sql.types.StringType;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.types.TimestampNTZType;
+import org.apache.spark.sql.types.TimestampType;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/** Converts Spark's stable V1 file filters into Vortex expressions. */
+public final class SparkFilterToVortexExpression {
+ private SparkFilterToVortexExpression() {}
+
+ /** Returns whether the complete filter can be evaluated by Vortex for {@code dataSchema}. */
+ public static boolean isPushable(Filter filter, StructType dataSchema) {
+ return convert(filter, dataSchema).isPresent();
+ }
+
+ /** Converts a filter, returning empty when its operator, column, or literal is unsupported. */
+ public static Optional convert(Filter filter, StructType dataSchema) {
+ if (filter instanceof AlwaysTrue) {
+ return Optional.of(Expression.literal(true));
+ }
+ if (filter instanceof AlwaysFalse) {
+ return Optional.of(Expression.literal(false));
+ }
+ if (filter instanceof And and) {
+ return combine(and.left(), and.right(), dataSchema, true);
+ }
+ if (filter instanceof Or or) {
+ return combine(or.left(), or.right(), dataSchema, false);
+ }
+ if (filter instanceof Not not) {
+ return convert(not.child(), dataSchema).map(Expression::not);
+ }
+ if (filter instanceof IsNull isNull) {
+ return column(isNull, isNull.attribute(), dataSchema).map(Expression::isNull);
+ }
+ if (filter instanceof IsNotNull isNotNull) {
+ return column(isNotNull, isNotNull.attribute(), dataSchema).map(Expression::isNotNull);
+ }
+ if (filter instanceof EqualTo equal) {
+ return comparison(equal, equal.attribute(), equal.value(), dataSchema, BinaryOp.EQ, false);
+ }
+ if (filter instanceof EqualNullSafe equal) {
+ if (equal.value() == null) {
+ return column(equal, equal.attribute(), dataSchema).map(Expression::isNull);
+ }
+ return comparison(equal, equal.attribute(), equal.value(), dataSchema, BinaryOp.EQ, true);
+ }
+ if (filter instanceof GreaterThan greater) {
+ return comparison(greater, greater.attribute(), greater.value(), dataSchema, BinaryOp.GT, false);
+ }
+ if (filter instanceof GreaterThanOrEqual greater) {
+ return comparison(greater, greater.attribute(), greater.value(), dataSchema, BinaryOp.GTE, false);
+ }
+ if (filter instanceof LessThan less) {
+ return comparison(less, less.attribute(), less.value(), dataSchema, BinaryOp.LT, false);
+ }
+ if (filter instanceof LessThanOrEqual less) {
+ return comparison(less, less.attribute(), less.value(), dataSchema, BinaryOp.LTE, false);
+ }
+ if (filter instanceof In in) {
+ return convertIn(in, dataSchema);
+ }
+ if (filter instanceof StringStartsWith startsWith) {
+ return stringMatch(startsWith, startsWith.attribute(), startsWith.value(), dataSchema, false, true);
+ }
+ if (filter instanceof StringEndsWith endsWith) {
+ return stringMatch(endsWith, endsWith.attribute(), endsWith.value(), dataSchema, true, false);
+ }
+ if (filter instanceof StringContains contains) {
+ return stringMatch(contains, contains.attribute(), contains.value(), dataSchema, true, true);
+ }
+ return Optional.empty();
+ }
+
+ private static Optional combine(Filter left, Filter right, StructType dataSchema, boolean conjunction) {
+ Optional convertedLeft = convert(left, dataSchema);
+ Optional convertedRight = convert(right, dataSchema);
+ if (convertedLeft.isEmpty() || convertedRight.isEmpty()) {
+ return Optional.empty();
+ }
+ return Optional.of(
+ conjunction
+ ? Expression.and(convertedLeft.get(), convertedRight.get())
+ : Expression.or(convertedLeft.get(), convertedRight.get()));
+ }
+
+ private static Optional comparison(
+ Filter filter, String attribute, Object value, StructType dataSchema, BinaryOp op, boolean nullSafe) {
+ Optional column = resolveColumn(filter, attribute, dataSchema);
+ if (column.isEmpty() || (value == null && !nullSafe)) {
+ return Optional.empty();
+ }
+ Optional literal = literal(value, column.get().dataType());
+ return literal.map(expression -> Expression.binary(op, column.get().expression(), expression));
+ }
+
+ private static Optional convertIn(In in, StructType dataSchema) {
+ Optional column = resolveColumn(in, in.attribute(), dataSchema);
+ if (column.isEmpty()) {
+ return Optional.empty();
+ }
+ List comparisons = new ArrayList<>();
+ for (Object value : in.values()) {
+ if (value == null) {
+ continue;
+ }
+ Optional literal = literal(value, column.get().dataType());
+ if (literal.isEmpty()) {
+ return Optional.empty();
+ }
+ comparisons.add(Expression.binary(BinaryOp.EQ, column.get().expression(), literal.get()));
+ }
+ if (comparisons.isEmpty()) {
+ return Optional.of(Expression.literal(false));
+ }
+ if (comparisons.size() == 1) {
+ return Optional.of(comparisons.get(0));
+ }
+ return Optional.of(Expression.or(comparisons.toArray(new Expression[0])));
+ }
+
+ private static Optional stringMatch(
+ Filter filter,
+ String attribute,
+ String value,
+ StructType dataSchema,
+ boolean leadingWildcard,
+ boolean trailingWildcard) {
+ Optional column = resolveColumn(filter, attribute, dataSchema);
+ if (column.isEmpty() || !(column.get().dataType() instanceof StringType)) {
+ return Optional.empty();
+ }
+ return Optional.of(Expression.like(
+ column.get().expression(),
+ Expression.literal(likePattern(value, leadingWildcard, trailingWildcard)),
+ false,
+ false));
+ }
+
+ private static String likePattern(String value, boolean leadingWildcard, boolean trailingWildcard) {
+ StringBuilder pattern = new StringBuilder(value.length() + 2);
+ if (leadingWildcard) {
+ pattern.append('%');
+ }
+ for (int i = 0; i < value.length(); i++) {
+ char character = value.charAt(i);
+ if (character == '%' || character == '_' || character == '\\') {
+ pattern.append('\\');
+ }
+ pattern.append(character);
+ }
+ if (trailingWildcard) {
+ pattern.append('%');
+ }
+ return pattern.toString();
+ }
+
+ private static Optional column(Filter filter, String attribute, StructType schema) {
+ return resolveColumn(filter, attribute, schema).map(ResolvedColumn::expression);
+ }
+
+ private static Optional resolveColumn(Filter filter, String attribute, StructType schema) {
+ String[][] references = filter.v2references();
+ if (references.length != 1 || references[0].length == 0) {
+ return Optional.empty();
+ }
+ String[] parts = references[0];
+ DataType current = schema;
+ for (String part : parts) {
+ if (!(current instanceof StructType struct)) {
+ return Optional.empty();
+ }
+ StructField field = findField(struct, part);
+ if (field == null) {
+ return Optional.empty();
+ }
+ current = field.dataType();
+ }
+ return Optional.of(new ResolvedColumn(Expression.column(parts), current));
+ }
+
+ private static StructField findField(StructType schema, String name) {
+ for (StructField field : schema.fields()) {
+ if (field.name().equals(name)) {
+ return field;
+ }
+ }
+ return null;
+ }
+
+ private static Optional literal(Object value, DataType dataType) {
+ if (dataType instanceof BooleanType) {
+ return value == null
+ ? Optional.of(Expression.nullLiteralBool())
+ : value instanceof Boolean booleanValue
+ ? Optional.of(Expression.literal(booleanValue))
+ : Optional.empty();
+ }
+ if (dataType instanceof ByteType) {
+ return numericLiteral(value, Expression.DType.I8, number -> Expression.literal(number.byteValue()));
+ }
+ if (dataType instanceof ShortType) {
+ return numericLiteral(value, Expression.DType.I16, number -> Expression.literal(number.shortValue()));
+ }
+ if (dataType instanceof IntegerType) {
+ return numericLiteral(value, Expression.DType.I32, number -> Expression.literal(number.intValue()));
+ }
+ if (dataType instanceof LongType) {
+ return numericLiteral(value, Expression.DType.I64, number -> Expression.literal(number.longValue()));
+ }
+ if (dataType instanceof FloatType) {
+ return numericLiteral(value, Expression.DType.F32, number -> Expression.literal(number.floatValue()));
+ }
+ if (dataType instanceof DoubleType) {
+ return numericLiteral(value, Expression.DType.F64, number -> Expression.literal(number.doubleValue()));
+ }
+ if (dataType instanceof StringType) {
+ if (value == null) {
+ return Optional.of(Expression.nullLiteral(Expression.DType.UTF8));
+ }
+ if (value instanceof CharSequence || value instanceof UTF8String) {
+ return Optional.of(Expression.literal(value.toString()));
+ }
+ return Optional.empty();
+ }
+ if (dataType instanceof BinaryType) {
+ if (value == null) {
+ return Optional.of(Expression.nullLiteral(Expression.DType.BINARY));
+ }
+ return value instanceof byte[] bytes ? Optional.of(Expression.literal(bytes)) : Optional.empty();
+ }
+ if (dataType instanceof DateType) {
+ if (value == null) {
+ return Optional.of(Expression.nullLiteralDate(TimeUnit.DAYS));
+ }
+ Optional days = dateDays(value);
+ return days.map(dayCount -> Expression.literalDate(dayCount, TimeUnit.DAYS));
+ }
+ if (dataType instanceof TimestampType) {
+ if (value == null) {
+ return Optional.of(Expression.nullLiteralTimestamp(TimeUnit.MICROSECONDS, "UTC"));
+ }
+ return timestampMicros(value, false)
+ .map(micros -> Expression.literalTimestamp(micros, TimeUnit.MICROSECONDS, "UTC"));
+ }
+ if (dataType instanceof TimestampNTZType) {
+ if (value == null) {
+ return Optional.of(Expression.nullLiteralTimestamp(TimeUnit.MICROSECONDS, null));
+ }
+ return timestampMicros(value, true)
+ .map(micros -> Expression.literalTimestamp(micros, TimeUnit.MICROSECONDS, null));
+ }
+ if (dataType instanceof DecimalType decimalType) {
+ if (value == null) {
+ return Optional.of(Expression.nullLiteralDecimal(decimalType.precision(), decimalType.scale()));
+ }
+ BigDecimal decimal;
+ if (value instanceof BigDecimal bigDecimal) {
+ decimal = bigDecimal;
+ } else if (value instanceof Decimal sparkDecimal) {
+ decimal = sparkDecimal.toJavaBigDecimal();
+ } else {
+ return Optional.empty();
+ }
+ try {
+ BigInteger unscaled = decimal.setScale(decimalType.scale()).unscaledValue();
+ return Optional.of(Expression.literalDecimal(unscaled, decimalType.precision(), decimalType.scale()));
+ } catch (ArithmeticException ignored) {
+ return Optional.empty();
+ }
+ }
+ return Optional.empty();
+ }
+
+ private static Optional numericLiteral(
+ Object value, Expression.DType nullType, NumericExpression factory) {
+ if (value == null) {
+ return Optional.of(Expression.nullLiteral(nullType));
+ }
+ return value instanceof Number number ? Optional.of(factory.create(number)) : Optional.empty();
+ }
+
+ private static Optional dateDays(Object value) {
+ if (value instanceof Date date) {
+ return Optional.of(date.toLocalDate().toEpochDay());
+ }
+ if (value instanceof LocalDate date) {
+ return Optional.of(date.toEpochDay());
+ }
+ if (value instanceof Number number) {
+ return Optional.of(number.longValue());
+ }
+ return Optional.empty();
+ }
+
+ private static Optional timestampMicros(Object value, boolean withoutTimeZone) {
+ if (value instanceof Number number) {
+ return Optional.of(number.longValue());
+ }
+ if (withoutTimeZone && value instanceof LocalDateTime localDateTime) {
+ return instantMicros(localDateTime.toInstant(ZoneOffset.UTC));
+ }
+ if (value instanceof Timestamp timestamp) {
+ return instantMicros(timestamp.toInstant());
+ }
+ if (value instanceof Instant instant) {
+ return instantMicros(instant);
+ }
+ if (value instanceof OffsetDateTime offsetDateTime) {
+ return instantMicros(offsetDateTime.toInstant());
+ }
+ if (value instanceof ZonedDateTime zonedDateTime) {
+ return instantMicros(zonedDateTime.toInstant());
+ }
+ return Optional.empty();
+ }
+
+ private static Optional instantMicros(Instant instant) {
+ try {
+ return Optional.of(Math.addExact(
+ Math.multiplyExact(instant.getEpochSecond(), 1_000_000L), instant.getNano() / 1_000L));
+ } catch (ArithmeticException ignored) {
+ return Optional.empty();
+ }
+ }
+
+ private static final class ResolvedColumn {
+ private final Expression expression;
+ private final DataType dataType;
+
+ private ResolvedColumn(Expression expression, DataType dataType) {
+ this.expression = expression;
+ this.dataType = dataType;
+ }
+
+ private Expression expression() {
+ return expression;
+ }
+
+ private DataType dataType() {
+ return dataType;
+ }
+ }
+
+ @FunctionalInterface
+ private interface NumericExpression {
+ Expression create(Number number);
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java
new file mode 100644
index 00000000000..fedb64d9a18
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java
@@ -0,0 +1,126 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read;
+
+import dev.vortex.spark.VortexOptions;
+import dev.vortex.spark.io.VortexFile;
+import dev.vortex.spark.io.VortexIo;
+import java.io.Serializable;
+import java.util.OptionalLong;
+import org.apache.spark.sql.catalyst.FileSourceOptions;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.expressions.Expression;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.aggregate.Aggregation;
+import org.apache.spark.sql.connector.read.InputPartition;
+import org.apache.spark.sql.connector.read.PartitionReader;
+import org.apache.spark.sql.execution.datasources.PartitionedFile;
+import org.apache.spark.sql.execution.datasources.v2.FilePartitionReaderFactory;
+import org.apache.spark.sql.execution.vectorized.ConstantColumnVector;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/** Produces one footer-backed partial COUNT(*) row per Vortex file. */
+public final class VortexAggregateReaderFactory extends FilePartitionReaderFactory implements Serializable {
+ private static final long serialVersionUID = 1L;
+
+ private final FileSourceOptions fileOptions;
+ private final VortexIo io;
+ private final VortexOptions formatOptions;
+ private final StructType aggregateSchema;
+ private final int[] groupByOrdinals;
+
+ public VortexAggregateReaderFactory(
+ FileSourceOptions fileOptions,
+ VortexIo io,
+ VortexOptions formatOptions,
+ StructType aggregateSchema,
+ StructType partitionSchema,
+ Aggregation aggregation) {
+ this.fileOptions = fileOptions;
+ this.io = io;
+ this.formatOptions = formatOptions;
+ this.aggregateSchema = aggregateSchema;
+ Expression[] groupBy = aggregation.groupByExpressions();
+ this.groupByOrdinals = new int[groupBy.length];
+ for (int i = 0; i < groupBy.length; i++) {
+ if (!(groupBy[i] instanceof NamedReference reference) || reference.fieldNames().length != 1) {
+ throw new IllegalArgumentException("COUNT(*) group-by must reference partition columns");
+ }
+ this.groupByOrdinals[i] = partitionSchema.fieldIndex(reference.fieldNames()[0]);
+ }
+ }
+
+ @Override
+ public FileSourceOptions options() {
+ return fileOptions;
+ }
+
+ @Override
+ public PartitionReader buildReader(PartitionedFile file) {
+ throw new UnsupportedOperationException("row-based aggregate reads are not supported");
+ }
+
+ @Override
+ public PartitionReader buildColumnarReader(PartitionedFile file) {
+ return new PartitionReader<>() {
+ private boolean emitted;
+ private ColumnarBatch batch;
+
+ @Override
+ public boolean next() {
+ if (emitted) {
+ return false;
+ }
+ emitted = true;
+ // An estimate is no answer to COUNT(*), so the footer must state the count exactly.
+ OptionalLong rowCount = VortexFooterReader.exactRowCount(
+ new VortexFile(file.toPath().toString(), file.fileSize()), io, formatOptions);
+ if (rowCount.isEmpty()) {
+ throw new IllegalStateException(String.format(
+ "Vortex footer states no exact row count for %s, so COUNT(*) cannot be answered from it. "
+ + "Set the vortex.aggregatePushdown option to false to count rows instead.",
+ file.toPath()));
+ }
+ StructField[] fields = aggregateSchema.fields();
+ ColumnVector[] vectors = new ColumnVector[fields.length];
+ for (int i = 0; i < groupByOrdinals.length; i++) {
+ vectors[i] =
+ PartitionColumnVectors.create(1, fields[i], file.partitionValues(), groupByOrdinals[i]);
+ }
+ for (int i = groupByOrdinals.length; i < fields.length; i++) {
+ ConstantColumnVector count = new ConstantColumnVector(1, fields[i].dataType());
+ count.setNotNull();
+ count.setLong(rowCount.getAsLong());
+ vectors[i] = count;
+ }
+ batch = new ColumnarBatch(vectors, 1);
+ return true;
+ }
+
+ @Override
+ public ColumnarBatch get() {
+ if (batch == null) {
+ throw new IllegalStateException("no aggregate row loaded; call next() first");
+ }
+ return batch;
+ }
+
+ @Override
+ public void close() {
+ if (batch != null) {
+ batch.close();
+ batch = null;
+ }
+ }
+ };
+ }
+
+ @Override
+ public boolean supportColumnarReads(InputPartition partition) {
+ return true;
+ }
+}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java
similarity index 100%
rename from java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java
rename to java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexFooterReader.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexFooterReader.java
new file mode 100644
index 00000000000..e646b2a7217
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexFooterReader.java
@@ -0,0 +1,335 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read;
+
+import dev.vortex.api.DataSource;
+import dev.vortex.api.DataSource.RowCount;
+import dev.vortex.api.Session;
+import dev.vortex.arrow.ArrowAllocation;
+import dev.vortex.io.NativeReadable;
+import dev.vortex.spark.ArrowUtils;
+import dev.vortex.spark.VortexOptions;
+import dev.vortex.spark.VortexSparkSession;
+import dev.vortex.spark.io.VortexFile;
+import dev.vortex.spark.io.VortexIo;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.OptionalLong;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.spark.sql.types.ArrayType;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.MapType;
+import org.apache.spark.sql.types.Metadata;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+
+/** Reads schema and row-count metadata from Vortex file footers. */
+public final class VortexFooterReader {
+ /** Option bounding how many footers scan statistics may read; {@code 0} removes the bound. */
+ public static final String MAX_FILES_OPTION = "vortex.stats.maxFiles";
+
+ /** Option turning schema merging off, leaving the first file's schema to stand for the dataset. */
+ public static final String MERGE_SCHEMA_OPTION = "vortex.mergeSchema";
+
+ /** Option bounding how many footers are read at once. */
+ public static final String FOOTER_PARALLELISM_OPTION = "vortex.footerParallelism";
+
+ private static final int DEFAULT_FOOTER_PARALLELISM = 8;
+ private static final int DEFAULT_MAX_FILES = 1000;
+
+ private VortexFooterReader() {}
+
+ /**
+ * Infers the data schema of a Vortex dataset, or returns null when the listing holds no files at all.
+ *
+ * Every file's footer is read and the schemas are merged, so a column only some files carry is still part of the
+ * dataset. A field missing from a file is nullable in the result, and the reader fills it with nulls for that
+ * file's rows. Set {@value #MERGE_SCHEMA_OPTION} to false to read one footer and let the first file's schema stand
+ * for the dataset; a dataset of uniform files then costs one footer read instead of one per file.
+ *
+ * @throws IllegalArgumentException if the listing holds files but none of them is a Vortex file, or if two files
+ * give a field types that cannot be merged
+ */
+ public static StructType inferSchema(List files, VortexIo io, VortexOptions options) {
+ List vortexFiles = new ArrayList<>();
+ boolean sawFile = false;
+ for (FileStatus status : files) {
+ if (!status.isFile()) {
+ continue;
+ }
+ sawFile = true;
+ if (VortexFile.hasVortexExtension(status.getPath().getName())) {
+ vortexFiles.add(new VortexFile(status.getPath().toString(), status.getLen()));
+ }
+ }
+
+ if (vortexFiles.isEmpty()) {
+ if (sawFile) {
+ throw new IllegalArgumentException(
+ "No Vortex file found to infer a schema from: every file in a Vortex dataset must end with "
+ + VortexFile.EXTENSION);
+ }
+ return null;
+ }
+
+ if (!options.getBoolean(MERGE_SCHEMA_OPTION, true)) {
+ return withDataSource(vortexFiles.get(0), io, options, VortexFooterReader::sparkSchema);
+ }
+
+ List schemas = mapFooters(vortexFiles, io, options, VortexFooterReader::sparkSchema);
+ StructType merged = schemas.get(0);
+ for (int i = 1; i < schemas.size(); i++) {
+ try {
+ merged = mergeStructs(merged, schemas.get(i));
+ } catch (IllegalArgumentException e) {
+ throw new IllegalArgumentException(
+ String.format(
+ Locale.ROOT,
+ "Cannot merge the schema of %s with the schemas of the files before it: %s. "
+ + "Set the %s option to false to read the dataset with one file's schema.",
+ vortexFiles.get(i).path(),
+ e.getMessage(),
+ MERGE_SCHEMA_OPTION),
+ e);
+ }
+ }
+ return merged;
+ }
+
+ /**
+ * Returns the footer row count for one file, exact or estimated.
+ *
+ * Only for Spark scan statistics, which are an estimate by contract. Anything that must answer a query with this
+ * number needs {@link #exactRowCount}.
+ */
+ public static OptionalLong estimatedRowCount(VortexFile file, VortexIo io, VortexOptions options) {
+ return withDataSource(file, io, options, source -> source.rowCount().asOptional());
+ }
+
+ /**
+ * Returns the footer row count for one file, and only when the footer states it exactly.
+ *
+ *
{@code COUNT(*)} pushdown answers the query from this number instead of reading the file, so an estimate would
+ * be returned to the user as fact.
+ */
+ public static OptionalLong exactRowCount(VortexFile file, VortexIo io, VortexOptions options) {
+ return withDataSource(
+ file,
+ io,
+ options,
+ source -> source.rowCount() instanceof RowCount.Exact exact
+ ? OptionalLong.of(exact.value())
+ : OptionalLong.empty());
+ }
+
+ /**
+ * Sums footer row counts in a bounded pool.
+ *
+ *
Returns empty if any footer has no count at all, or if the dataset holds more files than
+ * {@value #MAX_FILES_OPTION} allows. Each footer costs a read against storage on the driver, so a large dataset
+ * would otherwise pay for the whole listing before the job starts.
+ */
+ public static OptionalLong sumRowCounts(List files, VortexIo io, VortexOptions options) {
+ if (files.isEmpty()) {
+ return OptionalLong.of(0);
+ }
+ int maxFiles = options.getInt(MAX_FILES_OPTION, DEFAULT_MAX_FILES);
+ if (maxFiles > 0 && files.size() > maxFiles) {
+ return OptionalLong.empty();
+ }
+
+ long total = 0;
+ for (OptionalLong count :
+ mapFooters(files, io, options, source -> source.rowCount().asOptional())) {
+ if (count.isEmpty()) {
+ return OptionalLong.empty();
+ }
+ total = Math.addExact(total, count.getAsLong());
+ }
+ return OptionalLong.of(total);
+ }
+
+ /**
+ * Merges two schemas of the same dataset.
+ *
+ * Top-level fields are unioned and keep the order they were first seen in. A field only one side carries is
+ * nullable in the result, because the rows of the other side have no value for it, and the reader fills those rows
+ * with nulls.
+ *
+ *
Below the top level only nullability is merged. A struct that gained a field cannot be merged: the reader
+ * projects a struct column whole, as the file stores it, so it has no way to widen one file's struct to a shape
+ * another file agreed on.
+ *
+ * @throws IllegalArgumentException if a field has types that cannot be merged
+ */
+ static StructType mergeStructs(StructType left, StructType right) {
+ Map rightFields = indexByName(right);
+ Map leftFields = indexByName(left);
+
+ LinkedHashMap merged = new LinkedHashMap<>();
+ for (StructField field : left.fields()) {
+ StructField other = rightFields.get(field.name());
+ merged.put(field.name(), other == null ? asNullable(field) : mergeFields(field, other));
+ }
+ for (StructField field : right.fields()) {
+ if (!leftFields.containsKey(field.name())) {
+ merged.put(field.name(), asNullable(field));
+ }
+ }
+ return new StructType(merged.values().toArray(new StructField[0]));
+ }
+
+ private static StructField mergeFields(StructField left, StructField right) {
+ return new StructField(
+ left.name(),
+ mergeTypes(left.name(), left.dataType(), right.dataType()),
+ left.nullable() || right.nullable(),
+ left.metadata());
+ }
+
+ /** Merges the types of one field. Only nullability differences are reconcilable below the top level. */
+ private static DataType mergeTypes(String field, DataType left, DataType right) {
+ if (left.equals(right)) {
+ return left;
+ }
+ if (left instanceof StructType leftStruct && right instanceof StructType rightStruct) {
+ return mergeNestedStruct(field, leftStruct, rightStruct);
+ }
+ if (left instanceof ArrayType leftArray && right instanceof ArrayType rightArray) {
+ return new ArrayType(
+ mergeTypes(field, leftArray.elementType(), rightArray.elementType()),
+ leftArray.containsNull() || rightArray.containsNull());
+ }
+ if (left instanceof MapType leftMap && right instanceof MapType rightMap) {
+ return new MapType(
+ mergeTypes(field, leftMap.keyType(), rightMap.keyType()),
+ mergeTypes(field, leftMap.valueType(), rightMap.valueType()),
+ leftMap.valueContainsNull() || rightMap.valueContainsNull());
+ }
+ throw new IllegalArgumentException(String.format(
+ Locale.ROOT,
+ "field %s is %s in one file and %s in another",
+ field,
+ left.catalogString(),
+ right.catalogString()));
+ }
+
+ private static StructType mergeNestedStruct(String field, StructType left, StructType right) {
+ if (!Arrays.equals(left.fieldNames(), right.fieldNames())) {
+ throw new IllegalArgumentException(String.format(
+ Locale.ROOT,
+ "nested field %s holds %s in one file and %s in another, and a struct that gained or lost a "
+ + "field cannot be merged",
+ field,
+ left.catalogString(),
+ right.catalogString()));
+ }
+
+ StructField[] fields = new StructField[left.fields().length];
+ for (int i = 0; i < fields.length; i++) {
+ fields[i] = mergeFields(left.fields()[i], right.fields()[i]);
+ }
+ return new StructType(fields);
+ }
+
+ private static StructField asNullable(StructField field) {
+ return field.nullable() ? field : new StructField(field.name(), field.dataType(), true, field.metadata());
+ }
+
+ private static Map indexByName(StructType schema) {
+ Map byName = new LinkedHashMap<>();
+ for (StructField field : schema.fields()) {
+ byName.put(field.name(), field);
+ }
+ return byName;
+ }
+
+ private static StructType sparkSchema(DataSource source) {
+ StructField[] fields = source.arrowSchema(ArrowAllocation.rootAllocator()).getFields().stream()
+ .map(field -> new StructField(
+ field.getName(), ArrowUtils.fromArrowField(field), field.isNullable(), Metadata.empty()))
+ .toArray(StructField[]::new);
+ return new StructType(fields);
+ }
+
+ /** Applies {@code function} to every file's footer in a bounded pool, returning results in listing order. */
+ private static List mapFooters(
+ List files, VortexIo io, VortexOptions options, DataSourceFunction function) {
+ int configured = options.getInt(FOOTER_PARALLELISM_OPTION, DEFAULT_FOOTER_PARALLELISM);
+ if (configured < 1) {
+ throw new IllegalArgumentException(FOOTER_PARALLELISM_OPTION + " must be at least 1, got " + configured);
+ }
+
+ ExecutorService executor = Executors.newFixedThreadPool(Math.min(configured, files.size()));
+ try {
+ List> futures = new ArrayList<>(files.size());
+ for (VortexFile file : files) {
+ futures.add(executor.submit(() -> withDataSource(file, io, options, function)));
+ }
+
+ List results = new ArrayList<>(files.size());
+ for (Future future : futures) {
+ results.add(future.get());
+ }
+ return results;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while reading Vortex footers", e);
+ } catch (ExecutionException e) {
+ Throwable cause = e.getCause();
+ if (cause instanceof RuntimeException runtimeException) {
+ throw runtimeException;
+ }
+ if (cause instanceof Error error) {
+ throw error;
+ }
+ throw new RuntimeException("Failed to read Vortex footers", cause);
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ private static T withDataSource(
+ VortexFile file, VortexIo io, VortexOptions options, DataSourceFunction function) {
+ Session session = VortexSparkSession.get(options);
+ NativeReadable readable = io.openReadable(file);
+
+ T result;
+ try {
+ DataSource source = DataSource.open(session, List.of(readable), io.readConcurrency());
+ result = function.apply(source);
+ } catch (RuntimeException | Error e) {
+ // The footer read already failed. A close failure on top of it is a detail of that
+ // failure, never a replacement for it.
+ try {
+ readable.close();
+ } catch (IOException closeFailure) {
+ e.addSuppressed(closeFailure);
+ }
+ throw e;
+ }
+
+ try {
+ readable.close();
+ } catch (IOException e) {
+ throw new UncheckedIOException("Failed to close footer readable for " + file.path(), e);
+ }
+ return result;
+ }
+
+ @FunctionalInterface
+ private interface DataSourceFunction {
+ T apply(DataSource source);
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java
new file mode 100644
index 00000000000..63b14d04a4c
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java
@@ -0,0 +1,259 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read;
+
+import dev.vortex.api.DataSource;
+import dev.vortex.api.Expression;
+import dev.vortex.api.Partition;
+import dev.vortex.api.Scan;
+import dev.vortex.api.ScanOptions;
+import dev.vortex.api.Session;
+import dev.vortex.arrow.ArrowAllocation;
+import dev.vortex.io.NativeReadable;
+import dev.vortex.relocated.org.apache.arrow.memory.BufferAllocator;
+import dev.vortex.relocated.org.apache.arrow.vector.VectorSchemaRoot;
+import dev.vortex.relocated.org.apache.arrow.vector.ipc.ArrowReader;
+import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Field;
+import dev.vortex.spark.VortexOptions;
+import dev.vortex.spark.VortexSparkSession;
+import dev.vortex.spark.io.VortexFile;
+import dev.vortex.spark.io.VortexIo;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Optional;
+import java.util.Set;
+import org.apache.spark.sql.connector.read.PartitionReader;
+import org.apache.spark.sql.execution.datasources.PartitionedFile;
+import org.apache.spark.sql.execution.vectorized.ConstantColumnVector;
+import org.apache.spark.sql.sources.Filter;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/** Columnar reader over one Spark {@link PartitionedFile}. */
+public final class VortexPartitionReader implements PartitionReader {
+ private final PartitionedFile file;
+ private final StructType readDataSchema;
+ private final StructType readPartitionSchema;
+ private final BufferAllocator allocator;
+
+ /**
+ * For each field of {@code readDataSchema}, its position among the columns this file returns, or {@code -1} for a
+ * field the file does not carry. A merged dataset schema holds every field any file carries, so a file written
+ * before a column was added returns fewer columns than the query asked for.
+ */
+ private final int[] vectorSlots;
+
+ private NativeReadable readable;
+ private Session session;
+ private DataSource dataSource;
+ private Scan scan;
+ private Partition currentPartition;
+ private ArrowReader currentReader;
+ private ColumnarBatch currentBatch;
+ private boolean batchLoaded;
+ private boolean exhausted;
+
+ public VortexPartitionReader(
+ PartitionedFile file,
+ StructType dataSchema,
+ StructType readDataSchema,
+ StructType readPartitionSchema,
+ VortexIo io,
+ VortexOptions formatOptions,
+ Filter[] pushedFilters) {
+ this.file = file;
+ this.readDataSchema = readDataSchema;
+ this.readPartitionSchema = readPartitionSchema;
+ this.allocator = ArrowAllocation.rootAllocator();
+ try {
+ session = VortexSparkSession.get(formatOptions);
+ readable = io.openReadable(new VortexFile(file.toPath().toString(), file.fileSize()));
+ dataSource = DataSource.open(session, List.of(readable), io.readConcurrency());
+
+ Set fileFields = fieldNames(dataSource);
+ List projection = new ArrayList<>(readDataSchema.length());
+ this.vectorSlots = new int[readDataSchema.length()];
+ StructField[] fields = readDataSchema.fields();
+ for (int i = 0; i < fields.length; i++) {
+ if (fileFields.contains(fields[i].name())) {
+ vectorSlots[i] = projection.size();
+ projection.add(fields[i].name());
+ } else if (fields[i].nullable()) {
+ vectorSlots[i] = -1;
+ } else {
+ throw new IllegalArgumentException(String.format(
+ Locale.ROOT,
+ "%s does not carry the non-nullable column %s that the query requires",
+ file.toPath(),
+ fields[i].name()));
+ }
+ }
+
+ var options = ScanOptions.builder();
+ // Always project, an empty read schema included: a query that needs only partition columns or only a
+ // row count must not pull every data column off storage.
+ options.projection(Expression.select(projection.toArray(new String[0]), Expression.root()));
+ // Filters are converted against this file's own columns. Spark evaluates every data filter above the
+ // scan anyway, so one that reads a column this file lacks is dropped rather than pushed.
+ buildFilterExpression(pushedFilters, restrictTo(dataSchema, fileFields))
+ .ifPresent(options::filter);
+ scan = dataSource.scan(options.build());
+ } catch (RuntimeException e) {
+ closeReadableAfterFailure(e);
+ throw e;
+ }
+ }
+
+ private Set fieldNames(DataSource source) {
+ Set names = new HashSet<>();
+ for (Field field : source.arrowSchema(allocator).getFields()) {
+ names.add(field.getName());
+ }
+ return names;
+ }
+
+ private static StructType restrictTo(StructType schema, Set names) {
+ StructField[] present = Arrays.stream(schema.fields())
+ .filter(field -> names.contains(field.name()))
+ .toArray(StructField[]::new);
+ return present.length == schema.length() ? schema : new StructType(present);
+ }
+
+ private static Optional buildFilterExpression(Filter[] filters, StructType dataSchema) {
+ Expression combined = null;
+ if (filters != null) {
+ for (Filter filter : filters) {
+ Optional converted = SparkFilterToVortexExpression.convert(filter, dataSchema);
+ if (converted.isPresent()) {
+ combined = combined == null ? converted.get() : Expression.and(combined, converted.get());
+ }
+ }
+ }
+ return Optional.ofNullable(combined);
+ }
+
+ @Override
+ public boolean next() {
+ closeCurrentBatch();
+ batchLoaded = false;
+ if (exhausted) {
+ return false;
+ }
+ while (true) {
+ if (currentReader != null) {
+ try {
+ if (currentReader.loadNextBatch()) {
+ batchLoaded = true;
+ return true;
+ }
+ } catch (IOException e) {
+ throw failure("load a batch from", e);
+ }
+ closeCurrentReader();
+ }
+ if (!scan.hasNext()) {
+ exhausted = true;
+ return false;
+ }
+ currentPartition = scan.next();
+ currentReader = currentPartition.scanArrow(allocator);
+ }
+ }
+
+ @Override
+ public ColumnarBatch get() {
+ if (!batchLoaded) {
+ throw new IllegalStateException("no batch loaded; call next() first");
+ }
+ batchLoaded = false;
+ VectorSchemaRoot root;
+ try {
+ root = currentReader.getVectorSchemaRoot();
+ } catch (IOException e) {
+ throw failure("read the loaded batch of", e);
+ }
+
+ int rowCount = root.getRowCount();
+ StructField[] fields = readDataSchema.fields();
+ ColumnVector[] dataVectors = new ColumnVector[fields.length];
+ for (int i = 0; i < dataVectors.length; i++) {
+ dataVectors[i] = vectorSlots[i] < 0
+ ? nullColumn(rowCount, fields[i].dataType())
+ : new VortexArrowColumnVector(root.getFieldVectors().get(vectorSlots[i]));
+ }
+ ColumnVector[] partitionVectors =
+ PartitionColumnVectors.create(rowCount, readPartitionSchema, file.partitionValues());
+ ColumnVector[] vectors = Arrays.copyOf(dataVectors, dataVectors.length + partitionVectors.length);
+ System.arraycopy(partitionVectors, 0, vectors, dataVectors.length, partitionVectors.length);
+ currentBatch = new ColumnarBatch(vectors, rowCount);
+ return currentBatch;
+ }
+
+ private static ColumnVector nullColumn(int rowCount, DataType type) {
+ ConstantColumnVector vector = new ConstantColumnVector(rowCount, type);
+ vector.setNull();
+ return vector;
+ }
+
+ @Override
+ public void close() {
+ closeCurrentBatch();
+ closeCurrentReader();
+ scan = null;
+ dataSource = null;
+ session = null;
+ if (readable != null) {
+ try {
+ readable.close();
+ } catch (IOException e) {
+ throw failure("close the readable for", e);
+ } finally {
+ readable = null;
+ }
+ }
+ }
+
+ private RuntimeException failure(String action, IOException cause) {
+ return new UncheckedIOException(action + " " + file.toPath(), cause);
+ }
+
+ private void closeCurrentBatch() {
+ if (currentBatch != null) {
+ currentBatch.close();
+ currentBatch = null;
+ }
+ }
+
+ private void closeCurrentReader() {
+ if (currentReader != null) {
+ try {
+ currentReader.close();
+ } catch (IOException e) {
+ throw failure("close the reader for", e);
+ } finally {
+ currentReader = null;
+ currentPartition = null;
+ }
+ }
+ }
+
+ private void closeReadableAfterFailure(RuntimeException failure) {
+ if (readable != null) {
+ try {
+ readable.close();
+ } catch (IOException e) {
+ failure.addSuppressed(e);
+ }
+ readable = null;
+ }
+ }
+}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java
new file mode 100644
index 00000000000..42bfd9311ae
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java
@@ -0,0 +1,68 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read;
+
+import dev.vortex.spark.VortexOptions;
+import dev.vortex.spark.io.VortexIo;
+import java.io.Serializable;
+import org.apache.spark.sql.catalyst.FileSourceOptions;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.read.InputPartition;
+import org.apache.spark.sql.connector.read.PartitionReader;
+import org.apache.spark.sql.execution.datasources.PartitionedFile;
+import org.apache.spark.sql.execution.datasources.v2.FilePartitionReaderFactory;
+import org.apache.spark.sql.sources.Filter;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/** Produces one Vortex reader for each file selected by Spark's file index. */
+public final class VortexPartitionReaderFactory extends FilePartitionReaderFactory implements Serializable {
+ private static final long serialVersionUID = 1L;
+
+ private final FileSourceOptions fileOptions;
+ private final VortexIo io;
+ private final VortexOptions formatOptions;
+ private final StructType dataSchema;
+ private final StructType readDataSchema;
+ private final StructType readPartitionSchema;
+ private final Filter[] pushedFilters;
+
+ public VortexPartitionReaderFactory(
+ FileSourceOptions fileOptions,
+ VortexIo io,
+ VortexOptions formatOptions,
+ StructType dataSchema,
+ StructType readDataSchema,
+ StructType readPartitionSchema,
+ Filter[] pushedFilters) {
+ this.fileOptions = fileOptions;
+ this.io = io;
+ this.formatOptions = formatOptions;
+ this.dataSchema = dataSchema;
+ this.readDataSchema = readDataSchema;
+ this.readPartitionSchema = readPartitionSchema;
+ this.pushedFilters = pushedFilters == null ? new Filter[0] : pushedFilters.clone();
+ }
+
+ @Override
+ public FileSourceOptions options() {
+ return fileOptions;
+ }
+
+ @Override
+ public PartitionReader buildReader(PartitionedFile file) {
+ throw new UnsupportedOperationException("row-based V2 reads are not supported");
+ }
+
+ @Override
+ public PartitionReader buildColumnarReader(PartitionedFile file) {
+ return new VortexPartitionReader(
+ file, dataSchema, readDataSchema, readPartitionSchema, io, formatOptions, pushedFilters);
+ }
+
+ @Override
+ public boolean supportColumnarReads(InputPartition partition) {
+ return true;
+ }
+}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java
similarity index 82%
rename from java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java
rename to java/vortex-spark/common/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java
index 879db9fdd31..f2dda79d7f6 100644
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java
@@ -41,6 +41,37 @@ public final class SparkToArrowSchema {
private SparkToArrowSchema() {}
+ /** Returns whether Vortex's Spark writer can represent this Spark SQL type. */
+ public static boolean supportsDataType(DataType dataType) {
+ if (dataType instanceof StructType structType) {
+ for (StructField field : structType.fields()) {
+ if (!supportsDataType(field.dataType())) {
+ return false;
+ }
+ }
+ return true;
+ }
+ if (dataType instanceof ArrayType arrayType) {
+ return supportsDataType(arrayType.elementType());
+ }
+ if (dataType instanceof MapType mapType) {
+ return supportsDataType(mapType.keyType()) && supportsDataType(mapType.valueType());
+ }
+ return dataType instanceof BooleanType
+ || dataType instanceof ByteType
+ || dataType instanceof ShortType
+ || dataType instanceof IntegerType
+ || dataType instanceof LongType
+ || dataType instanceof FloatType
+ || dataType instanceof DoubleType
+ || dataType instanceof StringType
+ || dataType instanceof BinaryType
+ || dataType instanceof DateType
+ || dataType instanceof TimestampType
+ || dataType instanceof TimestampNTZType
+ || dataType instanceof DecimalType;
+ }
+
/**
* Converts a Spark StructType schema to an Arrow Schema.
*
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/write/VortexOutputWriter.java
similarity index 65%
rename from java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java
rename to java/vortex-spark/common/src/main/java/dev/vortex/spark/write/VortexOutputWriter.java
index 48f47284550..e7d25d3ea1b 100644
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/write/VortexOutputWriter.java
@@ -5,6 +5,7 @@
import dev.vortex.api.Session;
import dev.vortex.api.VortexWriter;
+import dev.vortex.io.NativeWritable;
import dev.vortex.relocated.org.apache.arrow.c.ArrowArray;
import dev.vortex.relocated.org.apache.arrow.c.ArrowSchema;
import dev.vortex.relocated.org.apache.arrow.c.Data;
@@ -28,18 +29,17 @@
import dev.vortex.relocated.org.apache.arrow.vector.complex.ListVector;
import dev.vortex.relocated.org.apache.arrow.vector.complex.MapVector;
import dev.vortex.relocated.org.apache.arrow.vector.complex.StructVector;
+import dev.vortex.spark.VortexOptions;
import dev.vortex.spark.VortexSparkSession;
import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Paths;
+import java.io.UncheckedIOException;
import java.util.ArrayList;
import java.util.List;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.catalyst.expressions.SpecializedGetters;
import org.apache.spark.sql.catalyst.util.ArrayData;
import org.apache.spark.sql.catalyst.util.MapData;
-import org.apache.spark.sql.connector.write.DataWriter;
-import org.apache.spark.sql.connector.write.WriterCommitMessage;
+import org.apache.spark.sql.execution.datasources.OutputWriter;
import org.apache.spark.sql.types.ArrayType;
import org.apache.spark.sql.types.BinaryType;
import org.apache.spark.sql.types.BooleanType;
@@ -58,7 +58,6 @@
import org.apache.spark.sql.types.StructType;
import org.apache.spark.sql.types.TimestampNTZType;
import org.apache.spark.sql.types.TimestampType;
-import org.apache.spark.sql.util.CaseInsensitiveStringMap;
import org.apache.spark.unsafe.types.UTF8String;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -69,8 +68,14 @@
* This writer converts Spark's internal row format to Arrow vectors and writes them to a Vortex file using the
* Vortex writer API.
*/
-public final class VortexDataWriter implements DataWriter, AutoCloseable {
- private static final Logger logger = LoggerFactory.getLogger(VortexDataWriter.class);
+public final class VortexOutputWriter extends OutputWriter {
+ private static final Logger logger = LoggerFactory.getLogger(VortexOutputWriter.class);
+
+ /** Option sizing the row batch converted to Arrow and handed to Vortex at a time. */
+ public static final String BATCH_SIZE_OPTION = "batch.size";
+
+ /** Option sizing the write batch for Vortex alone, overriding {@value #BATCH_SIZE_OPTION}. */
+ public static final String WRITE_BATCH_SIZE_OPTION = "vortex.write.batch.size";
private static final int DEFAULT_BATCH_SIZE = 2048;
private static final int MIN_BATCH_SIZE = 1;
@@ -78,34 +83,31 @@ public final class VortexDataWriter implements DataWriter, AutoClos
private final String filePath;
private final StructType schema;
- private final CaseInsensitiveStringMap options;
+ private final VortexOptions options;
private final int batchSize;
+ private NativeWritable writable;
private Session session;
private VortexWriter vortexWriter;
private BufferAllocator allocator;
private VectorSchemaRoot vectorSchemaRoot;
private final List batchRows = new ArrayList<>();
- private long recordCount = 0;
- private long bytesWritten = 0;
private boolean closed = false;
/**
- * Creates a new VortexDataWriter.
+ * Creates a writer for the task path assigned by Spark's commit protocol.
*
* @param filePath the path where the Vortex file will be written
* @param schema the schema of the data to write
* @param options additional write options
*/
- VortexDataWriter(String filePath, StructType schema, CaseInsensitiveStringMap options) {
+ public VortexOutputWriter(String filePath, StructType schema, VortexOptions options, NativeWritable writable) {
this.filePath = filePath;
this.schema = schema;
this.options = options;
+ this.writable = writable;
- // Get batch size from options with validation
- // Users can set this with: .option("vortex.write.batch.size", "4096")
- int configuredBatchSize =
- options.getInt("vortex.write.batch.size", options.getInt("batch.size", DEFAULT_BATCH_SIZE));
+ int configuredBatchSize = configuredBatchSize(options);
if (configuredBatchSize < MIN_BATCH_SIZE || configuredBatchSize > MAX_BATCH_SIZE) {
logger.warn(
"Batch size {} is out of valid range [{}, {}], using default: {}",
@@ -125,18 +127,33 @@ public final class VortexDataWriter implements DataWriter, AutoClos
this.allocator = new RootAllocator();
var arrowSchema = SparkToArrowSchema.convert(schema);
- this.session = VortexSparkSession.get(options.asCaseSensitiveMap());
- this.vortexWriter = VortexWriter.builder(session, filePath, arrowSchema, allocator)
- .options(options.asCaseSensitiveMap())
+ this.session = VortexSparkSession.get(options);
+ this.vortexWriter = VortexWriter.builder(session, writable, arrowSchema, allocator)
.build();
this.vectorSchemaRoot = VectorSchemaRoot.create(arrowSchema, allocator);
- logger.debug("Initialized VortexDataWriter for {}", filePath);
+ logger.debug("Initialized VortexOutputWriter for {}", filePath);
} catch (IOException e) {
- logger.error("Failed to initialize VortexDataWriter for {}", filePath, e);
- throw new RuntimeException("Failed to initialize VortexDataWriter", e);
+ closeAfterInitializationFailure(e);
+ throw new UncheckedIOException("Failed to initialize VortexOutputWriter for " + filePath, e);
+ }
+ }
+
+ /**
+ * The batch size this writer was asked for.
+ *
+ * {@value #BATCH_SIZE_OPTION} is the generic name, shared with whatever else a job writes.
+ * {@value #WRITE_BATCH_SIZE_OPTION} overrides it, so a job that sets one batch size across formats can still say
+ * something different for Vortex.
+ *
+ *
Package-private so the precedence between the two can be tested without a native writer.
+ */
+ static int configuredBatchSize(VortexOptions options) {
+ if (options.get(WRITE_BATCH_SIZE_OPTION).isPresent()) {
+ return options.getInt(WRITE_BATCH_SIZE_OPTION, DEFAULT_BATCH_SIZE);
}
+ return options.getInt(BATCH_SIZE_OPTION, DEFAULT_BATCH_SIZE);
}
/**
@@ -145,17 +162,19 @@ public final class VortexDataWriter implements DataWriter, AutoClos
* Rows are batched and converted to Arrow format before writing.
*
* @param row the row to write
- * @throws IOException if writing fails
*/
@Override
- public void write(InternalRow row) throws IOException {
+ public void write(InternalRow row) {
// Add row to current batch
batchRows.add(row.copy());
- recordCount++;
// Write batch if it's full
if (batchRows.size() >= batchSize) {
- writeBatch();
+ try {
+ writeBatch();
+ } catch (IOException e) {
+ throw new UncheckedIOException("Failed to write a Vortex batch to " + filePath, e);
+ }
}
}
@@ -173,17 +192,14 @@ private void writeBatch() throws IOException {
for (int fieldIndex = 0; fieldIndex < fields.length; fieldIndex++) {
FieldVector vector = vectorSchemaRoot.getVector(fieldIndex);
DataType dataType = fields[fieldIndex].dataType();
- boolean nullable = fields[fieldIndex].nullable();
// Populate this vector with data from all rows
for (int rowIndex = 0; rowIndex < batchRows.size(); rowIndex++) {
InternalRow row = batchRows.get(rowIndex);
- if (nullable && row.isNullAt(fieldIndex)) {
- // Set null value
+ if (row.isNullAt(fieldIndex)) {
vector.setNull(rowIndex);
} else {
- // Set actual value based on data type
populateVector(vector, dataType, row, fieldIndex, rowIndex);
}
}
@@ -236,12 +252,9 @@ private void populateVector(
} else if (dataType instanceof TimestampNTZType) {
((TimeStampMicroVector) vector).setSafe(rowIndex, row.getLong(fieldIndex));
} else if (dataType instanceof DecimalType decType) {
- if (decType.precision() <= 38) {
- // Use Decimal type from InternalRow
- java.math.BigDecimal decimal = row.getDecimal(fieldIndex, decType.precision(), decType.scale())
- .toJavaBigDecimal();
- ((DecimalVector) vector).setSafe(rowIndex, decimal);
- }
+ java.math.BigDecimal decimal = row.getDecimal(fieldIndex, decType.precision(), decType.scale())
+ .toJavaBigDecimal();
+ ((DecimalVector) vector).setSafe(rowIndex, decimal);
} else if (dataType instanceof StructType structType) {
populateStructVector(
(StructVector) vector, structType, row.getStruct(fieldIndex, structType.fields().length), rowIndex);
@@ -305,147 +318,84 @@ private void populateStructVector(StructVector vector, StructType dataType, Inte
}
}
- /**
- * Commits the write operation and returns a commit message.
- *
- *
This flushes any remaining rows and closes the Vortex writer.
- *
- * @return a commit message with file information
- * @throws IOException if commit fails
- */
@Override
- public WriterCommitMessage commit() throws IOException {
- if (!closed) {
- IOException exception = null;
-
- try {
- // Write any remaining rows
- if (!batchRows.isEmpty()) {
- writeBatch();
- }
-
- // Finalize the file; the summary carries its physical size
- if (vortexWriter != null) {
- try {
- bytesWritten = vortexWriter.finish().fileSize();
- } finally {
- vortexWriter = null; // Always null out the reference
- }
- }
- } catch (IOException e) {
- exception = e;
- }
-
- // Clean up Arrow resources - always attempt cleanup even if there was an error
- try {
- if (vectorSchemaRoot != null) {
- vectorSchemaRoot.close();
- vectorSchemaRoot = null;
- }
- } catch (Exception e) {
- if (exception == null) {
- exception = new IOException("Failed to close VectorSchemaRoot", e);
- } else {
- exception.addSuppressed(e);
- }
+ public void close() {
+ if (closed) {
+ return;
+ }
+ IOException failure = null;
+ try {
+ if (!batchRows.isEmpty()) {
+ writeBatch();
}
-
- // The Arrow C Data Interface export (Data.exportVectorSchemaRoot) creates structural
- // allocations from this allocator. When writeBatch passes the ArrowArray to Rust,
- // FFI_ArrowArray::from_raw() takes ownership and nullifies the release callback on
- // the Java side. The Rust side calls release asynchronously on its own thread, so
- // small structural allocations may still be outstanding when the allocator is closed.
- // These are reclaimed when the allocator is garbage collected.
- if (allocator != null) {
- try {
- allocator.close();
- } catch (IllegalStateException e) {
- logger.debug("Allocator closed with outstanding FFI allocations: {}", e.getMessage());
- }
- allocator = null;
+ if (vortexWriter != null) {
+ vortexWriter.finish();
}
+ } catch (IOException e) {
+ failure = e;
+ } finally {
+ vortexWriter = null;
+ }
- // Session is the JVM-wide singleton held by VortexSparkSession; we just
- // drop our local handle to it here.
- session = null;
-
- closed = true;
-
- // Throw any exception that occurred during cleanup
- if (exception != null) {
- throw exception;
+ failure = closeArrowResources(failure);
+ if (writable != null) {
+ try {
+ writable.close();
+ } catch (IOException e) {
+ failure = addFailure(failure, e);
+ } finally {
+ writable = null;
}
}
-
- return new VortexWriterCommitMessage(filePath, recordCount, bytesWritten);
+ session = null;
+ closed = true;
+ if (failure != null) {
+ throw new UncheckedIOException("Failed to close Vortex output " + filePath, failure);
+ }
}
- /**
- * Aborts the write operation and cleans up resources.
- *
- *
This deletes any partially written file.
- *
- * @throws IOException if abort fails
- */
@Override
- public void abort() throws IOException {
- if (!closed) {
- // Close resources
- if (vortexWriter != null) {
- try {
- vortexWriter.close();
- } catch (Exception e) {
- // Ignore errors during abort
- } finally {
- vortexWriter = null; // Always null out the reference
- }
- }
+ public String path() {
+ return filePath;
+ }
- if (vectorSchemaRoot != null) {
+ private IOException closeArrowResources(IOException failure) {
+ if (vectorSchemaRoot != null) {
+ try {
vectorSchemaRoot.close();
+ } catch (RuntimeException e) {
+ failure = addFailure(failure, new IOException("Failed to close VectorSchemaRoot", e));
+ } finally {
vectorSchemaRoot = null;
}
-
- if (allocator != null) {
- try {
- allocator.close();
- } catch (IllegalStateException e) {
- logger.debug("Allocator closed with outstanding FFI allocations: {}", e.getMessage());
- }
- allocator = null;
- }
-
- // Session is the JVM-wide singleton held by VortexSparkSession; we just
- // drop our local handle to it here.
- session = null;
-
- // Delete the partial file if it exists
+ }
+ if (allocator != null) {
try {
- Files.deleteIfExists(Paths.get(filePath));
- } catch (IOException e) {
- // Ignore - we're already aborting
+ allocator.close();
+ } catch (IllegalStateException e) {
+ logger.debug("Allocator closed with outstanding FFI allocations: {}", e.getMessage());
+ } finally {
+ allocator = null;
}
+ }
+ return failure;
+ }
- closed = true;
+ private void closeAfterInitializationFailure(IOException failure) {
+ closeArrowResources(failure);
+ try {
+ writable.close();
+ } catch (IOException e) {
+ failure.addSuppressed(e);
}
+ writable = null;
}
- /**
- * Closes the writer and releases resources.
- *
- *
This method ensures resources are cleaned up even if commit() or abort() were not called, making the class
- * safe for use with try-with-resources.
- */
- @Override
- public void close() throws IOException {
- if (!closed) {
- logger.warn("VortexDataWriter.close() called without commit() or abort() - cleaning up");
- try {
- abort();
- } catch (IOException e) {
- logger.error("Error during cleanup in close()", e);
- // Suppress the exception as we're already in close()
- }
+ private static IOException addFailure(IOException failure, IOException additional) {
+ if (failure == null) {
+ return additional;
}
+ failure.addSuppressed(additional);
+ return failure;
}
}
diff --git a/java/vortex-spark/common/src/main/java/dev/vortex/spark/write/VortexOutputWriterFactory.java b/java/vortex-spark/common/src/main/java/dev/vortex/spark/write/VortexOutputWriterFactory.java
new file mode 100644
index 00000000000..705893c594f
--- /dev/null
+++ b/java/vortex-spark/common/src/main/java/dev/vortex/spark/write/VortexOutputWriterFactory.java
@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.write;
+
+import dev.vortex.spark.VortexOptions;
+import dev.vortex.spark.io.HadoopWritable;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+import org.apache.spark.sql.execution.datasources.OutputWriter;
+import org.apache.spark.sql.execution.datasources.OutputWriterFactory;
+import org.apache.spark.sql.types.StructType;
+
+/** Creates Vortex output writers at paths assigned by Spark's file commit protocol. */
+public final class VortexOutputWriterFactory extends OutputWriterFactory {
+ private static final long serialVersionUID = 1L;
+
+ private final StructType schema;
+ private final VortexOptions options;
+
+ public VortexOutputWriterFactory(StructType schema, VortexOptions options) {
+ this.schema = schema;
+ this.options = options;
+ }
+
+ @Override
+ public String getFileExtension(TaskAttemptContext context) {
+ return ".vortex";
+ }
+
+ @Override
+ public OutputWriter newInstance(String path, StructType dataSchema, TaskAttemptContext context) {
+ return new VortexOutputWriter(path, schema, options, HadoopWritable.create(context.getConfiguration(), path));
+ }
+}
diff --git a/java/vortex-spark/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister b/java/vortex-spark/common/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister
similarity index 100%
rename from java/vortex-spark/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister
rename to java/vortex-spark/common/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister
diff --git a/java/vortex-spark/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister.license b/java/vortex-spark/common/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister.license
similarity index 100%
rename from java/vortex-spark/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister.license
rename to java/vortex-spark/common/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister.license
diff --git a/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexDataSourceV2.scala b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexDataSourceV2.scala
new file mode 100644
index 00000000000..2637450bd96
--- /dev/null
+++ b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexDataSourceV2.scala
@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark
+
+import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.execution.datasources.FileFormat
+import org.apache.spark.sql.execution.datasources.v2.FileDataSourceV2
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/** Spark file data source for Vortex files. */
+final class VortexDataSourceV2 extends FileDataSourceV2 {
+ override def fallbackFileFormat: Class[_ <: FileFormat] =
+ classOf[VortexFileFormat]
+
+ override def shortName(): String = "vortex"
+
+ override protected def getTable(options: CaseInsensitiveStringMap): Table = {
+ val paths = getPaths(options)
+ new VortexTable(
+ getTableName(options, paths),
+ sparkSession,
+ getOptionsWithoutPaths(options),
+ paths,
+ None,
+ fallbackFileFormat
+ )
+ }
+
+ override protected def getTable(
+ options: CaseInsensitiveStringMap,
+ schema: StructType
+ ): Table = {
+ val paths = getPaths(options)
+ new VortexTable(
+ getTableName(options, paths),
+ sparkSession,
+ getOptionsWithoutPaths(options),
+ paths,
+ Some(schema),
+ fallbackFileFormat
+ )
+ }
+}
diff --git a/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexFileFormat.scala b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexFileFormat.scala
new file mode 100644
index 00000000000..c1427f48a78
--- /dev/null
+++ b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexFileFormat.scala
@@ -0,0 +1,220 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark
+
+import scala.collection.JavaConverters._
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileStatus, Path}
+import org.apache.hadoop.mapreduce.Job
+
+import org.apache.spark.TaskContext
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.execution.vectorized.ConstantColumnVector
+import org.apache.spark.sql.execution.datasources.{
+ FileFormat,
+ OutputWriterFactory,
+ PartitionedFile
+}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.sources.{DataSourceRegister, Filter}
+import org.apache.spark.sql.types.{DataType, StructType}
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import dev.vortex.spark.VortexOptions
+import dev.vortex.spark.io.VortexIo
+import dev.vortex.spark.read.{
+ VortexArrowColumnVector,
+ VortexFooterReader,
+ VortexPartitionReader
+}
+import dev.vortex.spark.write.{SparkToArrowSchema, VortexOutputWriterFactory}
+
+/** V1 fallback used for catalog tables and when Vortex is listed in
+ * useV1SourceList.
+ */
+final class VortexFileFormat
+ extends FileFormat
+ with DataSourceRegister
+ with Serializable {
+ override def shortName(): String = "vortex"
+
+ override def toString: String = "Vortex"
+
+ override def equals(other: Any): Boolean =
+ other.isInstanceOf[VortexFileFormat]
+
+ override def hashCode(): Int = getClass.hashCode()
+
+ override def inferSchema(
+ spark: SparkSession,
+ options: Map[String, String],
+ files: Seq[FileStatus]
+ ): Option[StructType] = {
+ val hadoopConf = spark.sessionState.newHadoopConfWithOptions(options)
+ val vortexOptions = VortexOptions.of(options.asJava)
+ Option(
+ VortexFooterReader.inferSchema(
+ files.asJava,
+ VortexIo.create(vortexOptions, hadoopConf),
+ vortexOptions
+ )
+ )
+ }
+
+ override def isSplitable(
+ sparkSession: SparkSession,
+ options: Map[String, String],
+ path: Path
+ ): Boolean = false
+
+ override def supportBatch(
+ sparkSession: SparkSession,
+ dataSchema: StructType
+ ): Boolean = true
+
+ override def vectorTypes(
+ requiredSchema: StructType,
+ partitionSchema: StructType,
+ sqlConf: SQLConf
+ ): Option[Seq[String]] =
+ Some(
+ Seq.fill(requiredSchema.length)(
+ classOf[VortexArrowColumnVector].getName
+ ) ++
+ Seq.fill(partitionSchema.length)(classOf[ConstantColumnVector].getName)
+ )
+
+ override def prepareWrite(
+ spark: SparkSession,
+ job: Job,
+ options: Map[String, String],
+ dataSchema: StructType
+ ): OutputWriterFactory =
+ new VortexOutputWriterFactory(dataSchema, VortexOptions.of(options.asJava))
+
+ override def buildReader(
+ spark: SparkSession,
+ dataSchema: StructType,
+ partitionSchema: StructType,
+ requiredSchema: StructType,
+ filters: Seq[Filter],
+ options: Map[String, String],
+ hadoopConf: Configuration
+ ): PartitionedFile => Iterator[InternalRow] = {
+ val optionMap = VortexOptions.of(options.asJava)
+ val io = VortexIo.create(optionMap, hadoopConf)
+ (file: PartitionedFile) => {
+ val reader = new VortexPartitionReader(
+ file,
+ dataSchema,
+ requiredSchema,
+ new StructType(),
+ io,
+ optionMap,
+ filters.toArray
+ )
+ Option(TaskContext.get())
+ .foreach(_.addTaskCompletionListener[Unit](_ => reader.close()))
+ // Reached whenever Spark asks for rows rather than batches, which `FileSourceScanExec` does when
+ // whole-stage codegen is off or the schema has more fields than `spark.sql.codegen.maxFields`.
+ new Iterator[InternalRow] {
+ private var batch = Option.empty[ColumnarBatch]
+ private var rowIndex = 0
+
+ override def hasNext: Boolean = {
+ // The reader owns each batch: `next()` releases the one before it, and `close()` releases the last.
+ while (batch.forall(rowIndex >= _.numRows()) && reader.next()) {
+ batch = Some(reader.get())
+ rowIndex = 0
+ }
+ if (batch.exists(rowIndex < _.numRows())) {
+ true
+ } else {
+ batch = None
+ reader.close()
+ false
+ }
+ }
+
+ override def next(): InternalRow = {
+ if (!hasNext) {
+ throw new NoSuchElementException("end of Vortex file")
+ }
+ val row = batch.get.getRow(rowIndex)
+ rowIndex += 1
+ row
+ }
+ }
+ }
+ }
+
+ override def buildReaderWithPartitionValues(
+ spark: SparkSession,
+ dataSchema: StructType,
+ partitionSchema: StructType,
+ requiredSchema: StructType,
+ filters: Seq[Filter],
+ options: Map[String, String],
+ hadoopConf: Configuration
+ ): PartitionedFile => Iterator[InternalRow] = {
+ if (
+ !options.getOrElse(FileFormat.OPTION_RETURNING_BATCH, "false").toBoolean
+ ) {
+ return super.buildReaderWithPartitionValues(
+ spark,
+ dataSchema,
+ partitionSchema,
+ requiredSchema,
+ filters,
+ options,
+ hadoopConf
+ )
+ }
+
+ val optionMap = VortexOptions.of(options.asJava)
+ val io = VortexIo.create(optionMap, hadoopConf)
+ (file: PartitionedFile) => {
+ val reader = new VortexPartitionReader(
+ file,
+ dataSchema,
+ requiredSchema,
+ partitionSchema,
+ io,
+ optionMap,
+ filters.toArray
+ )
+ Option(TaskContext.get()).foreach(
+ _.addTaskCompletionListener[Unit](_ => reader.close())
+ )
+ new Iterator[ColumnarBatch] {
+ private var loaded = false
+ private var exhausted = false
+
+ override def hasNext: Boolean = {
+ if (!loaded && !exhausted) {
+ loaded = reader.next()
+ if (!loaded) {
+ exhausted = true
+ reader.close()
+ }
+ }
+ loaded
+ }
+
+ override def next(): ColumnarBatch = {
+ if (!hasNext) {
+ throw new NoSuchElementException("end of Vortex file")
+ }
+ loaded = false
+ reader.get()
+ }
+ }.asInstanceOf[Iterator[InternalRow]]
+ }
+ }
+
+ override def supportDataType(dataType: DataType): Boolean =
+ SparkToArrowSchema.supportsDataType(dataType)
+}
diff --git a/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexTable.scala b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexTable.scala
new file mode 100644
index 00000000000..8e69c9f548b
--- /dev/null
+++ b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexTable.scala
@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark
+
+import scala.collection.JavaConverters._
+
+import org.apache.hadoop.fs.FileStatus
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.connector.write.{
+ LogicalWriteInfo,
+ Write,
+ WriteBuilder
+}
+import org.apache.spark.sql.execution.datasources.FileFormat
+import org.apache.spark.sql.execution.datasources.v2.FileTable
+import org.apache.spark.sql.types.{DataType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+import dev.vortex.spark.VortexOptions
+import dev.vortex.spark.io.VortexIo
+import dev.vortex.spark.read.{VortexFooterReader, VortexScanBuilder}
+import dev.vortex.spark.write.{SparkToArrowSchema, VortexWrite}
+
+/** A Vortex table backed by Spark's file index and commit protocol. */
+final class VortexTable(
+ tableName: String,
+ sparkSession: SparkSession,
+ options: CaseInsensitiveStringMap,
+ paths: Seq[String],
+ userSpecifiedSchema: Option[StructType],
+ override val fallbackFileFormat: Class[_ <: FileFormat]
+) extends FileTable(sparkSession, options, paths, userSpecifiedSchema) {
+
+ override def name(): String = tableName
+
+ override def newScanBuilder(
+ operationOptions: CaseInsensitiveStringMap
+ ): VortexScanBuilder =
+ new VortexScanBuilder(
+ sparkSession,
+ fileIndex,
+ schema,
+ dataSchema,
+ VortexOptions.of(mergeOptions(operationOptions).asCaseSensitiveMap)
+ )
+
+ override def inferSchema(files: Seq[FileStatus]): Option[StructType] = {
+ val vortexOptions = VortexOptions.of(options.asCaseSensitiveMap)
+ val hadoopConf = sparkSession.sessionState.newHadoopConfWithOptions(
+ vortexOptions.asCaseSensitiveMap.asScala.toMap
+ )
+ Option(
+ VortexFooterReader.inferSchema(
+ files.asJava,
+ VortexIo.create(vortexOptions, hadoopConf),
+ vortexOptions
+ )
+ )
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
+ val mergedInfo = new LogicalWriteInfo {
+ override def queryId(): String = info.queryId()
+ override def schema(): StructType = info.schema()
+ override def options(): CaseInsensitiveStringMap = mergeOptions(
+ info.options()
+ )
+ }
+ new WriteBuilder {
+ override def build(): Write =
+ VortexWrite(paths, formatName, supportsDataType, mergedInfo)
+ }
+ }
+
+ override def supportsDataType(dataType: DataType): Boolean =
+ SparkToArrowSchema.supportsDataType(dataType)
+
+ override def formatName: String = "VORTEX"
+
+ private def mergeOptions(
+ operationOptions: CaseInsensitiveStringMap
+ ): CaseInsensitiveStringMap = {
+ val merged =
+ options.asCaseSensitiveMap.asScala ++ operationOptions.asCaseSensitiveMap.asScala
+ new CaseInsensitiveStringMap(merged.asJava)
+ }
+}
diff --git a/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala
new file mode 100644
index 00000000000..bc8df2f85d8
--- /dev/null
+++ b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala
@@ -0,0 +1,139 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read
+
+import java.util.OptionalLong
+
+import scala.collection.JavaConverters._
+
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.FileSourceOptions
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
+import org.apache.spark.sql.connector.read.{PartitionReaderFactory, Statistics}
+import org.apache.spark.sql.execution.datasources.{
+ AggregatePushDownUtils,
+ PartitioningAwareFileIndex
+}
+import org.apache.spark.sql.execution.datasources.v2.FileScan
+import org.apache.spark.sql.sources.Filter
+import org.apache.spark.sql.types.StructType
+
+import dev.vortex.spark.VortexOptions
+import dev.vortex.spark.io.{VortexFile, VortexIo}
+
+object VortexScan {
+
+ /** Option turning footer-backed row-count statistics off. */
+ val RowCountStatisticsOption = "vortex.stats.rowCount"
+}
+
+/** A Vortex scan planned by Spark's file source framework. */
+final class VortexScan(
+ override val sparkSession: SparkSession,
+ override val fileIndex: PartitioningAwareFileIndex,
+ override val dataSchema: StructType,
+ override val readDataSchema: StructType,
+ override val readPartitionSchema: StructType,
+ val options: VortexOptions,
+ val pushedFilters: Array[Filter],
+ val pushedAggregation: Option[Aggregation],
+ override val partitionFilters: Seq[Expression],
+ override val dataFilters: Seq[Expression]
+) extends FileScan {
+
+ private val caseSensitiveOptions = options.asCaseSensitiveMap.asScala.toMap
+ private val hadoopConf =
+ sparkSession.sessionState.newHadoopConfWithOptions(caseSensitiveOptions)
+ private val io = VortexIo.create(options, hadoopConf)
+
+ override def isSplitable(path: Path): Boolean = false
+
+ override def readSchema(): StructType =
+ if (pushedAggregation.nonEmpty) readDataSchema else super.readSchema()
+
+ override def createReaderFactory(): PartitionReaderFactory = {
+ val fileOptions = new FileSourceOptions(caseSensitiveOptions)
+ pushedAggregation match {
+ case Some(aggregation) =>
+ new VortexAggregateReaderFactory(
+ fileOptions,
+ io,
+ options,
+ readDataSchema,
+ readPartitionSchema,
+ aggregation
+ )
+ case None =>
+ new VortexPartitionReaderFactory(
+ fileOptions,
+ io,
+ options,
+ dataSchema,
+ readDataSchema,
+ readPartitionSchema,
+ pushedFilters
+ )
+ }
+ }
+
+ override def estimateStatistics(): Statistics = cachedStatistics
+
+ private lazy val cachedStatistics: Statistics = {
+ val base = super.estimateStatistics()
+ val rows = pushedAggregation match {
+ // The aggregate reader answers COUNT(*) from the same footers, and emits one row per file. Summing them
+ // here as well would pay for every footer twice over -- once on the driver, once on the executors -- and
+ // would describe rows the scan does not emit.
+ case Some(_) => OptionalLong.of(partitions.map(_.files.length.toLong).sum)
+ case None
+ if options.getBoolean(VortexScan.RowCountStatisticsOption, true) =>
+ // Count over the same file set the scan will read, so statistics and execution cannot
+ // disagree about what belongs to the dataset.
+ val files = partitions
+ .flatMap(_.files)
+ .map(file => new VortexFile(file.toPath.toString, file.fileSize))
+ VortexFooterReader.sumRowCounts(files.asJava, io, options)
+ case None => OptionalLong.empty()
+ }
+ new Statistics {
+ override def sizeInBytes(): OptionalLong = base.sizeInBytes()
+ override def numRows(): OptionalLong = rows
+ }
+ }
+
+ override def equals(other: Any): Boolean = other match {
+ case scan: VortexScan =>
+ val aggregationsEqual =
+ (pushedAggregation, scan.pushedAggregation) match {
+ case (Some(left), Some(right)) =>
+ AggregatePushDownUtils.equivalentAggregations(left, right)
+ case (None, None) => true
+ case _ => false
+ }
+ super.equals(
+ scan
+ ) && dataSchema == scan.dataSchema && options == scan.options &&
+ equivalentFilters(pushedFilters, scan.pushedFilters) && aggregationsEqual
+ case _ => false
+ }
+
+ override def hashCode(): Int = getClass.hashCode()
+
+ override def getMetaData(): Map[String, String] = {
+ val aggregation = pushedAggregation
+ .map(value => seqToString(value.aggregateExpressions().toSeq))
+ .getOrElse("[]")
+ val groupBy = pushedAggregation
+ .map(value => seqToString(value.groupByExpressions().toSeq))
+ .getOrElse("[]")
+ super.getMetaData() ++ Map(
+ "PushedFilters" -> seqToString(pushedFilters.toSeq),
+ "PushedAggregation" -> aggregation,
+ "PushedGroupBy" -> groupBy
+ )
+ }
+}
diff --git a/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScanBuilder.scala b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScanBuilder.scala
new file mode 100644
index 00000000000..70200a6ec34
--- /dev/null
+++ b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScanBuilder.scala
@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.connector.expressions.aggregate.{
+ Aggregation,
+ CountStar
+}
+import org.apache.spark.sql.connector.read.SupportsPushDownAggregates
+import org.apache.spark.sql.execution.datasources.{
+ AggregatePushDownUtils,
+ PartitioningAwareFileIndex
+}
+import org.apache.spark.sql.execution.datasources.v2.FileScanBuilder
+import org.apache.spark.sql.sources.Filter
+import org.apache.spark.sql.types.StructType
+
+import dev.vortex.spark.VortexOptions
+
+object VortexScanBuilder {
+
+ /** Option turning COUNT(*) pushdown off. */
+ val AggregatePushdownOption = "vortex.aggregatePushdown"
+}
+
+/** Builds Vortex file scans with projection, filter, and COUNT(*) pushdown. */
+final class VortexScanBuilder(
+ sparkSession: SparkSession,
+ fileIndex: PartitioningAwareFileIndex,
+ schema: StructType,
+ dataSchema: StructType,
+ options: VortexOptions
+) extends FileScanBuilder(sparkSession, fileIndex, dataSchema)
+ with SupportsPushDownAggregates {
+
+ private var finalSchema = new StructType()
+ private var pushedAggregation = Option.empty[Aggregation]
+
+ override protected def pushDataFilters(
+ filters: Array[Filter]
+ ): Array[Filter] =
+ filters.filter(SparkFilterToVortexExpression.isPushable(_, dataSchema))
+
+ override def pushAggregation(aggregation: Aggregation): Boolean = {
+ if (!aggregatePushdownEnabled) {
+ return false
+ }
+ pushedAggregationSchema(aggregation) match {
+ case Some(aggregateSchema) =>
+ finalSchema = aggregateSchema
+ pushedAggregation = Some(aggregation)
+ true
+ case None => false
+ }
+ }
+
+ // The reader emits one partial aggregate per file. Spark must combine files that share a group.
+ override def supportCompletePushDown(aggregation: Aggregation): Boolean =
+ false
+
+ override def build(): VortexScan = {
+ if (pushedAggregation.isEmpty) {
+ finalSchema = readDataSchema()
+ }
+ new VortexScan(
+ sparkSession,
+ fileIndex,
+ dataSchema,
+ finalSchema,
+ readPartitionSchema(),
+ options,
+ pushedDataFilters,
+ pushedAggregation,
+ partitionFilters,
+ dataFilters
+ )
+ }
+
+ /** Escape hatch matching `spark.sql.parquet.aggregatePushdown`, for when a
+ * footer count must not be trusted.
+ */
+ private def aggregatePushdownEnabled: Boolean =
+ options.getBoolean(VortexScanBuilder.AggregatePushdownOption, true)
+
+ private def pushedAggregationSchema(
+ aggregation: Aggregation
+ ): Option[StructType] = {
+ if (
+ pushedDataFilters.nonEmpty ||
+ !aggregation.aggregateExpressions().forall(_.isInstanceOf[CountStar])
+ ) {
+ None
+ } else {
+ AggregatePushDownUtils.getSchemaForPushedAggregation(
+ aggregation,
+ schema,
+ partitionNameSet,
+ dataFilters
+ )
+ }
+ }
+}
diff --git a/java/vortex-spark/common/src/main/scala/dev/vortex/spark/write/VortexWrite.scala b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/write/VortexWrite.scala
new file mode 100644
index 00000000000..21dd4b572ee
--- /dev/null
+++ b/java/vortex-spark/common/src/main/scala/dev/vortex/spark/write/VortexWrite.scala
@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.write
+
+import scala.collection.JavaConverters._
+
+import org.apache.hadoop.mapreduce.Job
+
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.execution.datasources.OutputWriterFactory
+import org.apache.spark.sql.execution.datasources.v2.FileWrite
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{DataType, StructType}
+
+import dev.vortex.spark.VortexOptions
+
+/** Vortex write support using Spark's file commit protocol. */
+final case class VortexWrite(
+ paths: Seq[String],
+ formatName: String,
+ supportsDataType: DataType => Boolean,
+ info: LogicalWriteInfo
+) extends FileWrite {
+
+ override def prepareWrite(
+ sqlConf: SQLConf,
+ job: Job,
+ options: Map[String, String],
+ dataSchema: StructType
+ ): OutputWriterFactory =
+ new VortexOutputWriterFactory(dataSchema, VortexOptions.of(options.asJava))
+}
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/ArrowUtilsTest.java
similarity index 100%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/ArrowUtilsTest.java
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexAggregatePushdownTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexAggregatePushdownTest.java
new file mode 100644
index 00000000000..4de243dadf4
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexAggregatePushdownTest.java
@@ -0,0 +1,145 @@
+// 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 static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Path;
+import java.util.List;
+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.catalyst.plans.logical.Statistics;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.io.TempDir;
+
+/** COUNT(*) footer pushdown behavior and rejection cases. */
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public final class VortexAggregatePushdownTest {
+ private SparkSession spark;
+
+ @TempDir
+ Path tempDir;
+
+ @BeforeAll
+ public void setUp() {
+ spark = SparkSession.builder()
+ .appName("VortexAggregatePushdownTest")
+ .master("local[2]")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.sql.adaptive.enabled", "false")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ }
+
+ @AfterAll
+ public void tearDown() {
+ if (spark != null) {
+ spark.stop();
+ }
+ }
+
+ @Test
+ void countStarPushesDownAcrossMultipleFiles() {
+ Dataset data = writeAndRead(321, "count", 5, false);
+ Dataset count = data.selectExpr("count(*) AS count");
+
+ assertEquals(321L, count.first().getLong(0));
+ String plan = count.queryExecution().executedPlan().toString();
+ assertTrue(plan.contains("PushedAggregation: [COUNT(*)]"), plan);
+ }
+
+ @Test
+ void filteredCountAndCountColumnAreRejected() {
+ Dataset data = writeAndRead(30, "rejected", 3, false);
+ Dataset filtered = data.filter("id > 10").selectExpr("count(*) AS count");
+ Dataset countColumn = data.selectExpr("count(value) AS count");
+
+ assertEquals(19L, filtered.first().getLong(0));
+ assertFalse(filtered.queryExecution().executedPlan().toString().contains("PushedAggregation: [count(*)]"));
+ assertEquals(30L, countColumn.first().getLong(0));
+ assertFalse(
+ countColumn.queryExecution().executedPlan().toString().contains("PushedAggregation: [count(value)]"));
+ }
+
+ @Test
+ void countStarGroupsByHivePartitionValues() {
+ Dataset data = writeAndRead(12, "grouped", 2, true);
+
+ List rows = data.groupBy("group").count().orderBy("group").collectAsList();
+
+ assertEquals(2, rows.size());
+ assertEquals(0, rows.get(0).getInt(0));
+ assertEquals(6L, rows.get(0).getLong(1));
+ assertEquals(1, rows.get(1).getInt(0));
+ assertEquals(6L, rows.get(1).getLong(1));
+ }
+
+ @Test
+ void countStarPushdownCanBeDisabled() {
+ Path output = tempDir.resolve("pushdown_off");
+ spark.range(0, 47)
+ .selectExpr("cast(id as int) as id")
+ .repartition(3)
+ .write()
+ .format("vortex")
+ .mode(SaveMode.Overwrite)
+ .save(output.toString());
+
+ Dataset count = spark.read()
+ .format("vortex")
+ .option("vortex.aggregatePushdown", "false")
+ .load(output.toString())
+ .selectExpr("count(*) AS count");
+
+ assertEquals(47L, count.first().getLong(0));
+ String plan = count.queryExecution().executedPlan().toString();
+ assertFalse(plan.contains("PushedAggregation: [COUNT(*)]"), plan);
+ }
+
+ @Test
+ void aPushedCountReportsTheRowsItEmitsAndReadsNoFooterTwice() {
+ Dataset data = writeAndRead(400, "pushed_stats", 4, false);
+ Dataset count = data.selectExpr("count(*) AS count");
+
+ Statistics statistics =
+ count.queryExecution().optimizedPlan().collectLeaves().head().stats();
+
+ // The reader answers from one footer per file and emits one row for each. Summing those footers here as
+ // well would pay for every footer twice and describe rows the scan never emits.
+ assertTrue(statistics.rowCount().isDefined());
+ assertEquals(4L, statistics.rowCount().get().longValue());
+ assertEquals(400L, count.first().getLong(0));
+ }
+
+ private Dataset writeAndRead(int count, String name, int partitions, boolean partitioned) {
+ Path output = tempDir.resolve(name);
+ Dataset data = spark.range(0, count)
+ .selectExpr(
+ "cast(id as int) as id",
+ "concat('value_', cast(id as string)) as value",
+ "cast(id % 2 as int) as group");
+ if (partitioned) {
+ data.repartition(partitions)
+ .write()
+ .format("vortex")
+ .partitionBy("group")
+ .mode(SaveMode.Overwrite)
+ .save(output.toString());
+ } else {
+ data.repartition(partitions)
+ .write()
+ .format("vortex")
+ .mode(SaveMode.Overwrite)
+ .save(output.toString());
+ }
+ return spark.read().format("vortex").load(output.toString());
+ }
+}
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceBasicTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceBasicTest.java
similarity index 88%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceBasicTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceBasicTest.java
index f22246d43f2..2c9482319ff 100644
--- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceBasicTest.java
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceBasicTest.java
@@ -9,7 +9,6 @@
import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType;
import dev.vortex.spark.write.SparkToArrowSchema;
-import dev.vortex.spark.write.VortexWriterCommitMessage;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
@@ -92,18 +91,4 @@ public void testNestedSparkToArrowSchemaConversion() {
assertEquals("active", nestedFields.get(3).getName());
assertInstanceOf(ArrowType.Bool.class, nestedFields.get(3).getType());
}
-
- @Test
- @DisplayName("VortexWriterCommitMessage should store metadata correctly")
- public void testWriterCommitMessage() {
- String testPath = "/test/path/file.vortex";
- long recordCount = 1000;
- long bytesWritten = 50000;
-
- var message = new VortexWriterCommitMessage(testPath, recordCount, bytesWritten);
-
- assertEquals(testPath, message.filePath());
- assertEquals(recordCount, message.recordCount());
- assertEquals(bytesWritten, message.bytesWritten());
- }
}
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java
similarity index 84%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java
index 8c103cea9b9..349760cc10f 100644
--- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java
@@ -42,18 +42,16 @@ public void tearDown() {
}
@Test
- @DisplayName("Reading a directory without Vortex files reports the offending path")
- public void inferSchemaFailureNamesThePath() {
+ @DisplayName("Reading a directory without Vortex files reports schema inference failure")
+ public void inferSchemaFailureUsesSparkFileSourceError() {
String emptyDir = tempDir.toString();
Throwable thrown = assertThrows(
Throwable.class, () -> spark.read().format("vortex").load(emptyDir));
String allMessages = messagesOf(thrown);
- assertTrue(
- allMessages.contains("no .vortex files found"),
- "error should explain that no Vortex files were found, got: " + allMessages);
- assertTrue(allMessages.contains(emptyDir), "error should name the offending path, got: " + allMessages);
+ assertTrue(allMessages.contains("Unable to infer schema"), "error should report schema inference failure");
+ assertTrue(allMessages.contains("VORTEX"), "error should identify the Vortex file format");
}
/** Concatenates the messages of the whole cause chain, since Spark may wrap data source exceptions. */
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceS3MockTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceS3MockTest.java
similarity index 94%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceS3MockTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceS3MockTest.java
index a324e64478f..5ae2215a267 100644
--- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceS3MockTest.java
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceS3MockTest.java
@@ -27,7 +27,7 @@
* This test verifies that Vortex can correctly read and write files from S3-compatible storage by using S3Mock
* running as a Testcontainer.
*/
-@Testcontainers
+@Testcontainers(disabledWithoutDocker = true)
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
public final class VortexDataSourceS3MockTest {
@@ -147,14 +147,10 @@ private void assertSchemaEquals(StructType expected, StructType actual) {
StructField actualField = actual.fields()[i];
assertEquals(expectedField.name(), actualField.name(), "Field names should match at position " + i);
- assertEquals(
- expectedField.dataType(),
- actualField.dataType(),
- "Field types should match for field: " + expectedField.name());
- assertEquals(
- expectedField.nullable(),
- actualField.nullable(),
- "Field nullability should match for field: " + expectedField.name());
+ org.junit.jupiter.api.Assertions.assertTrue(
+ expectedField.dataType().sameType(actualField.dataType()),
+ "Field types should match ignoring file-source nullability for field: " + expectedField.name());
+ // Spark's FileTable intentionally makes inferred file fields nullable.
}
}
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java
new file mode 100644
index 00000000000..943803ef973
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java
@@ -0,0 +1,117 @@
+// 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 static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.IOException;
+import java.nio.file.Path;
+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.catalyst.plans.logical.Statistics;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.io.TempDir;
+
+/** Integration tests for row-count and byte-size statistics exposed to Catalyst. */
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public final class VortexDataSourceStatsTest {
+ private SparkSession spark;
+
+ @TempDir
+ Path tempDir;
+
+ @BeforeAll
+ public void setUp() {
+ spark = SparkSession.builder()
+ .appName("VortexStatsTest")
+ .master("local[2]")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.sql.adaptive.enabled", "false")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ }
+
+ @AfterAll
+ public void tearDown() {
+ if (spark != null) {
+ spark.stop();
+ }
+ }
+
+ @Test
+ void reportsExactRowCountAcrossFiles() throws IOException {
+ Path output = writeRows(400, "row_count", 4);
+
+ Statistics statistics = statistics(read(output));
+
+ assertTrue(statistics.rowCount().isDefined());
+ assertEquals(400L, statistics.rowCount().get().longValue());
+ assertTrue(statistics.sizeInBytes().longValue() > 0);
+ }
+
+ @Test
+ void rowCountStatisticsCanBeDisabled() throws IOException {
+ Path output = writeRows(25, "disabled", 2);
+
+ Dataset data = spark.read()
+ .format("vortex")
+ .option("vortex.stats.rowCount", "false")
+ .load(output.toString());
+
+ assertFalse(statistics(data).rowCount().isDefined());
+ assertEquals(25, data.count());
+ }
+
+ @Test
+ void rowCountIsSkippedWhenTheDatasetHasTooManyFiles() throws IOException {
+ Path output = writeRows(40, "too_many_files", 4);
+
+ Dataset capped = spark.read()
+ .format("vortex")
+ .option("vortex.stats.maxFiles", "2")
+ .load(output.toString());
+
+ // Each footer costs a read on the driver, so a large dataset reports no row count rather than
+ // paying for the whole listing before the job starts.
+ assertFalse(statistics(capped).rowCount().isDefined());
+ assertEquals(40, capped.count());
+ }
+
+ @Test
+ void projectedScanHasSmallerByteEstimate() throws IOException {
+ Path output = writeRows(120, "projected", 3);
+ Dataset full = read(output);
+ Dataset projected = full.select("id");
+
+ assertTrue(statistics(projected).sizeInBytes().longValue()
+ < statistics(full).sizeInBytes().longValue());
+ }
+
+ private Statistics statistics(Dataset data) {
+ return data.queryExecution().optimizedPlan().stats();
+ }
+
+ private Dataset read(Path output) {
+ return spark.read().format("vortex").load(output.toString());
+ }
+
+ private Path writeRows(int count, String name, int partitions) throws IOException {
+ Path output = tempDir.resolve(name);
+ spark.range(0, count)
+ .selectExpr("cast(id as int) as id", "concat('value_', cast(id as string)) as value")
+ .repartition(partitions)
+ .write()
+ .format("vortex")
+ .mode(SaveMode.Overwrite)
+ .save(output.toString());
+ return output;
+ }
+}
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java
similarity index 88%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java
index 14b66c9e10a..5e57328fb9a 100644
--- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java
@@ -26,6 +26,7 @@
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.MapType;
+import org.apache.spark.sql.types.Metadata;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.junit.jupiter.api.AfterAll;
@@ -112,8 +113,8 @@ public void testWriteAndReadVortexFiles() throws IOException {
// Verify files have expected naming pattern
for (Path file : vortexFiles) {
assertTrue(
- file.getFileName().toString().matches("part-\\d{5}-\\d+\\.vortex"),
- "File should match pattern part-XXXXX-Y.vortex: " + file.getFileName());
+ file.getFileName().toString().matches("part-\\d{5}-.+\\.vortex"),
+ "File should use Spark's part-file naming: " + file.getFileName());
}
// When: Read the Vortex files back
@@ -171,6 +172,29 @@ public void testWriteAndReadWithBarePath() throws IOException {
assertEquals(25, reread.count(), "Should have data from second write after overwrite");
}
+ @Test
+ @DisplayName("A null in a nullable column round-trips as a null")
+ void nullInNullableColumnRoundTrips() {
+ StructType schema = new StructType(new StructField[] {
+ new StructField("id", DataTypes.IntegerType, false, Metadata.empty()),
+ new StructField("amount", DataTypes.LongType, true, Metadata.empty())
+ });
+ Dataset df =
+ spark.createDataFrame(Arrays.asList(RowFactory.create(1, 10L), RowFactory.create(2, null)), schema);
+ Path output = tempDir.resolve("nullable_null");
+
+ df.write().format("vortex").mode(SaveMode.Overwrite).save(output.toString());
+
+ List rows = spark.read()
+ .format("vortex")
+ .load(output.toString())
+ .orderBy("id")
+ .collectAsList();
+ assertEquals(2, rows.size());
+ assertEquals(10L, rows.get(0).get(1));
+ assertTrue(rows.get(1).isNullAt(1), "expected a null for the second row");
+ }
+
@Test
@DisplayName("Write empty DataFrame as Vortex")
public void testWriteEmptyDataFrame() throws IOException {
@@ -278,12 +302,54 @@ public void testPartitionedWrite() throws IOException {
// Verify partition values are correct
Dataset groupA = readDf.filter(readDf.col("group").equalTo("A")).orderBy("id");
+ String plan = groupA.queryExecution().executedPlan().toString();
+ assertTrue(
+ plan.contains("PartitionFilters:") && plan.contains("group"), "Expected partition pruning:\n" + plan);
assertEquals(3, groupA.count(), "Group A should have 3 rows");
assertEquals(1, (int) groupA.collectAsList().get(0).getAs("id"));
assertEquals(3, (int) groupA.collectAsList().get(1).getAs("id"));
assertEquals(5, (int) groupA.collectAsList().get(2).getAs("id"));
}
+ @Test
+ @DisplayName("Dynamic partition overwrite only replaces touched partitions")
+ public void testDynamicPartitionOverwrite() {
+ StructType schema = DataTypes.createStructType(Arrays.asList(
+ DataTypes.createStructField("id", DataTypes.IntegerType, false),
+ DataTypes.createStructField("group", DataTypes.StringType, false)));
+ Dataset initial =
+ spark.createDataFrame(Arrays.asList(RowFactory.create(1, "A"), RowFactory.create(2, "B")), schema);
+ Dataset replacement = spark.createDataFrame(List.of(RowFactory.create(3, "A")), schema);
+ String outputPath =
+ tempDir.resolve("dynamic_partition_overwrite").toUri().toString();
+
+ initial.write()
+ .format("vortex")
+ .partitionBy("group")
+ .mode(SaveMode.Overwrite)
+ .save(outputPath);
+
+ spark.conf().set("spark.sql.sources.partitionOverwriteMode", "dynamic");
+ try {
+ replacement
+ .write()
+ .format("vortex")
+ .partitionBy("group")
+ .mode(SaveMode.Overwrite)
+ .save(outputPath);
+ } finally {
+ spark.conf().set("spark.sql.sources.partitionOverwriteMode", "static");
+ }
+
+ List rows = spark.read()
+ .format("vortex")
+ .load(outputPath)
+ .select("id", "group")
+ .orderBy("id")
+ .collectAsList();
+ assertEquals(List.of(RowFactory.create(2, "B"), RowFactory.create(3, "A")), rows);
+ }
+
@Test
@DisplayName("Write and read with multiple partition columns")
public void testMultiColumnPartitionedWrite() throws IOException {
@@ -570,14 +636,10 @@ private void assertSchemaEquals(StructType expected, StructType actual) {
StructField actualField = actual.fields()[i];
assertEquals(expectedField.name(), actualField.name(), "Field names should match at position " + i);
- assertEquals(
- expectedField.dataType(),
- actualField.dataType(),
- "Field types should match for field: " + expectedField.name());
- assertEquals(
- expectedField.nullable(),
- actualField.nullable(),
- "Field nullability should match for field: " + expectedField.name());
+ assertTrue(
+ expectedField.dataType().sameType(actualField.dataType()),
+ "Field types should match ignoring file-source nullability for field: " + expectedField.name());
+ // Spark's FileTable intentionally makes inferred file fields nullable.
}
}
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexFileExtensionTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexFileExtensionTest.java
new file mode 100644
index 00000000000..4681f494a19
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexFileExtensionTest.java
@@ -0,0 +1,132 @@
+// 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 static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.stream.Stream;
+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.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.io.TempDir;
+
+/**
+ * Every file in a Vortex dataset must end with {@code .vortex}.
+ *
+ * Spark's file index keeps {@code _metadata} and {@code _common_metadata}, and it keeps every extension it does not
+ * recognise, so the connector is the only thing that can decide what belongs to the dataset. These tests hold the three
+ * paths that see the listing — schema inference, scan statistics, and the scan itself — to the same answer.
+ */
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public final class VortexFileExtensionTest {
+ private SparkSession spark;
+
+ @TempDir
+ Path tempDir;
+
+ @BeforeAll
+ public void setUp() {
+ spark = SparkSession.builder()
+ .appName("VortexFileExtensionTest")
+ .master("local[2]")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.sql.adaptive.enabled", "false")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ }
+
+ @AfterAll
+ public void tearDown() {
+ if (spark != null) {
+ spark.stop();
+ }
+ }
+
+ @Test
+ void schemaInferenceRejectsADirectoryOfNonVortexFiles() throws IOException {
+ Path output = tempDir.resolve("renamed");
+ write(output, 20);
+ renameEveryVortexFile(output, ".dat");
+
+ Exception failure = assertThrows(
+ Exception.class,
+ () -> spark.read().format("vortex").load(output.toString()).schema());
+
+ assertTrue(rootMessage(failure).contains(".vortex"), rootMessage(failure));
+ }
+
+ @Test
+ void aStrayNonVortexFileFailsTheScan() throws IOException {
+ Path output = tempDir.resolve("stray");
+ write(output, 20);
+ // Spark's listing hides names that begin with `_` or `.`, and keeps everything else.
+ Files.write(output.resolve("stray.dat"), new byte[] {1, 2, 3});
+
+ // Inference still finds a Vortex file, so the stray one survives until the scan reaches it.
+ Dataset data = spark.read().format("vortex").load(output.toString());
+
+ Exception failure = assertThrows(Exception.class, data::count);
+ assertTrue(rootMessage(failure).contains("stray.dat"), rootMessage(failure));
+ }
+
+ @Test
+ void aDirectPathToANonVortexFileIsRejected() throws IOException {
+ Path file = tempDir.resolve("bare.dat");
+ Files.write(file, new byte[] {1, 2, 3});
+
+ Exception failure = assertThrows(
+ Exception.class,
+ () -> spark.read().format("vortex").load(file.toString()).count());
+
+ assertTrue(rootMessage(failure).contains(".vortex"), rootMessage(failure));
+ }
+
+ @Test
+ void anUppercaseExtensionIsStillAVortexFile() throws IOException {
+ Path output = tempDir.resolve("uppercase");
+ write(output, 15);
+ renameEveryVortexFile(output, ".VORTEX");
+
+ assertEquals(15, spark.read().format("vortex").load(output.toString()).count());
+ }
+
+ private void write(Path output, int rows) {
+ spark.range(0, rows)
+ .selectExpr("cast(id as int) as id", "concat('value_', cast(id as string)) as value")
+ .repartition(2)
+ .write()
+ .format("vortex")
+ .mode(SaveMode.Overwrite)
+ .save(output.toString());
+ }
+
+ private static void renameEveryVortexFile(Path directory, String extension) throws IOException {
+ try (Stream files = Files.list(directory)) {
+ for (Path file : files.toList()) {
+ String name = file.getFileName().toString();
+ if (name.endsWith(".vortex")) {
+ Files.move(file, file.resolveSibling(name.replace(".vortex", extension)));
+ }
+ }
+ }
+ }
+
+ private static String rootMessage(Throwable failure) {
+ StringBuilder messages = new StringBuilder();
+ for (Throwable cause = failure; cause != null; cause = cause.getCause()) {
+ messages.append(cause.getMessage()).append('\n');
+ }
+ return messages.toString();
+ }
+}
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexFilterPushdownTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexFilterPushdownTest.java
similarity index 99%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/VortexFilterPushdownTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexFilterPushdownTest.java
index 61087837042..fea685adb91 100644
--- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexFilterPushdownTest.java
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexFilterPushdownTest.java
@@ -378,7 +378,7 @@ public void testPushedFiltersInPlan() throws IOException {
SparkPlan plan = qe.executedPlan();
String planString = plan.toString();
assertTrue(
- planString.contains("id > 1"),
+ planString.contains("PushedFilters: [IsNotNull(id), GreaterThan(id,1)]"),
"Expected pushed predicate for id > 1 in the executed plan: " + planString);
}
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexOptionsTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexOptionsTest.java
new file mode 100644
index 00000000000..5c3fa8ff511
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexOptionsTest.java
@@ -0,0 +1,96 @@
+// 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 static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.util.Map;
+import java.util.Optional;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+final class VortexOptionsTest {
+ @ParameterizedTest
+ @ValueSource(
+ strings = {
+ "vortex.readConcurrency",
+ "vortex.readconcurrency",
+ "VORTEX.READCONCURRENCY",
+ "Vortex.ReadConcurrency"
+ })
+ void anyCasingFindsTheValue(String key) {
+ VortexOptions options = VortexOptions.of(Map.of("vortex.readConcurrency", "4"));
+
+ assertEquals(Optional.of("4"), options.get(key));
+ assertEquals(4, options.getInt(key, 1));
+ }
+
+ @Test
+ void anUnsetOptionFallsBack() {
+ VortexOptions options = VortexOptions.empty();
+
+ assertEquals(Optional.empty(), options.get("vortex.absent"));
+ assertEquals("fallback", options.get("vortex.absent", "fallback"));
+ assertEquals(7, options.getInt("vortex.absent", 7));
+ assertTrue(options.getBoolean("vortex.absent", true));
+ }
+
+ @Test
+ void booleansAreReadWithoutRegardToCase() {
+ VortexOptions options = VortexOptions.of(Map.of("on", "TRUE", "off", " False "));
+
+ assertTrue(options.getBoolean("on", false));
+ assertFalse(options.getBoolean("off", true));
+ }
+
+ @Test
+ void aValueOfTheWrongShapeNamesItsOption() {
+ VortexOptions options = VortexOptions.of(Map.of("vortex.count", "many", "vortex.flag", "yes"));
+
+ assertTrue(assertThrows(IllegalArgumentException.class, () -> options.getInt("vortex.count", 1))
+ .getMessage()
+ .contains("vortex.count"));
+ assertTrue(assertThrows(IllegalArgumentException.class, () -> options.getBoolean("vortex.flag", true))
+ .getMessage()
+ .contains("vortex.flag"));
+ }
+
+ @Test
+ void hadoopKeysKeepTheCasingTheyWereGivenIn() {
+ Map given = Map.of("fs.s3a.Endpoint", "https://storage.example");
+
+ // Hadoop configuration keys are case-sensitive, so the map that feeds a Configuration must not be
+ // lower-cased along with the vortex.* lookups.
+ assertEquals(given, VortexOptions.of(given).asCaseSensitiveMap());
+ }
+
+ @Test
+ void optionsSurviveSerialization() throws IOException, ClassNotFoundException {
+ VortexOptions options = VortexOptions.of(Map.of("vortex.workerThreads", "2"));
+
+ VortexOptions shipped = roundTrip(options);
+
+ assertEquals(2, shipped.getInt("vortex.workerthreads", 1));
+ assertEquals(options, shipped);
+ }
+
+ private static VortexOptions roundTrip(VortexOptions options) throws IOException, ClassNotFoundException {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ try (ObjectOutputStream out = new ObjectOutputStream(bytes)) {
+ out.writeObject(options);
+ }
+ try (ObjectInputStream in = new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) {
+ return (VortexOptions) in.readObject();
+ }
+ }
+}
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexProjectionTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexProjectionTest.java
new file mode 100644
index 00000000000..ff56f78967b
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexProjectionTest.java
@@ -0,0 +1,115 @@
+// 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 static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Path;
+import java.util.List;
+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.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.io.TempDir;
+
+/**
+ * Reads that need no data column at all.
+ *
+ * A query over partition columns alone, and a count Spark declines to push down, both leave the read data schema
+ * empty. The scan must still push an empty projection, or it pulls every column off storage to answer them.
+ */
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public final class VortexProjectionTest {
+ private SparkSession spark;
+
+ @TempDir
+ Path tempDir;
+
+ @BeforeAll
+ public void setUp() {
+ spark = SparkSession.builder()
+ .appName("VortexProjectionTest")
+ .master("local[2]")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.sql.adaptive.enabled", "false")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ }
+
+ @AfterAll
+ public void tearDown() {
+ if (spark != null) {
+ spark.stop();
+ }
+ }
+
+ @Test
+ void selectingOnlyAPartitionColumnReadsNoDataColumn() {
+ Dataset data = writeAndRead("partition_only", true).select("group");
+
+ List rows = data.collectAsList();
+
+ assertEquals(30, rows.size());
+ assertEquals(15, rows.stream().filter(row -> row.getInt(0) == 0).count());
+ assertEquals(15, rows.stream().filter(row -> row.getInt(0) == 1).count());
+ // No data column is named anywhere in the scan, so nothing but the partition value was read.
+ assertFalse(plan(data).contains("value"), plan(data));
+ }
+
+ @Test
+ void countWithAggregatePushdownOffStillCountsEveryRow() {
+ Path output = tempDir.resolve("count_no_pushdown");
+ write(output, false);
+ Dataset data = spark.read()
+ .format("vortex")
+ .option("vortex.aggregatePushdown", "false")
+ .load(output.toString());
+
+ Dataset count = data.selectExpr("count(*) AS count");
+
+ assertEquals(30L, count.first().getLong(0));
+ assertTrue(plan(count).contains("PushedAggregation: []"), plan(count));
+ }
+
+ @Test
+ void constantProjectionStillReportsEveryRow() {
+ Dataset data = writeAndRead("constant", false);
+
+ assertEquals(30, data.selectExpr("1 AS one").collectAsList().size());
+ }
+
+ private String plan(Dataset data) {
+ return data.queryExecution().executedPlan().toString();
+ }
+
+ private Dataset writeAndRead(String name, boolean partitioned) {
+ Path output = tempDir.resolve(name);
+ write(output, partitioned);
+ return spark.read().format("vortex").load(output.toString());
+ }
+
+ private void write(Path output, boolean partitioned) {
+ Dataset data = spark.range(0, 30)
+ .selectExpr(
+ "cast(id as int) as id",
+ "concat('value_', cast(id as string)) as value",
+ "cast(id % 2 as int) as group")
+ .repartition(3);
+ if (partitioned) {
+ data.write()
+ .format("vortex")
+ .partitionBy("group")
+ .mode(SaveMode.Overwrite)
+ .save(output.toString());
+ } else {
+ data.write().format("vortex").mode(SaveMode.Overwrite).save(output.toString());
+ }
+ }
+}
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexSchemaMergeTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexSchemaMergeTest.java
new file mode 100644
index 00000000000..cf0fb0221ba
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexSchemaMergeTest.java
@@ -0,0 +1,210 @@
+// 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 static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Path;
+import java.util.List;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
+import org.apache.spark.sql.SaveMode;
+import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.StructType;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.io.TempDir;
+
+/**
+ * Datasets whose files do not all carry the same columns.
+ *
+ * Schema inference merges every footer, so a column added by a later write belongs to the dataset, and the files
+ * written before it read as null.
+ */
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public final class VortexSchemaMergeTest {
+ private static final StructType NARROW =
+ new StructType().add("a", DataTypes.IntegerType, false).add("b", DataTypes.StringType, true);
+
+ private static final StructType WIDE = new StructType()
+ .add("a", DataTypes.IntegerType, false)
+ .add("b", DataTypes.StringType, true)
+ .add("c", DataTypes.DoubleType, true);
+
+ private SparkSession spark;
+
+ @TempDir
+ Path tempDir;
+
+ @BeforeAll
+ public void setUp() {
+ spark = SparkSession.builder()
+ .appName("VortexSchemaMergeTest")
+ .master("local[2]")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.sql.adaptive.enabled", "false")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ }
+
+ @AfterAll
+ public void tearDown() {
+ if (spark != null) {
+ spark.stop();
+ }
+ }
+
+ @Test
+ void aColumnOnlyOneFileCarriesBelongsToTheDataset() {
+ Path dir = tempDir.resolve("widened");
+ append(dir, NARROW, RowFactory.create(1, "one"));
+ append(dir, WIDE, RowFactory.create(2, "two", 2.5));
+
+ Dataset data = read(dir).orderBy("a");
+
+ assertEquals("struct", data.schema().simpleString());
+ List rows = data.collectAsList();
+ assertEquals(2, rows.size());
+ // The file written before `c` existed has no value for it.
+ assertNull(rows.get(0).get(2));
+ assertEquals(2.5, rows.get(1).getDouble(2));
+ }
+
+ @Test
+ void aColumnOnlyOneFileCarriesIsNullableEvenWhenThatFileRequiresIt() {
+ Path dir = tempDir.resolve("nullability");
+ append(dir, NARROW, RowFactory.create(1, "one"));
+ append(
+ dir,
+ new StructType()
+ .add("a", DataTypes.IntegerType, false)
+ .add("b", DataTypes.StringType, true)
+ .add("d", DataTypes.IntegerType, false),
+ RowFactory.create(2, "two", 9));
+
+ assertTrue(read(dir).schema().apply("d").nullable());
+ }
+
+ @Test
+ void selectingOnlyTheAddedColumnStillReadsEveryFile() {
+ Path dir = tempDir.resolve("projected");
+ append(dir, NARROW, RowFactory.create(1, "one"));
+ append(dir, WIDE, RowFactory.create(2, "two", 2.5));
+
+ List rows = read(dir).selectExpr("c").collectAsList();
+
+ assertEquals(2, rows.size());
+ assertEquals(1, rows.stream().filter(row -> row.isNullAt(0)).count());
+ }
+
+ @Test
+ void filteringOnTheAddedColumnKeepsOnlyTheFileThatCarriesIt() {
+ Path dir = tempDir.resolve("filtered");
+ append(dir, NARROW, RowFactory.create(1, "one"));
+ append(dir, WIDE, RowFactory.create(2, "two", 2.5));
+
+ // The filter cannot be pushed into the file that lacks `c`, so Spark's own filter above the scan is what
+ // drops that file's rows.
+ List rows = read(dir).where("c > 1.0").collectAsList();
+
+ assertEquals(1, rows.size());
+ assertEquals(2, rows.get(0).getInt(0));
+ }
+
+ @Test
+ void mergingCanBeTurnedOffToReadOneFooter() {
+ Path dir = tempDir.resolve("unmerged");
+ append(dir, NARROW, RowFactory.create(1, "one"));
+ append(dir, NARROW, RowFactory.create(2, "two"));
+
+ // An unusable footer parallelism proves which path ran: merging reads footers in a pool, while a single
+ // footer read never asks for one.
+ Dataset data = spark.read()
+ .format("vortex")
+ .option("vortex.mergeSchema", "false")
+ .option("vortex.stats.rowCount", "false")
+ .option("vortex.footerParallelism", "0")
+ .load(dir.toString());
+
+ assertEquals("struct", data.schema().simpleString());
+ assertEquals(2, data.count());
+ }
+
+ @Test
+ void mergingIsFoundUnderAnyOptionCasing() {
+ Path dir = tempDir.resolve("option_case");
+ append(dir, NARROW, RowFactory.create(1, "one"));
+
+ Dataset data = spark.read()
+ .format("vortex")
+ .option("VORTEX.MERGESCHEMA", "false")
+ .option("VORTEX.STATS.ROWCOUNT", "false")
+ .option("VORTEX.FOOTERPARALLELISM", "0")
+ .load(dir.toString());
+
+ assertEquals(1, data.count());
+ }
+
+ @Test
+ void aFieldWithTwoTypesNamesTheFileThatDisagrees() {
+ Path dir = tempDir.resolve("conflict");
+ append(dir, NARROW, RowFactory.create(1, "one"));
+ append(
+ dir,
+ new StructType().add("a", DataTypes.StringType, false).add("b", DataTypes.StringType, true),
+ RowFactory.create("two", "two"));
+
+ String message = assertThrows(
+ IllegalArgumentException.class, () -> read(dir).schema())
+ .getMessage();
+
+ assertTrue(message.contains(".vortex"), message);
+ assertTrue(message.contains("field a is"), message);
+ assertTrue(message.contains("vortex.mergeSchema"), message);
+ }
+
+ @Test
+ void aStructThatGainedAFieldCannotBeMerged() {
+ Path dir = tempDir.resolve("nested_conflict");
+ StructType narrowStruct = new StructType()
+ .add("id", DataTypes.IntegerType, false)
+ .add("s", new StructType().add("x", DataTypes.IntegerType, true), true);
+ StructType widerStruct = new StructType()
+ .add("id", DataTypes.IntegerType, false)
+ .add(
+ "s",
+ new StructType().add("x", DataTypes.IntegerType, true).add("y", DataTypes.StringType, true),
+ true);
+ append(dir, narrowStruct, RowFactory.create(1, RowFactory.create(1)));
+ append(dir, widerStruct, RowFactory.create(2, RowFactory.create(2, "two")));
+
+ String message = assertThrows(
+ IllegalArgumentException.class, () -> read(dir).schema())
+ .getMessage();
+
+ // A struct column is projected whole, as the file stores it, so a widened struct cannot be read back
+ // through one merged schema.
+ assertTrue(message.contains("nested field s"), message);
+ }
+
+ private Dataset read(Path dir) {
+ return spark.read().format("vortex").load(dir.toString());
+ }
+
+ private void append(Path dir, StructType schema, Row... rows) {
+ spark.createDataFrame(List.of(rows), schema)
+ .coalesce(1)
+ .write()
+ .format("vortex")
+ .mode(SaveMode.Append)
+ .save(dir.toString());
+ }
+}
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSqlTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexSqlTest.java
similarity index 76%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/VortexSqlTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexSqlTest.java
index 1a9b21ef779..1f00f9ef47f 100644
--- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSqlTest.java
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexSqlTest.java
@@ -7,7 +7,6 @@
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assumptions.assumeTrue;
import java.io.IOException;
import java.nio.file.Files;
@@ -32,7 +31,7 @@
/**
* Integration tests for Spark SQL access to Vortex: managed tables created without a {@code LOCATION} clause, and
- * direct file queries through {@link VortexCatalog}.
+ * direct file queries through Spark's file-source SQL plumbing.
*/
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
public final class VortexSqlTest {
@@ -51,7 +50,6 @@ public void setUp() throws IOException {
.master("local[2]")
.config("spark.driver.host", "127.0.0.1")
.config("spark.sql.warehouse.dir", warehouseDir.toUri().toString())
- .config("spark.sql.catalog.vortex", VortexCatalog.class.getName())
.config("spark.ui.enabled", "false")
.getOrCreate();
}
@@ -63,23 +61,9 @@ public void tearDown() {
}
}
- /**
- * Spark 3.5's built-in session catalog cannot read tables backed by a DataSource-V2-only provider: its
- * {@code FindDataSourceTable} rule falls back to the V1 {@code DataSource} path, which rejects the provider with
- * "vortex is not a valid Spark SQL Data Source". Spark 4 resolves such tables through the provider directly. On
- * Spark 3.5 the {@link VortexSessionCatalog} extension provides the same support — see
- * {@link VortexSessionCatalogTest}, which runs these scenarios on both versions.
- */
- private void assumeSupportsSqlTables() {
- assumeTrue(
- spark.version().startsWith("4."),
- "CREATE TABLE ... USING vortex requires Spark 4 or the VortexSessionCatalog extension");
- }
-
@Test
@DisplayName("Managed table lifecycle: CREATE, SELECT while empty, INSERT, INSERT OVERWRITE, DROP")
public void testManagedTableLifecycle() {
- assumeSupportsSqlTables();
spark.sql("CREATE TABLE managed_students (id INT, name STRING, age INT) USING vortex");
assertEquals(0, spark.sql("SELECT * FROM managed_students").count(), "New managed table should be empty");
@@ -102,7 +86,6 @@ public void testManagedTableLifecycle() {
@Test
@DisplayName("CREATE TABLE AS SELECT without a LOCATION clause")
public void testCreateManagedTableAsSelect() {
- assumeSupportsSqlTables();
spark.sql("CREATE TABLE ctas_source (id INT, name STRING) USING vortex");
spark.sql("INSERT INTO ctas_source VALUES (1, 'Alice'), (2, 'Bob')");
@@ -119,36 +102,27 @@ public void testCreateManagedTableAsSelect() {
@DisplayName("Reading the vortex format without a path option still fails")
public void testReadWithoutPathStillThrows() {
assertThrows(
- IllegalArgumentException.class,
+ AnalysisException.class,
() -> spark.read().format("vortex").load(),
"A read with no path should not silently return an empty DataFrame");
}
@Test
- @DisplayName("Direct file query through the vortex catalog")
+ @DisplayName("Direct file query through Spark's file-source fallback")
public void testDirectPathQuery() {
Path dataDir = tempDir.resolve("direct_query");
writeTestData(dataDir);
- List rows = spark.sql(String.format("SELECT name FROM vortex.`%s` WHERE age > 30 ORDER BY name", dataDir))
- .collectAsList();
+ Dataset query =
+ spark.sql(String.format("SELECT name FROM vortex.`%s` WHERE age > 30 ORDER BY name", dataDir));
+ String plan = query.queryExecution().executedPlan().toString();
+ assertTrue(plan.contains("Batched: true"), "V1 fallback should return columnar batches:\n" + plan);
+ List rows = query.collectAsList();
assertEquals(2, rows.size());
assertEquals("Alice", rows.get(0).getString(0));
assertEquals("Carol", rows.get(1).getString(0));
}
- @Test
- @DisplayName("INSERT INTO a direct path through the vortex catalog")
- public void testDirectPathInsert() {
- Path dataDir = tempDir.resolve("direct_insert");
- writeTestData(dataDir);
-
- spark.sql(String.format("INSERT INTO vortex.`%s` VALUES (4, 'Dave', 50)", dataDir));
- assertEquals(
- 4,
- spark.sql(String.format("SELECT * FROM vortex.`%s`", dataDir)).count());
- }
-
@Test
@DisplayName("Direct queries of missing paths and non-path names fail as table-not-found")
public void testDirectPathNotFound() {
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexV1FallbackTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexV1FallbackTest.java
new file mode 100644
index 00000000000..0fa2cfcabeb
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexV1FallbackTest.java
@@ -0,0 +1,146 @@
+// 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 static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Path;
+import java.util.List;
+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.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.junit.jupiter.api.io.TempDir;
+
+/**
+ * The V1 {@link VortexFileFormat} path.
+ *
+ * Catalog tables reach it, and so does {@code spark.sql.sources.useV1SourceList}. Within it, Spark asks for rows
+ * instead of batches when whole-stage codegen is off or the schema carries more fields than
+ * {@code spark.sql.codegen.maxFields}, so both readers in the format need cover.
+ */
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public final class VortexV1FallbackTest {
+ private SparkSession spark;
+
+ @TempDir
+ Path tempDir;
+
+ @BeforeAll
+ public void setUp() {
+ spark = SparkSession.builder()
+ .appName("VortexV1FallbackTest")
+ .master("local[2]")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.sql.adaptive.enabled", "false")
+ .config("spark.ui.enabled", "false")
+ .config("spark.sql.sources.useV1SourceList", "vortex")
+ .getOrCreate();
+ // Set it again at runtime: `getOrCreate` reuses a session another test class in this JVM may have
+ // built, and then it ignores the builder options.
+ spark.conf().set("spark.sql.sources.useV1SourceList", "vortex");
+ }
+
+ @AfterAll
+ public void tearDown() {
+ if (spark != null) {
+ spark.stop();
+ }
+ }
+
+ @Test
+ void readsAndWritesThroughTheV1FileFormat() {
+ Path output = write("v1_roundtrip", false);
+ Dataset data = spark.read().format("vortex").load(output.toString());
+
+ // `FileScan vortex` is the V1 plan node. The V2 path would read `BatchScan`.
+ assertTrue(plan(data).contains("FileScan vortex"), plan(data));
+ assertEquals(40, data.count());
+ assertEquals(
+ List.of(0, 1, 2),
+ data.orderBy("id").limit(3).collectAsList().stream()
+ .map(row -> row.getInt(0))
+ .toList());
+ }
+
+ @Test
+ void prunesColumnsAndPushesFiltersThroughTheV1FileFormat() {
+ Path output = write("v1_pushdown", false);
+
+ Dataset data = spark.read()
+ .format("vortex")
+ .load(output.toString())
+ .select("value")
+ .filter("value = 'value_7'");
+
+ assertEquals(1, data.count());
+ assertEquals("value_7", data.first().getString(0));
+ }
+
+ @Test
+ void readsHivePartitionValuesThroughTheV1FileFormat() {
+ Path output = write("v1_partitioned", true);
+
+ Dataset data = spark.read().format("vortex").load(output.toString());
+
+ assertEquals(40, data.count());
+ assertEquals(20, data.filter("grp = 0").count());
+ assertEquals(20, data.select("grp").filter("grp = 1").count());
+ }
+
+ @Test
+ void readsRowsWhenSparkDeclinesTheColumnarPath() {
+ Path output = write("v1_rows", true);
+
+ // Whole-stage codegen off makes `FileSourceScanExec` ask the format for rows, not batches.
+ spark.conf().set("spark.sql.codegen.wholeStage", "false");
+ try {
+ Dataset data = spark.read().format("vortex").load(output.toString());
+
+ assertEquals(40, data.count());
+ List rows = data.orderBy("id").limit(2).collectAsList();
+ assertEquals(0, intOf(rows.get(0), "id"));
+ assertEquals("value_0", rows.get(0).getString(rows.get(0).fieldIndex("value")));
+ // Partition values are joined onto the row by Spark, not by the reader.
+ assertEquals(0, intOf(rows.get(0), "grp"));
+ assertEquals(1, intOf(rows.get(1), "id"));
+ assertEquals(20, data.filter("grp = 1").count());
+ } finally {
+ spark.conf().unset("spark.sql.codegen.wholeStage");
+ }
+ }
+
+ private static int intOf(Row row, String name) {
+ return row.getInt(row.fieldIndex(name));
+ }
+
+ private String plan(Dataset data) {
+ return data.queryExecution().executedPlan().toString();
+ }
+
+ private Path write(String name, boolean partitioned) {
+ Path output = tempDir.resolve(name);
+ Dataset data = spark.range(0, 40)
+ .selectExpr(
+ "cast(id as int) as id",
+ "concat('value_', cast(id as string)) as value",
+ "cast(id % 2 as int) as grp")
+ .repartition(2);
+ if (partitioned) {
+ data.write()
+ .format("vortex")
+ .partitionBy("grp")
+ .mode(SaveMode.Overwrite)
+ .save(output.toString());
+ } else {
+ data.write().format("vortex").mode(SaveMode.Overwrite).save(output.toString());
+ }
+ return output;
+ }
+}
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java
new file mode 100644
index 00000000000..afef7223b2a
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java
@@ -0,0 +1,177 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.io;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import dev.vortex.io.NativeReadable;
+import java.io.EOFException;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import org.apache.hadoop.conf.Configuration;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+final class HadoopReadableTest {
+ private static final int SIZE = 512 * 1024;
+
+ private final Configuration conf = new Configuration();
+
+ @TempDir
+ Path tempDir;
+
+ @Test
+ void readsEveryRangeOfTheFile() throws IOException {
+ byte[] content = content(SIZE);
+ Path file = write("data.vortex", content);
+
+ try (NativeReadable readable = HadoopReadable.open(conf, file.toString())) {
+ assertEquals(SIZE, readable.length());
+ assertEquals(file.toString(), readable.name());
+
+ for (int[] range : new int[][] {{0, SIZE}, {0, 1}, {SIZE - 1, 1}, {1234, 65536}, {SIZE / 2, SIZE / 2}}) {
+ ByteBuffer buffer = ByteBuffer.allocateDirect(range[1]);
+ readable.readFully(range[0], buffer);
+ assertArrayEquals(
+ slice(content, range[0], range[1]), drain(buffer), "range " + range[0] + "+" + range[1]);
+ }
+ }
+ }
+
+ @Test
+ void suppliedLengthIsTrustedOverAnExtraStat() throws IOException {
+ Path file = write("data.vortex", content(64));
+
+ try (NativeReadable readable = HadoopReadable.open(conf, file.toString(), 64)) {
+ assertEquals(64, readable.length());
+ }
+ }
+
+ @Test
+ void concurrentReadsEachGetTheirOwnStream() throws Exception {
+ byte[] content = content(SIZE);
+ Path file = write("data.vortex", content);
+ int readers = 8;
+ int chunk = SIZE / readers;
+
+ try (NativeReadable readable = HadoopReadable.open(conf, file.toString())) {
+ ExecutorService pool = Executors.newFixedThreadPool(readers);
+ try {
+ List> reads = new ArrayList<>();
+ for (int i = 0; i < readers; i++) {
+ int offset = i * chunk;
+ reads.add(() -> {
+ byte[] read = new byte[chunk];
+ // Read the same range repeatedly so streams keep being returned to and taken
+ // from the pool while other threads are doing the same.
+ for (int round = 0; round < 20; round++) {
+ ByteBuffer buffer = ByteBuffer.allocateDirect(chunk);
+ readable.readFully(offset, buffer);
+ read = drain(buffer);
+ }
+ return read;
+ });
+ }
+
+ List> results = pool.invokeAll(reads);
+ for (int i = 0; i < readers; i++) {
+ assertArrayEquals(
+ slice(content, i * chunk, chunk), results.get(i).get(), "reader " + i);
+ }
+ } finally {
+ pool.shutdownNow();
+ }
+ }
+ }
+
+ @Test
+ void readsPastTheEndFail() throws IOException {
+ Path file = write("data.vortex", content(1024));
+
+ try (NativeReadable readable = HadoopReadable.open(conf, file.toString())) {
+ assertThrows(EOFException.class, () -> readable.readFully(512, ByteBuffer.allocateDirect(1024)));
+ assertThrows(EOFException.class, () -> readable.readFully(1024, ByteBuffer.allocateDirect(1)));
+ }
+ }
+
+ @Test
+ void readingAfterCloseFails() throws IOException {
+ Path file = write("data.vortex", content(1024));
+
+ NativeReadable readable = HadoopReadable.open(conf, file.toString());
+ readable.readFully(0, ByteBuffer.allocateDirect(16));
+ readable.close();
+
+ assertThrows(IllegalStateException.class, () -> readable.readFully(0, ByteBuffer.allocateDirect(16)));
+ }
+
+ @Test
+ void writtenBytesAreReadBackThroughTheHadoopBridge() throws IOException {
+ byte[] content = content(4096);
+ Path file = tempDir.resolve("written.vortex");
+
+ try (HadoopWritable writable = HadoopWritable.create(conf, file.toString())) {
+ writable.write(content, 0, 1024);
+ writable.flush();
+ writable.write(content, 1024, content.length - 1024);
+ }
+
+ assertArrayEquals(content, Files.readAllBytes(file));
+
+ try (NativeReadable readable = HadoopReadable.open(conf, file.toString())) {
+ ByteBuffer buffer = ByteBuffer.allocateDirect(content.length);
+ readable.readFully(0, buffer);
+ assertArrayEquals(content, drain(buffer));
+ }
+ }
+
+ @Test
+ void createReplacesAnExistingFile() throws IOException {
+ Path file = tempDir.resolve("replaced.vortex");
+ Files.write(file, new byte[] {1, 2, 3, 4, 5, 6, 7, 8});
+
+ try (HadoopWritable writable = HadoopWritable.create(conf, file.toString())) {
+ writable.write(new byte[] {9}, 0, 1);
+ }
+
+ assertArrayEquals(new byte[] {9}, Files.readAllBytes(file));
+ }
+
+ private Path write(String name, byte[] content) throws IOException {
+ Path file = tempDir.resolve(name);
+ Files.write(file, content);
+ return file;
+ }
+
+ private static byte[] content(int size) {
+ byte[] content = new byte[size];
+ for (int i = 0; i < size; i++) {
+ content[i] = (byte) (i % 251);
+ }
+ return content;
+ }
+
+ private static byte[] slice(byte[] content, int offset, int length) {
+ byte[] expected = new byte[length];
+ System.arraycopy(content, offset, expected, 0, length);
+ return expected;
+ }
+
+ private static byte[] drain(ByteBuffer buffer) {
+ byte[] read = new byte[buffer.position()];
+ buffer.flip();
+ buffer.get(read);
+ return read;
+ }
+}
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/VortexIoTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/VortexIoTest.java
new file mode 100644
index 00000000000..a7800f25786
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/io/VortexIoTest.java
@@ -0,0 +1,78 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.io;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import dev.vortex.spark.VortexOptions;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.util.Map;
+import org.apache.hadoop.conf.Configuration;
+import org.junit.jupiter.api.Test;
+
+final class VortexIoTest {
+ @Test
+ void readConcurrencyDefaultsToNativeDefault() {
+ VortexIo io = VortexIo.create(VortexOptions.empty(), new Configuration());
+
+ assertEquals(0, io.readConcurrency());
+ }
+
+ @Test
+ void readConcurrencyIsParsedAndValidated() {
+ Configuration conf = new Configuration();
+
+ assertEquals(
+ 8,
+ VortexIo.create(options(VortexIo.READ_CONCURRENCY_OPTION, "8"), conf)
+ .readConcurrency());
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> VortexIo.create(options(VortexIo.READ_CONCURRENCY_OPTION, "-1"), conf));
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> VortexIo.create(options(VortexIo.READ_CONCURRENCY_OPTION, "many"), conf));
+ }
+
+ @Test
+ void readConcurrencyIsFoundUnderAnyKeyCase() {
+ Configuration conf = new Configuration();
+
+ assertEquals(
+ 6, VortexIo.create(options("VORTEX.READCONCURRENCY", "6"), conf).readConcurrency());
+ assertEquals(
+ 6, VortexIo.create(options("vortex.readconcurrency", "6"), conf).readConcurrency());
+ }
+
+ @Test
+ void settingsAndConfigurationSurviveSerialization() throws IOException, ClassNotFoundException {
+ Configuration conf = new Configuration();
+ conf.set("fs.s3a.endpoint", "https://storage.example");
+ VortexIo io = VortexIo.create(options(VortexIo.READ_CONCURRENCY_OPTION, "4"), conf);
+
+ VortexIo shipped = roundTrip(io);
+
+ assertEquals(4, shipped.readConcurrency());
+ assertEquals("https://storage.example", shipped.hadoopConf().get("fs.s3a.endpoint"));
+ }
+
+ private static VortexOptions options(String key, String value) {
+ return VortexOptions.of(Map.of(key, value));
+ }
+
+ private static VortexIo roundTrip(VortexIo io) throws IOException, ClassNotFoundException {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ try (ObjectOutputStream out = new ObjectOutputStream(bytes)) {
+ out.writeObject(io);
+ }
+ try (ObjectInputStream in = new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) {
+ return (VortexIo) in.readObject();
+ }
+ }
+}
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/read/SparkFilterToVortexExpressionTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/read/SparkFilterToVortexExpressionTest.java
new file mode 100644
index 00000000000..ed3583a7fcd
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/read/SparkFilterToVortexExpressionTest.java
@@ -0,0 +1,101 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.read;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.math.BigDecimal;
+import java.sql.Date;
+import java.sql.Timestamp;
+import org.apache.spark.sql.sources.And;
+import org.apache.spark.sql.sources.EqualNullSafe;
+import org.apache.spark.sql.sources.EqualTo;
+import org.apache.spark.sql.sources.GreaterThan;
+import org.apache.spark.sql.sources.In;
+import org.apache.spark.sql.sources.IsNotNull;
+import org.apache.spark.sql.sources.IsNull;
+import org.apache.spark.sql.sources.Not;
+import org.apache.spark.sql.sources.Or;
+import org.apache.spark.sql.sources.StringContains;
+import org.apache.spark.sql.sources.StringStartsWith;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.StructType;
+import org.junit.jupiter.api.Test;
+
+final class SparkFilterToVortexExpressionTest {
+ private static final StructType SCHEMA = new StructType()
+ .add("id", DataTypes.IntegerType)
+ .add("date", DataTypes.DateType)
+ .add("timestamp", DataTypes.TimestampType)
+ .add("decimal", DataTypes.createDecimalType(10, 2))
+ .add("binary", DataTypes.BinaryType)
+ .add("name", DataTypes.StringType)
+ .add("nested", new StructType().add("value", DataTypes.LongType));
+
+ @Test
+ void convertsExternalTemporalDecimalAndBinaryLiterals() {
+ assertPushable(new EqualTo("date", Date.valueOf("2024-01-02")));
+ assertPushable(new GreaterThan("timestamp", Timestamp.valueOf("2024-01-02 03:04:05.123456")));
+ assertPushable(new EqualTo("decimal", new BigDecimal("123.45")));
+ assertPushable(new EqualTo("binary", new byte[] {1, 2, 3}));
+ }
+
+ @Test
+ void convertsNestedLogicalAndStringFilters() {
+ assertPushable(new EqualTo("nested.value", 42L));
+ assertPushable(new And(new GreaterThan("id", 1), new StringContains("name", "%literal_")));
+ assertPushable(new In("id", new Object[] {1, 2, null}));
+ }
+
+ @Test
+ void convertsNullSafeEqualityAndNegation() {
+ assertPushable(new EqualNullSafe("id", null));
+ assertPushable(new EqualNullSafe("id", 7));
+ assertPushable(new Not(new EqualTo("id", 7)));
+ assertPushable(new IsNull("name"));
+ assertPushable(new IsNotNull("name"));
+ }
+
+ @Test
+ void convertsInListsHoldingNulls() {
+ // SQL never matches a null through IN, so the null values drop out of the disjunction.
+ assertPushable(new In("id", new Object[] {1, null}));
+ // Nothing but nulls can never match, so the filter becomes a constant false.
+ assertPushable(new In("id", new Object[] {null}));
+ }
+
+ @Test
+ void rejectsDecimalLiteralsThatCannotHoldTheirScale() {
+ // The column keeps two decimal places, so a third would be lost in the conversion.
+ assertFalse(SparkFilterToVortexExpression.isPushable(new EqualTo("decimal", new BigDecimal("1.005")), SCHEMA));
+ assertPushable(new EqualTo("decimal", new BigDecimal("1.00")));
+ }
+
+ @Test
+ void rejectsAPartlyConvertibleConjunctionAndDisjunction() {
+ assertFalse(SparkFilterToVortexExpression.isPushable(
+ new And(new GreaterThan("id", 1), new EqualTo("missing", 2)), SCHEMA));
+ assertFalse(SparkFilterToVortexExpression.isPushable(
+ new Or(new GreaterThan("id", 1), new EqualTo("missing", 2)), SCHEMA));
+ }
+
+ @Test
+ void rejectsStringMatchesOnNonStringColumns() {
+ assertFalse(SparkFilterToVortexExpression.isPushable(new StringContains("id", "1"), SCHEMA));
+ assertFalse(SparkFilterToVortexExpression.isPushable(new StringStartsWith("id", "1"), SCHEMA));
+ }
+
+ @Test
+ void rejectsMissingColumnsAndMismatchedLiteralTypes() {
+ assertFalse(SparkFilterToVortexExpression.isPushable(new EqualTo("missing", 1), SCHEMA));
+ assertFalse(SparkFilterToVortexExpression.isPushable(new EqualTo("id", "one"), SCHEMA));
+ assertFalse(SparkFilterToVortexExpression.isPushable(new EqualTo("id", null), SCHEMA));
+ }
+
+ private static void assertPushable(org.apache.spark.sql.sources.Filter filter) {
+ assertTrue(SparkFilterToVortexExpression.isPushable(filter, SCHEMA));
+ assertTrue(SparkFilterToVortexExpression.convert(filter, SCHEMA).isPresent());
+ }
+}
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java
similarity index 100%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java
diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java
similarity index 100%
rename from java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java
rename to java/vortex-spark/common/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java
diff --git a/java/vortex-spark/common/src/test/java/dev/vortex/spark/write/VortexOutputWriterTest.java b/java/vortex-spark/common/src/test/java/dev/vortex/spark/write/VortexOutputWriterTest.java
new file mode 100644
index 00000000000..d179fd92026
--- /dev/null
+++ b/java/vortex-spark/common/src/test/java/dev/vortex/spark/write/VortexOutputWriterTest.java
@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.spark.write;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import dev.vortex.spark.VortexOptions;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+final class VortexOutputWriterTest {
+ private static final int DEFAULT = 2048;
+
+ @Test
+ void anUnsetBatchSizeFallsBackToTheDefault() {
+ assertEquals(DEFAULT, VortexOutputWriter.configuredBatchSize(VortexOptions.empty()));
+ }
+
+ @Test
+ void theGenericBatchSizeIsHonoured() {
+ assertEquals(4096, VortexOutputWriter.configuredBatchSize(options(Map.of("batch.size", "4096"))));
+ }
+
+ @Test
+ void theVortexBatchSizeOverridesTheGenericOne() {
+ VortexOptions options = options(Map.of("batch.size", "4096", "vortex.write.batch.size", "512"));
+
+ // A job that sets one batch size across formats can still say something different for Vortex.
+ assertEquals(512, VortexOutputWriter.configuredBatchSize(options));
+ }
+
+ @Test
+ void theOverrideStandsEvenWhenTheGenericValueIsUnusable() {
+ VortexOptions options = options(Map.of("batch.size", "lots", "vortex.write.batch.size", "512"));
+
+ assertEquals(512, VortexOutputWriter.configuredBatchSize(options));
+ }
+
+ @Test
+ void eitherNameIsFoundUnderAnyCasing() {
+ assertEquals(64, VortexOutputWriter.configuredBatchSize(options(Map.of("BATCH.SIZE", "64"))));
+ assertEquals(64, VortexOutputWriter.configuredBatchSize(options(Map.of("Vortex.Write.Batch.Size", "64"))));
+ }
+
+ @Test
+ void aBatchSizeOfTheWrongShapeNamesItsOption() {
+ VortexOptions options = options(Map.of("vortex.write.batch.size", "lots"));
+
+ assertTrue(assertThrows(IllegalArgumentException.class, () -> VortexOutputWriter.configuredBatchSize(options))
+ .getMessage()
+ .contains("vortex.write.batch.size"));
+ }
+
+ private static VortexOptions options(Map values) {
+ return VortexOptions.of(values);
+ }
+}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java
deleted file mode 100644
index 2bd4fd3273d..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java
+++ /dev/null
@@ -1,129 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark;
-
-import java.util.Map;
-import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
-import org.apache.spark.sql.connector.catalog.Identifier;
-import org.apache.spark.sql.connector.catalog.Table;
-import org.apache.spark.sql.connector.catalog.TableCatalog;
-import org.apache.spark.sql.connector.catalog.TableChange;
-import org.apache.spark.sql.connector.expressions.Transform;
-import org.apache.spark.sql.types.StructType;
-import org.apache.spark.sql.util.CaseInsensitiveStringMap;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * A path-based Spark catalog for querying Vortex files directly from SQL.
- *
- * Spark only supports {@code SELECT * FROM format.`path`} syntax for built-in file formats, so this catalog provides
- * the equivalent for Vortex. Register it under the name {@code vortex}:
- *
- *
spark.sql.catalog.vortex=dev.vortex.spark.VortexCatalog
- *
- * then query a Vortex file, or a directory of Vortex files, directly by path:
- *
- *
SELECT * FROM vortex.`/path/to/data`;
- *
- * The table identifier must look like a path — contain a {@code /} — and resolves to the same table a
- * {@code spark.read.format("vortex")} load of that path would produce, so reads, writes ({@code INSERT INTO
- * vortex.`/path/to/data`}), and pushdown all behave identically. The catalog holds no state and supports no DDL.
- */
-public final class VortexCatalog implements TableCatalog {
- private static final Logger log = LoggerFactory.getLogger(VortexCatalog.class);
-
- private static final String PATH_KEY = "path";
-
- private String name = "vortex";
-
- /**
- * Creates a new catalog instance.
- *
- *
This no-argument constructor is required for Spark to instantiate the catalog through reflection from the
- * {@code spark.sql.catalog.} configuration.
- */
- public VortexCatalog() {}
-
- @Override
- public void initialize(String name, CaseInsensitiveStringMap options) {
- this.name = name;
- }
-
- @Override
- public String name() {
- return name;
- }
-
- /**
- * Returns no identifiers: this catalog holds no state, tables are addressed by path.
- *
- * @param namespace the namespace to list, ignored
- * @return an empty array
- */
- @Override
- public Identifier[] listTables(String[] namespace) {
- return new Identifier[0];
- }
-
- /**
- * Loads the Vortex file or directory of Vortex files at the path given by the identifier name.
- *
- * @param ident identifier whose name is a filesystem path or URL, e.g. {@code vortex.`/path/to/data`}
- * @return a table backed by the Vortex files at the path
- * @throws NoSuchTableException if the identifier does not look like a path, or the path cannot be read
- */
- @SuppressWarnings("deprecation")
- @Override
- public Table loadTable(Identifier ident) throws NoSuchTableException {
- String path = ident.name();
- if (ident.namespace().length != 0 || !path.contains("/")) {
- throw new NoSuchTableException(ident);
- }
- var options = new CaseInsensitiveStringMap(Map.of(PATH_KEY, path));
- var provider = new VortexDataSourceV2();
- StructType schema;
- Transform[] partitioning;
- try {
- schema = provider.inferSchema(options);
- partitioning = provider.inferPartitioning(options);
- } catch (RuntimeException e) {
- // Missing or unreadable paths surface as "table not found" to SQL users. NoSuchTableException
- // has no cause-accepting constructor common to the supported Spark versions, so carry the
- // original failure as a suppressed exception and log it: it names the offending path and
- // distinguishes an empty directory from a credentials or unsupported-type error.
- log.warn("Cannot load {} as a Vortex table, reporting it as not found", path, e);
- NoSuchTableException notFound = new NoSuchTableException(ident);
- notFound.addSuppressed(e);
- throw notFound;
- }
- return provider.getTable(schema, partitioning, Map.of(PATH_KEY, path));
- }
-
- /** Unsupported: tables are addressed by path, create them by writing data with the {@code vortex} format. */
- @Override
- public Table createTable(
- Identifier ident, StructType schema, Transform[] partitions, Map properties) {
- throw new UnsupportedOperationException(
- "VortexCatalog does not support CREATE TABLE, write data to the path instead");
- }
-
- /** Unsupported: this catalog holds no table metadata to alter. */
- @Override
- public Table alterTable(Identifier ident, TableChange... changes) {
- throw new UnsupportedOperationException("VortexCatalog does not support ALTER TABLE");
- }
-
- /** Unsupported: this catalog never drops data, returns false. */
- @Override
- public boolean dropTable(Identifier ident) {
- return false;
- }
-
- /** Unsupported: this catalog holds no table metadata to rename. */
- @Override
- public void renameTable(Identifier oldIdent, Identifier newIdent) {
- throw new UnsupportedOperationException("VortexCatalog does not support RENAME TABLE");
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java
deleted file mode 100644
index bd434bcc8d2..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java
+++ /dev/null
@@ -1,249 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark;
-
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
-import dev.vortex.api.DataSource;
-import dev.vortex.jni.NativeFiles;
-import dev.vortex.spark.config.HadoopUtils;
-import dev.vortex.spark.read.PartitionPathUtils;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Optional;
-import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
-import org.apache.spark.sql.SparkSession;
-import org.apache.spark.sql.connector.catalog.Table;
-import org.apache.spark.sql.connector.catalog.TableProvider;
-import org.apache.spark.sql.connector.expressions.Expressions;
-import org.apache.spark.sql.connector.expressions.Transform;
-import org.apache.spark.sql.sources.DataSourceRegister;
-import org.apache.spark.sql.types.DataType;
-import org.apache.spark.sql.types.Metadata;
-import org.apache.spark.sql.types.StructField;
-import org.apache.spark.sql.types.StructType;
-import org.apache.spark.sql.util.CaseInsensitiveStringMap;
-import scala.Option;
-
-/**
- * Spark V2 data source for reading and writing Vortex files.
- *
- * This class is automatically registered so it can be discovered by the Spark runtime. For reading:
- * {@link org.apache.spark.sql.SparkSession#read} and specify the format as "vortex". For writing:
- * {@link org.apache.spark.sql.Dataset#write} and specify the format as "vortex".
- */
-public final class VortexDataSourceV2 implements TableProvider, DataSourceRegister {
- private static final ObjectMapper MAPPER = new ObjectMapper();
-
- private static final String PATH_KEY = "path";
- private static final String PATHS_KEY = "paths";
-
- private final Option sparkSession;
-
- /**
- * Creates a new instance of the Vortex data source.
- *
- * This no-argument constructor is required for Spark to instantiate the data source through reflection.
- */
- public VortexDataSourceV2() {
- this.sparkSession = SparkSession.getActiveSession();
- }
-
- /**
- * Infers the schema of the Vortex files specified in the options.
- *
- *
This method examines the last file in the provided paths to determine the schema. Currently, schema evolution
- * and merging across multiple files is not supported.
- *
- * @param options the data source options containing file paths
- * @return the inferred Spark SQL schema
- * @throws IllegalArgumentException if no Vortex files can be found under the supplied paths
- * @throws RuntimeException if there's an error reading the file or converting the schema
- */
- @Override
- public StructType inferSchema(CaseInsensitiveStringMap options) {
- // For write operations, the path might not exist yet
- // In that case, return an empty schema to signal Spark to use the DataFrame's schema
- var paths = getPaths(options);
-
- // If path is not found, we report empty schema.
- // This will be replaced with whatever the DataFrame schema is
- if (paths.isEmpty()) {
- return new StructType();
- }
-
- var formatOptions = buildDataSourceOptions(options.asCaseSensitiveMap());
-
- var pathToInfer = Objects.requireNonNull(Iterables.getLast(paths));
- // If the path is a directory, scan the directory for a file and use that file
- if (!pathToInfer.endsWith(".vortex")) {
- Optional firstFile =
- NativeFiles.listFiles(VortexSparkSession.get(formatOptions), pathToInfer, formatOptions).stream()
- .findFirst();
-
- if (firstFile.isEmpty()) {
- throw new IllegalArgumentException(String.format(
- "Unable to infer schema for %s: no .vortex files found under path %s. "
- + "Check that the path is correct and contains at least one Vortex file, "
- + "or provide an explicit schema.",
- shortName(), pathToInfer));
- } else {
- pathToInfer = firstFile.get();
- }
- }
-
- StructType dataSchema;
- {
- DataSource ds = DataSource.open(VortexSparkSession.get(formatOptions), pathToInfer, formatOptions);
- var arrowSchema = ds.arrowSchema(dev.vortex.arrow.ArrowAllocation.rootAllocator());
- StructField[] fields = arrowSchema.getFields().stream()
- .map(f -> new StructField(
- f.getName(), ArrowUtils.fromArrowField(f), f.isNullable(), Metadata.empty()))
- .toArray(StructField[]::new);
- dataSchema = new StructType(fields);
- }
-
- // Discover partition columns from Hive-style directory paths and append them.
- Map partitionValues = PartitionPathUtils.parsePartitionValues(pathToInfer);
- if (!partitionValues.isEmpty()) {
- Set dataColumnNames = Stream.of(dataSchema.fieldNames()).collect(Collectors.toSet());
- for (Map.Entry entry : partitionValues.entrySet()) {
- if (!dataColumnNames.contains(entry.getKey())) {
- DataType type = PartitionPathUtils.inferPartitionColumnType(entry.getValue());
- dataSchema = dataSchema.add(entry.getKey(), type, true);
- }
- }
- }
-
- return dataSchema;
- }
-
- /**
- * Infers partition transforms by inspecting Hive-style {@code key=value} segments in the first listed file path.
- *
- * Spark calls this before {@link #getTable(StructType, Transform[], Map)} when the caller did not provide
- * explicit partitioning. Returning identity transforms here lets downstream components (notably
- * {@link dev.vortex.spark.read.VortexScanBuilder}) tell which schema columns are encoded in the directory layout
- * rather than stored inside the Vortex files, which matters for predicate pushdown.
- *
- *
The options may contain no path at all: when a managed table is created ({@code CREATE TABLE ... USING vortex}
- * without a {@code LOCATION} clause), Spark's session catalog calls this before it has assigned the table's
- * warehouse location. No transforms are inferred in that case.
- */
- @Override
- public Transform[] inferPartitioning(CaseInsensitiveStringMap options) {
- var paths = getPathsOrEmpty(options);
- if (paths.isEmpty()) {
- return new Transform[0];
- }
- var formatOptions = buildDataSourceOptions(options.asCaseSensitiveMap());
- String pathToInfer = Objects.requireNonNull(Iterables.getLast(paths));
- if (!pathToInfer.endsWith(".vortex")) {
- Optional firstFile =
- NativeFiles.listFiles(VortexSparkSession.get(formatOptions), pathToInfer, formatOptions).stream()
- .findFirst();
- if (firstFile.isEmpty()) {
- return new Transform[0];
- }
- pathToInfer = firstFile.get();
- }
- Map partitionValues = PartitionPathUtils.parsePartitionValues(pathToInfer);
- if (partitionValues.isEmpty()) {
- return new Transform[0];
- }
- return partitionValues.keySet().stream().map(Expressions::identity).toArray(Transform[]::new);
- }
-
- /**
- * Creates a Vortex table instance with the given schema and properties.
- *
- * This method creates a VortexWritableTable that can be used to both read from and write to Vortex files. The
- * partitioning parameter is currently ignored.
- *
- *
The properties may contain no path at all: when a managed table is created ({@code CREATE TABLE ... USING
- * vortex} without a {@code LOCATION} clause), Spark's session catalog validates the table before it has assigned
- * the table's warehouse location. The returned table then has no paths; once the table is loaded for reads or
- * writes, Spark always supplies the resolved table location as the {@code path} property.
- *
- * @param schema the table schema
- * @param partitioning table partitioning transforms
- * @param properties the table properties containing file paths and other options
- * @return a VortexTable instance for reading and writing data
- */
- @Override
- public Table getTable(StructType schema, Transform[] partitioning, Map properties) {
- var uncased = new CaseInsensitiveStringMap(properties);
- ImmutableList paths = getPathsOrEmpty(uncased);
- return new VortexTable(paths, schema, buildDataSourceOptions(properties), partitioning);
- }
-
- /**
- * Indicates whether this data source supports external metadata (schemas).
- *
- * Returns true to indicate that this data source accepts external schemas, which is necessary for write
- * operations where the DataFrame provides the schema.
- *
- * @return true to accept external schemas
- */
- @Override
- public boolean supportsExternalMetadata() {
- return true;
- }
-
- /**
- * Returns the short name identifier for this data source.
- *
- *
This name is used by Spark when registering the data source and can be used in SQL queries and DataFrame read
- * operations to specify this format.
- *
- * @return the short name "vortex"
- */
- @Override
- public String shortName() {
- return "vortex";
- }
-
- private Map buildDataSourceOptions(Map properties) {
- var hadoopConf = sparkSession.get().sessionState().newHadoopConf();
-
- var options = ImmutableMap.builder();
- options.putAll(properties);
-
- // Forward any S3-relevant properties from hadoopConf to the reader config.
- options.putAll(HadoopUtils.s3PropertiesFromHadoopConf(hadoopConf));
- // Forward any Azure-relevant properties from hadoopConf to the reader config.
- options.putAll(HadoopUtils.azurePropertiesFromHadoopConf(hadoopConf));
-
- return options.build();
- }
-
- private static ImmutableList getPathsOrEmpty(CaseInsensitiveStringMap uncased) {
- if (!uncased.containsKey(PATH_KEY) && !uncased.containsKey(PATHS_KEY)) {
- return ImmutableList.of();
- }
- return getPaths(uncased);
- }
-
- private static ImmutableList getPaths(CaseInsensitiveStringMap uncased) {
- if (uncased.containsKey(PATH_KEY)) {
- return ImmutableList.of(uncased.get(PATH_KEY));
- } else if (uncased.containsKey(PATHS_KEY)) {
- return decodePathsSafe(uncased.get(PATHS_KEY));
- } else {
- throw new IllegalArgumentException("Missing required option: \"path\" or \"paths\"");
- }
- }
-
- private static ImmutableList decodePathsSafe(String pathsJson) {
- try {
- return ImmutableList.copyOf(MAPPER.readValue(pathsJson, String[].class));
- } catch (Exception e) {
- throw new IllegalArgumentException("Failed to decode \"paths\" option", e);
- }
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexFilePartition.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexFilePartition.java
deleted file mode 100644
index 4247788a887..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexFilePartition.java
+++ /dev/null
@@ -1,32 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark;
-
-import java.io.Serializable;
-import java.util.List;
-import java.util.Map;
-import org.apache.spark.sql.connector.read.InputPartition;
-import org.apache.spark.sql.types.StructType;
-
-/**
- * An {@link InputPartition} describing a group of Vortex files that a single reader should handle together.
- *
- * Each executor opens a single Vortex {@code Session}, {@code DataSource} and {@code Scan} over the partition's
- * {@link #paths()} and consumes every Vortex partition produced by that scan before moving on to the next Spark
- * {@code InputPartition}.
- *
- *
The requested output schema is carried as a {@link StructType} rather than a list of {@code Column} objects:
- * {@code StructType} is the stable serialization surface in Spark and survives shipping to executors reliably.
- *
- * @param paths the Vortex file paths (or globs) belonging to this input partition
- * @param readSchema the requested output schema (data columns + partition columns)
- * @param formatOptions object-store properties used to open the files
- * @param partitionValues Hive-style partition column values shared by all {@link #paths()}
- */
-public record VortexFilePartition(
- List paths,
- StructType readSchema,
- Map formatOptions,
- Map partitionValues)
- implements InputPartition, Serializable {}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java
deleted file mode 100644
index d1162bbf391..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java
+++ /dev/null
@@ -1,85 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark;
-
-import java.util.Map;
-import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException;
-import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
-import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException;
-import org.apache.spark.sql.connector.catalog.DelegatingCatalogExtension;
-import org.apache.spark.sql.connector.catalog.Identifier;
-import org.apache.spark.sql.connector.catalog.Table;
-import org.apache.spark.sql.connector.catalog.TableCatalog;
-import org.apache.spark.sql.connector.expressions.Transform;
-import org.apache.spark.sql.types.StructType;
-
-/**
- * A session catalog extension that resolves {@code USING vortex} tables through the Vortex DataSource V2 connector.
- *
- * Spark 3.5's built-in session catalog resolves the tables it stores through the V1 {@code DataSource} path, which
- * rejects DataSource-V2-only connectors like Vortex, so {@code CREATE TABLE ... USING vortex} tables cannot be read
- * back. (Spark 4 resolves them through the V2 provider directly and needs none of this.) Registering this extension as
- * the session catalog fixes that on Spark 3.5:
- *
- *
spark.sql.catalog.spark_catalog=dev.vortex.spark.VortexSessionCatalog
- *
- * All operations are delegated to the built-in session catalog — table metadata lives wherever it normally would,
- * including the Hive metastore — but any table whose provider is {@code vortex} is loaded as a Vortex DataSource V2
- * table, backed by the files at the table's location. Tables of every other provider are untouched.
- */
-public final class VortexSessionCatalog extends DelegatingCatalogExtension {
-
- /**
- * Creates a new session catalog extension.
- *
- *
This no-argument constructor is required for Spark to instantiate the catalog through reflection from the
- * {@code spark.sql.catalog.spark_catalog} configuration.
- */
- public VortexSessionCatalog() {}
-
- @Override
- public Table loadTable(Identifier ident) throws NoSuchTableException {
- return asVortexTableIfVortex(super.loadTable(ident));
- }
-
- /**
- * Creates the table in the delegate session catalog, then returns it resolved through the Vortex connector when its
- * provider is {@code vortex}.
- *
- *
Spark does not route table creation through this overload — {@link DelegatingCatalogExtension} sends the
- * {@code Column[]} overload straight to the delegate, and that is the one both supported Spark versions call. It is
- * kept because the delegate may return {@code null} on the normal path, which {@link #asVortexTableIfVortex} now
- * tolerates; {@code CREATE TABLE ... AS SELECT} gets its Vortex table from {@link #loadTable} instead.
- */
- @SuppressWarnings("deprecation")
- @Override
- public Table createTable(
- Identifier ident, StructType schema, Transform[] partitions, Map properties)
- throws TableAlreadyExistsException, NoSuchNamespaceException {
- return asVortexTableIfVortex(super.createTable(ident, schema, partitions, properties));
- }
-
- /**
- * Rebuilds a session-catalog table as a Vortex DataSource V2 table when its provider is {@code vortex} and it has a
- * location; returns every other table unchanged, and {@code null} for a {@code null} input. The schema and
- * partitioning stored in the catalog are used as-is, no file needs to be opened.
- */
- @SuppressWarnings("deprecation")
- private static Table asVortexTableIfVortex(Table table) {
- if (table == null) {
- // V2SessionCatalog.createTable returns null on its normal path.
- return null;
- }
- Map properties = table.properties();
- VortexDataSourceV2 provider = new VortexDataSourceV2();
- if (!provider.shortName().equalsIgnoreCase(properties.get(TableCatalog.PROP_PROVIDER))) {
- return table;
- }
- String location = properties.get(TableCatalog.PROP_LOCATION);
- if (location == null) {
- return table;
- }
- return provider.getTable(table.schema(), table.partitioning(), Map.of("path", location));
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexTable.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexTable.java
deleted file mode 100644
index f65f74ccf19..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexTable.java
+++ /dev/null
@@ -1,127 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark;
-
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Maps;
-import dev.vortex.spark.read.VortexScanBuilder;
-import dev.vortex.spark.write.VortexWriteBuilder;
-import java.util.Arrays;
-import java.util.Map;
-import java.util.Set;
-import org.apache.spark.sql.connector.catalog.CatalogV2Util;
-import org.apache.spark.sql.connector.catalog.SupportsRead;
-import org.apache.spark.sql.connector.catalog.SupportsWrite;
-import org.apache.spark.sql.connector.catalog.Table;
-import org.apache.spark.sql.connector.catalog.TableCapability;
-import org.apache.spark.sql.connector.expressions.Transform;
-import org.apache.spark.sql.connector.read.ScanBuilder;
-import org.apache.spark.sql.connector.write.LogicalWriteInfo;
-import org.apache.spark.sql.connector.write.WriteBuilder;
-import org.apache.spark.sql.types.StructType;
-import org.apache.spark.sql.util.CaseInsensitiveStringMap;
-
-/** Spark V2 {@link Table} of Vortex files that supports both reading and writing. */
-public final class VortexTable implements Table, SupportsRead, SupportsWrite {
- private static final String SHORT_NAME = "vortex";
-
- private final ImmutableList paths;
- private final StructType schema;
- private final Map formatOptions;
- private final Transform[] partitionTransforms;
-
- /** Creates a new VortexTable with read/write support. */
- public VortexTable(
- ImmutableList paths,
- StructType schema,
- Map formatOptions,
- Transform[] partitionTransforms) {
- this.paths = paths;
- this.schema = schema;
- this.formatOptions = formatOptions;
- this.partitionTransforms = partitionTransforms;
- }
-
- /**
- * Creates a new ScanBuilder for this table.
- *
- * The scan builder is pre-configured with all the file paths and columns from this table.
- *
- * @param options scan options
- * @return a new VortexScanBuilder configured for this table
- */
- @Override
- public ScanBuilder newScanBuilder(CaseInsensitiveStringMap options) {
- Map opts = Maps.newHashMap();
- opts.putAll(formatOptions);
- opts.putAll(options);
- return new VortexScanBuilder(opts, partitionTransforms)
- .addAllPaths(paths)
- .addAllColumns(Arrays.asList(CatalogV2Util.structTypeToV2Columns(schema)));
- }
-
- /**
- * Returns the name of this table.
- *
- * The name includes the "vortex" prefix and a comma-separated list of all file paths that comprise this table.
- *
- * @return the table name in the format: vortex."path1,path2,..."
- */
- @Override
- public String name() {
- return String.format("%s.\"%s\"", SHORT_NAME, String.join(",", paths));
- }
-
- /**
- * Returns the schema of this table.
- *
- *
The schema is derived from the columns available for reading, or from the explicit write schema if this table
- * is being used for writing.
- *
- * @return the StructType representing the table schema
- */
- @Override
- public StructType schema() {
- return schema;
- }
-
- /**
- * Creates a new WriteBuilder for writing data to this table.
- *
- *
The WriteBuilder is responsible for configuring and executing write operations to create new Vortex files.
- *
- * @param info logical information about the write operation
- * @return a new VortexWriteBuilder configured for this table
- */
- @Override
- public WriteBuilder newWriteBuilder(LogicalWriteInfo info) {
- // Make sure only one write path was provided.
- String writePath = Iterables.getOnlyElement(paths);
- return new VortexWriteBuilder(writePath, info, formatOptions, partitionTransforms);
- }
-
- /**
- * Returns the partitioning transforms for this table.
- *
- * @return an array of partition transforms
- */
- @Override
- public Transform[] partitioning() {
- return partitionTransforms;
- }
-
- /**
- * Returns the capabilities supported by this table.
- *
- *
Vortex tables support batch reading and batch writing.
- *
- * @return a set containing TableCapability.BATCH_READ and BATCH_WRITE
- */
- @Override
- public Set capabilities() {
- return ImmutableSet.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE);
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/config/HadoopUtils.java b/java/vortex-spark/src/main/java/dev/vortex/spark/config/HadoopUtils.java
deleted file mode 100644
index 2f3c4d78d90..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/config/HadoopUtils.java
+++ /dev/null
@@ -1,77 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark.config;
-
-import java.util.Map;
-import org.apache.hadoop.conf.Configuration;
-
-public final class HadoopUtils {
- private HadoopUtils() {}
-
- static final String FS_S3A_ACCESS_KEY = "fs.s3a.access.key";
- static final String FS_S3A_SECRET_KEY = "fs.s3a.secret.key";
- static final String FS_S3A_SESSION_TOKEN = "fs.s3a.session.token";
- static final String FS_S3A_ENDPOINT = "fs.s3a.endpoint";
- static final String FS_S3A_ENDPOINT_REGION = "fs.s3a.endpoint.region";
-
- public static Map s3PropertiesFromHadoopConf(Configuration hadoopConf) {
- VortexS3Properties properties = new VortexS3Properties();
-
- for (Map.Entry entry : hadoopConf) {
- switch (entry.getKey()) {
- case FS_S3A_ACCESS_KEY:
- properties.setAccessKeyId(entry.getValue());
- break;
- case FS_S3A_SECRET_KEY:
- properties.setSecretAccessKey(entry.getValue());
- break;
- case FS_S3A_SESSION_TOKEN:
- properties.setSessionToken(entry.getValue());
- break;
- case FS_S3A_ENDPOINT:
- String qualified = entry.getValue();
- if (!qualified.startsWith("http")) {
- qualified = "https://" + qualified;
- }
- properties.setEndpoint(qualified);
- // object_store rejects plain-HTTP endpoints (LocalStack, MinIO, S3Mock)
- // unless explicitly allowed.
- if (qualified.startsWith("http://")) {
- properties.setAllowHttp(true);
- }
- break;
- case FS_S3A_ENDPOINT_REGION:
- properties.setRegion(entry.getValue());
- break;
- default:
- break;
- }
- }
-
- return properties.asProperties();
- }
-
- static final String ACCESS_KEY_PREFIX = "fs.azure.account.key";
- static final String FIXED_TOKEN_PREFIX = "fs.azure.sas.fixed.token.";
-
- public static Map azurePropertiesFromHadoopConf(Configuration hadoopConf) {
- VortexAzureProperties properties = new VortexAzureProperties();
-
- // TODO(aduffy): match on storage account name.
- for (Map.Entry entry : hadoopConf) {
- String configKey = entry.getKey();
- if (configKey.startsWith(ACCESS_KEY_PREFIX)) {
- properties.setAccessKey(entry.getValue());
- } else if (configKey.startsWith(FIXED_TOKEN_PREFIX)) {
- properties.setSasKey(entry.getValue());
- }
- }
-
- if (properties.accessKey().isEmpty()) {
- properties.setSkipSignature(true);
- }
-
- return properties.asProperties();
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexAzureProperties.java b/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexAzureProperties.java
deleted file mode 100644
index b064575c311..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexAzureProperties.java
+++ /dev/null
@@ -1,48 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark.config;
-
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Maps;
-import java.util.Map;
-import java.util.Optional;
-
-public final class VortexAzureProperties {
- private static final String ACCOUNT_KEY = "azure_storage_account_key";
- private static final String SAS_KEY = "azure_storage_sas_key";
- private static final String SKIP_SIGNATURE = "azure_skip_signature";
-
- private final Map properties = Maps.newHashMap();
-
- public Optional accessKey() {
- return Optional.ofNullable(properties.get(ACCOUNT_KEY));
- }
-
- public Optional sasKey() {
- return Optional.ofNullable(properties.get(SAS_KEY));
- }
-
- public boolean skipSignature() {
- return Boolean.parseBoolean(properties.getOrDefault(SKIP_SIGNATURE, "false"));
- }
-
- public VortexAzureProperties setAccessKey(String accountKey) {
- properties.put(ACCOUNT_KEY, accountKey);
- return this;
- }
-
- public VortexAzureProperties setSasKey(String sasKey) {
- properties.put(SAS_KEY, sasKey);
- return this;
- }
-
- public VortexAzureProperties setSkipSignature(boolean skipSignature) {
- properties.put(SKIP_SIGNATURE, String.valueOf(skipSignature));
- return this;
- }
-
- public Map asProperties() {
- return ImmutableMap.copyOf(properties);
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexS3Properties.java b/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexS3Properties.java
deleted file mode 100644
index 2bd2221e413..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexS3Properties.java
+++ /dev/null
@@ -1,77 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark.config;
-
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Maps;
-import java.util.Map;
-import java.util.Optional;
-
-public final class VortexS3Properties {
- private static final String ACCESS_KEY = "aws_access_key_id";
- private static final String SECRET_KEY = "aws_secret_access_key";
- private static final String SESSION_TOKEN = "aws_session_token";
- private static final String REGION = "aws_region";
- private static final String ENDPOINT = "aws_endpoint";
- private static final String ALLOW_HTTP = "aws_allow_http";
- private static final String SKIP_SIGNATURE = "aws_skip_signature";
-
- private final Map properties = Maps.newHashMap();
-
- public Optional accessKeyId() {
- return Optional.ofNullable(properties.get(ACCESS_KEY));
- }
-
- public Optional secretAccessKey() {
- return Optional.ofNullable(properties.get(SECRET_KEY));
- }
-
- public Optional sessionToken() {
- return Optional.ofNullable(properties.get(SESSION_TOKEN));
- }
-
- public Optional region() {
- return Optional.ofNullable(properties.get(REGION));
- }
-
- public Optional endpoint() {
- return Optional.ofNullable(properties.get(ENDPOINT));
- }
-
- public boolean skipSignature() {
- return Boolean.parseBoolean(properties.getOrDefault(SKIP_SIGNATURE, "false"));
- }
-
- public void setAccessKeyId(String accessKeyId) {
- properties.put(ACCESS_KEY, accessKeyId);
- }
-
- public void setSecretAccessKey(String secretAccessKey) {
- properties.put(SECRET_KEY, secretAccessKey);
- }
-
- public void setSessionToken(String sessionToken) {
- properties.put(SESSION_TOKEN, sessionToken);
- }
-
- public void setRegion(String region) {
- properties.put(REGION, region);
- }
-
- public void setEndpoint(String endpoint) {
- properties.put(ENDPOINT, endpoint);
- }
-
- public void setAllowHttp(boolean allowHttp) {
- properties.put(ALLOW_HTTP, Boolean.toString(allowHttp));
- }
-
- public void setSkipSignature(boolean skipSignature) {
- properties.put(SKIP_SIGNATURE, Boolean.toString(skipSignature));
- }
-
- public Map asProperties() {
- return ImmutableMap.copyOf(properties);
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/PartitionPathUtils.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/PartitionPathUtils.java
deleted file mode 100644
index 24a5c8aa953..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/PartitionPathUtils.java
+++ /dev/null
@@ -1,110 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark.read;
-
-import com.google.common.base.Splitter;
-import com.google.common.primitives.Doubles;
-import com.google.common.primitives.Ints;
-import com.google.common.primitives.Longs;
-import java.net.URLDecoder;
-import java.nio.charset.StandardCharsets;
-import java.util.LinkedHashMap;
-import java.util.Map;
-import org.apache.spark.sql.execution.vectorized.ConstantColumnVector;
-import org.apache.spark.sql.types.BooleanType;
-import org.apache.spark.sql.types.ByteType;
-import org.apache.spark.sql.types.DataType;
-import org.apache.spark.sql.types.DataTypes;
-import org.apache.spark.sql.types.DateType;
-import org.apache.spark.sql.types.DoubleType;
-import org.apache.spark.sql.types.FloatType;
-import org.apache.spark.sql.types.IntegerType;
-import org.apache.spark.sql.types.LongType;
-import org.apache.spark.sql.types.ShortType;
-import org.apache.spark.sql.types.StringType;
-import org.apache.spark.sql.types.TimestampNTZType;
-import org.apache.spark.sql.types.TimestampType;
-import org.apache.spark.unsafe.types.UTF8String;
-
-/** Utilities for discovering and materializing Hive-style partition columns from file paths. */
-public final class PartitionPathUtils {
- private static final String HIVE_DEFAULT_PARTITION = "__HIVE_DEFAULT_PARTITION__";
- private static final Splitter PATH_SPLITTER = Splitter.on('/');
-
- private PartitionPathUtils() {}
-
- /**
- * Parses Hive-style {@code key=value} segments from a file path.
- *
- * @return an ordered map of partition column names to their string values
- */
- public static Map parsePartitionValues(String filePath) {
- LinkedHashMap values = new LinkedHashMap<>();
- for (String segment : PATH_SPLITTER.split(filePath)) {
- int eqIdx = segment.indexOf('=');
- if (eqIdx > 0 && eqIdx < segment.length() - 1) {
- String key = URLDecoder.decode(segment.substring(0, eqIdx), StandardCharsets.UTF_8);
- String val = URLDecoder.decode(segment.substring(eqIdx + 1), StandardCharsets.UTF_8);
- values.put(key, val);
- }
- }
- return values;
- }
-
- /**
- * Infers a Spark {@link DataType} from a partition value string. Tries integer, long, double, boolean, and falls
- * back to string.
- */
- public static DataType inferPartitionColumnType(String value) {
- if (value == null || HIVE_DEFAULT_PARTITION.equals(value)) {
- return DataTypes.StringType;
- }
- if (Ints.tryParse(value) != null) {
- return DataTypes.IntegerType;
- }
- if (Longs.tryParse(value) != null) {
- return DataTypes.LongType;
- }
- if (Doubles.tryParse(value) != null) {
- return DataTypes.DoubleType;
- }
- if ("true".equalsIgnoreCase(value) || "false".equalsIgnoreCase(value)) {
- return DataTypes.BooleanType;
- }
- return DataTypes.StringType;
- }
-
- /**
- * Creates a Spark {@link ConstantColumnVector} populated with the given partition value, parsed according to the
- * target {@link DataType}.
- */
- public static ConstantColumnVector createConstantVector(int numRows, DataType type, String value) {
- ConstantColumnVector vec = new ConstantColumnVector(numRows, type);
- if (value == null || HIVE_DEFAULT_PARTITION.equals(value)) {
- vec.setNull();
- return vec;
- }
- vec.setNotNull();
- if (type instanceof StringType) {
- vec.setUtf8String(UTF8String.fromString(value));
- } else if (type instanceof IntegerType || type instanceof DateType) {
- vec.setInt(Integer.parseInt(value));
- } else if (type instanceof LongType || type instanceof TimestampType || type instanceof TimestampNTZType) {
- vec.setLong(Long.parseLong(value));
- } else if (type instanceof ShortType) {
- vec.setShort(Short.parseShort(value));
- } else if (type instanceof ByteType) {
- vec.setByte(Byte.parseByte(value));
- } else if (type instanceof BooleanType) {
- vec.setBoolean(Boolean.parseBoolean(value));
- } else if (type instanceof FloatType) {
- vec.setFloat(Float.parseFloat(value));
- } else if (type instanceof DoubleType) {
- vec.setDouble(Double.parseDouble(value));
- } else {
- vec.setUtf8String(UTF8String.fromString(value));
- }
- return vec;
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/SparkPredicateToVortexExpression.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/SparkPredicateToVortexExpression.java
deleted file mode 100644
index 781c008c79a..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/SparkPredicateToVortexExpression.java
+++ /dev/null
@@ -1,517 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark.read;
-
-import dev.vortex.api.Expression;
-import dev.vortex.api.Expression.BinaryOp;
-import dev.vortex.api.Expression.TimeUnit;
-import java.math.BigDecimal;
-import java.math.BigInteger;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import org.apache.spark.sql.connector.expressions.Literal;
-import org.apache.spark.sql.connector.expressions.NamedReference;
-import org.apache.spark.sql.connector.expressions.filter.AlwaysFalse;
-import org.apache.spark.sql.connector.expressions.filter.AlwaysTrue;
-import org.apache.spark.sql.connector.expressions.filter.And;
-import org.apache.spark.sql.connector.expressions.filter.Not;
-import org.apache.spark.sql.connector.expressions.filter.Or;
-import org.apache.spark.sql.connector.expressions.filter.Predicate;
-import org.apache.spark.sql.types.BinaryType;
-import org.apache.spark.sql.types.BooleanType;
-import org.apache.spark.sql.types.ByteType;
-import org.apache.spark.sql.types.DataType;
-import org.apache.spark.sql.types.DateType;
-import org.apache.spark.sql.types.Decimal;
-import org.apache.spark.sql.types.DecimalType;
-import org.apache.spark.sql.types.DoubleType;
-import org.apache.spark.sql.types.FloatType;
-import org.apache.spark.sql.types.IntegerType;
-import org.apache.spark.sql.types.LongType;
-import org.apache.spark.sql.types.ShortType;
-import org.apache.spark.sql.types.StringType;
-import org.apache.spark.sql.types.StructField;
-import org.apache.spark.sql.types.StructType;
-import org.apache.spark.sql.types.TimestampNTZType;
-import org.apache.spark.sql.types.TimestampType;
-import org.apache.spark.unsafe.types.UTF8String;
-
-/**
- * Translates {@link Predicate Spark V2 predicates} into Vortex {@link Expression}s for predicate pushdown.
- *
- * The translator aims to express every Spark predicate Vortex can evaluate. Predicates that cannot be translated
- * (unsupported functions, literals on user-defined types, references to columns not present in the file, etc.) are left
- * to Spark for post-scan evaluation.
- */
-final class SparkPredicateToVortexExpression {
-
- private SparkPredicateToVortexExpression() {}
-
- /**
- * Returns true if the given Spark predicate can be translated to a Vortex expression and every named reference
- * resolves to a real field path under {@code dataColumnTypes}.
- *
- *
{@code dataColumnTypes} maps each pushable top-level column name to its top-level Spark {@link DataType};
- * partition columns and columns the scan does not project should not appear in the map. For nested references (for
- * example {@code info.email}) the validator walks the named reference part by part, descending into
- * {@link StructType} fields so that {@code info} must be a struct that contains an {@code email} field.
- *
- *
This is the cheap check used in {@code SupportsPushDownV2Filters.pushPredicates} to decide which predicates
- * Spark can drop. It does not allocate any native expressions; if it returns true, {@link #convert(Predicate)} must
- * succeed (otherwise callers would silently drop predicates).
- */
- static boolean isPushable(Predicate predicate, Map dataColumnTypes) {
- for (NamedReference ref : predicate.references()) {
- if (!resolveFieldPath(ref.fieldNames(), dataColumnTypes)) {
- return false;
- }
- }
- return isStructurallyPushable(predicate);
- }
-
- /**
- * Walks {@code parts} against {@code dataColumnTypes}, descending through {@link StructType} fields for
- * dot-separated nested references. Returns true only when every part resolves to an actual field in the schema.
- */
- private static boolean resolveFieldPath(String[] parts, Map dataColumnTypes) {
- if (parts.length == 0) {
- return false;
- }
- DataType current = dataColumnTypes.get(parts[0]);
- if (current == null) {
- return false;
- }
- for (int i = 1; i < parts.length; i++) {
- if (!(current instanceof StructType struct)) {
- return false;
- }
- Optional field = findField(struct, parts[i]);
- if (field.isEmpty()) {
- return false;
- }
- current = field.get().dataType();
- }
- return true;
- }
-
- private static Optional findField(StructType struct, String name) {
- return Arrays.stream(struct.fields())
- .filter(structField -> structField.name().equals(name))
- .findFirst();
- }
-
- private static boolean isStructurallyPushable(Predicate predicate) {
- if (predicate instanceof AlwaysTrue || predicate instanceof AlwaysFalse) {
- return true;
- }
- if (predicate instanceof And a) {
- return isStructurallyPushable(a.left()) && isStructurallyPushable(a.right());
- }
- if (predicate instanceof Or o) {
- return isStructurallyPushable(o.left()) && isStructurallyPushable(o.right());
- }
- if (predicate instanceof Not n) {
- return isStructurallyPushable(n.child());
- }
-
- org.apache.spark.sql.connector.expressions.Expression[] children = predicate.children();
- return switch (predicate.name()) {
- case "=", "<>", "!=", ">", ">=", "<", "<=" -> isPushableComparison(children);
- case "IS_NULL", "IS_NOT_NULL" -> children.length == 1 && isPushableFieldRef(children[0]);
- case "IN" -> {
- if (children.length < 2 || !isPushableFieldRef(children[0])) {
- yield false;
- }
- for (int i = 1; i < children.length; i++) {
- if (!isPushableLiteral(children[i])) {
- yield false;
- }
- }
- yield true;
- }
- case "STARTS_WITH", "ENDS_WITH", "CONTAINS" ->
- children.length == 2 && isPushableFieldRef(children[0]) && isPushableStringLiteral(children[1]);
- // `BOOLEAN_EXPRESSION` wraps a bare boolean-valued child. We only handle the case
- // where the child itself is a field reference (e.g. `WHERE bool_col`).
- case "BOOLEAN_EXPRESSION" -> children.length == 1 && isPushableFieldRef(children[0]);
- default -> false;
- };
- }
-
- /**
- * Converts a Spark predicate to a Vortex expression. Returns {@link Optional#empty()} if the predicate cannot be
- * translated; callers should normally pre-check with {@link #isPushable}.
- */
- static Optional convert(Predicate predicate) {
- if (predicate instanceof AlwaysTrue) {
- return Optional.of(Expression.literal(true));
- }
- if (predicate instanceof AlwaysFalse) {
- return Optional.of(Expression.literal(false));
- }
- if (predicate instanceof And a) {
- Optional left = convert(a.left());
- Optional right = convert(a.right());
- if (left.isPresent() && right.isPresent()) {
- return Optional.of(Expression.and(left.get(), right.get()));
- }
- return Optional.empty();
- }
- if (predicate instanceof Or o) {
- Optional left = convert(o.left());
- Optional right = convert(o.right());
- if (left.isPresent() && right.isPresent()) {
- return Optional.of(Expression.or(left.get(), right.get()));
- }
- return Optional.empty();
- }
- if (predicate instanceof Not n) {
- return convert(n.child()).map(Expression::not);
- }
- org.apache.spark.sql.connector.expressions.Expression[] children = predicate.children();
- return switch (predicate.name()) {
- case "=", "<>", "!=", ">", ">=", "<", "<=" -> convertComparison(predicate.name(), children);
- case "IS_NULL" -> children.length == 1 ? columnOf(children[0]).map(Expression::isNull) : Optional.empty();
- case "IS_NOT_NULL" ->
- children.length == 1 ? columnOf(children[0]).map(Expression::isNotNull) : Optional.empty();
- case "IN" -> convertIn(children);
- case "STARTS_WITH" ->
- convertStringMatch(children, /* leadingWildcard= */ false, /* trailingWildcard= */ true);
- case "ENDS_WITH" ->
- convertStringMatch(children, /* leadingWildcard= */ true, /* trailingWildcard= */ false);
- case "CONTAINS" -> convertStringMatch(children, /* leadingWildcard= */ true, /* trailingWildcard= */ true);
- case "BOOLEAN_EXPRESSION" -> children.length == 1 ? columnOf(children[0]) : Optional.empty();
- default -> Optional.empty();
- };
- }
-
- private static Optional convertComparison(
- String op, org.apache.spark.sql.connector.expressions.Expression[] children) {
- if (children.length != 2) {
- return Optional.empty();
- }
- // Allow either side to be the column; Spark's V2 builder sometimes commutes.
- Optional lhs = exprOf(children[0]);
- Optional rhs = exprOf(children[1]);
- if (lhs.isEmpty() || rhs.isEmpty()) {
- return Optional.empty();
- }
- // We require at least one side to be a column reference to keep the surface small and to
- // match what Vortex pushdown understands.
- boolean lhsIsCol = isFieldRefExpr(children[0]);
- boolean rhsIsCol = isFieldRefExpr(children[1]);
- if (!lhsIsCol && !rhsIsCol) {
- return Optional.empty();
- }
- BinaryOp binaryOp = toBinaryOp(op);
- // Canonicalize so the column is on the left when only one side is a column.
- if (!lhsIsCol) {
- binaryOp = swap(binaryOp);
- Expression tmp = lhs.get();
- return Optional.of(Expression.binary(binaryOp, rhs.get(), tmp));
- }
- return Optional.of(Expression.binary(binaryOp, lhs.get(), rhs.get()));
- }
-
- private static Optional convertIn(org.apache.spark.sql.connector.expressions.Expression[] children) {
- if (children.length < 2) {
- return Optional.empty();
- }
- Optional column = columnOf(children[0]);
- if (column.isEmpty()) {
- return Optional.empty();
- }
- Expression columnExpr = column.get();
- List eqs = new ArrayList<>(children.length - 1);
- for (int i = 1; i < children.length; i++) {
- Optional literal = literalOf(children[i]);
- if (literal.isEmpty()) {
- return Optional.empty();
- }
- eqs.add(Expression.binary(BinaryOp.EQ, columnExpr, literal.get()));
- }
- if (eqs.size() == 1) {
- return Optional.of(eqs.get(0));
- }
- return Optional.of(Expression.or(eqs.toArray(new Expression[0])));
- }
-
- private static Optional convertStringMatch(
- org.apache.spark.sql.connector.expressions.Expression[] children,
- boolean leadingWildcard,
- boolean trailingWildcard) {
- if (children.length != 2) {
- return Optional.empty();
- }
- Optional column = columnOf(children[0]);
- Optional needle = stringValueOf(children[1]);
- if (column.isEmpty() || needle.isEmpty()) {
- return Optional.empty();
- }
- String pattern = buildLikePattern(needle.get(), leadingWildcard, trailingWildcard);
- return Optional.of(Expression.like(
- column.get(), Expression.literal(pattern), /* negated= */ false, /* caseInsensitive= */ false));
- }
-
- /**
- * Build a LIKE pattern from a literal substring, escaping the {@code %}, {@code _}, and {@code \} meta-characters
- * so the Spark {@code STARTS_WITH}/{@code ENDS_WITH}/{@code CONTAINS} semantics (exact substring match) are
- * preserved.
- */
- private static String buildLikePattern(String literal, boolean leadingWildcard, boolean trailingWildcard) {
- StringBuilder sb = new StringBuilder(literal.length() + 2);
- if (leadingWildcard) {
- sb.append('%');
- }
- for (int i = 0; i < literal.length(); i++) {
- char c = literal.charAt(i);
- if (c == '%' || c == '_' || c == '\\') {
- sb.append('\\');
- }
- sb.append(c);
- }
- if (trailingWildcard) {
- sb.append('%');
- }
- return sb.toString();
- }
-
- private static BinaryOp toBinaryOp(String name) {
- return switch (name) {
- case "=" -> BinaryOp.EQ;
- case "<>", "!=" -> BinaryOp.NOT_EQ;
- case ">" -> BinaryOp.GT;
- case ">=" -> BinaryOp.GTE;
- case "<" -> BinaryOp.LT;
- case "<=" -> BinaryOp.LTE;
- default -> throw new IllegalArgumentException("not a pushable comparison operator: " + name);
- };
- }
-
- private static BinaryOp swap(BinaryOp op) {
- return switch (op) {
- case EQ, NOT_EQ -> op;
- case GT -> BinaryOp.LT;
- case GTE -> BinaryOp.LTE;
- case LT -> BinaryOp.GT;
- case LTE -> BinaryOp.GTE;
- default -> throw new IllegalArgumentException("not a comparison operator: " + op);
- };
- }
-
- private static boolean isPushableComparison(org.apache.spark.sql.connector.expressions.Expression[] children) {
- if (children.length != 2) {
- return false;
- }
- boolean lhsCol = isPushableFieldRef(children[0]);
- boolean lhsLit = isPushableLiteral(children[0]);
- boolean rhsCol = isPushableFieldRef(children[1]);
- boolean rhsLit = isPushableLiteral(children[1]);
- boolean lhsOk = lhsCol || lhsLit;
- boolean rhsOk = rhsCol || rhsLit;
- // We need at least one column reference; otherwise the predicate is comparing two
- // constants — Spark normally folds those, so we don't bother.
- return lhsOk && rhsOk && (lhsCol || rhsCol);
- }
-
- private static boolean isPushableFieldRef(org.apache.spark.sql.connector.expressions.Expression expr) {
- return expr instanceof NamedReference && ((NamedReference) expr).fieldNames().length >= 1;
- }
-
- private static boolean isFieldRefExpr(org.apache.spark.sql.connector.expressions.Expression expr) {
- return expr instanceof NamedReference;
- }
-
- /** Returns the Vortex column expression for a Spark named reference, walking nested struct fields. */
- private static Optional columnOf(org.apache.spark.sql.connector.expressions.Expression expr) {
- if (!(expr instanceof NamedReference)) {
- return Optional.empty();
- }
- String[] parts = ((NamedReference) expr).fieldNames();
- if (parts.length == 0) {
- return Optional.empty();
- }
- return Optional.of(Expression.column(parts));
- }
-
- private static Optional exprOf(org.apache.spark.sql.connector.expressions.Expression expr) {
- Optional col = columnOf(expr);
- if (col.isPresent()) {
- return col;
- }
- return literalOf(expr);
- }
-
- private static Optional stringValueOf(org.apache.spark.sql.connector.expressions.Expression expr) {
- if (!(expr instanceof Literal)) {
- return Optional.empty();
- }
- Object value = ((Literal) expr).value();
- if (value == null) {
- return Optional.empty();
- }
- if (value instanceof UTF8String) {
- return Optional.of(value.toString());
- }
- if (value instanceof CharSequence) {
- return Optional.of(value.toString());
- }
- return Optional.empty();
- }
-
- private static boolean isPushableStringLiteral(org.apache.spark.sql.connector.expressions.Expression expr) {
- return stringValueOf(expr).isPresent();
- }
-
- private static boolean isPushableLiteral(org.apache.spark.sql.connector.expressions.Expression expr) {
- if (!(expr instanceof Literal)) {
- return false;
- }
- Literal lit = (Literal) expr;
- DataType dataType = lit.dataType();
- // Null literals are pushable (we emit a typed null literal).
- if (lit.value() == null) {
- return dataType instanceof BooleanType
- || dataType instanceof ByteType
- || dataType instanceof ShortType
- || dataType instanceof IntegerType
- || dataType instanceof LongType
- || dataType instanceof FloatType
- || dataType instanceof DoubleType
- || dataType instanceof StringType
- || dataType instanceof BinaryType
- || dataType instanceof DateType
- || dataType instanceof TimestampType
- || dataType instanceof TimestampNTZType
- || dataType instanceof DecimalType;
- }
- return literalOf(expr).isPresent();
- }
-
- private static Optional literalOf(org.apache.spark.sql.connector.expressions.Expression expr) {
- if (!(expr instanceof Literal)) {
- return Optional.empty();
- }
- Literal lit = (Literal) expr;
- Object value = lit.value();
- DataType dataType = lit.dataType();
- return convertLiteral(value, dataType);
- }
-
- private static Optional convertLiteral(Object value, DataType dataType) {
- if (dataType instanceof BooleanType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteralBool());
- }
- return Optional.of(Expression.literal((Boolean) value));
- }
- if (dataType instanceof ByteType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteral(Expression.DType.I8));
- }
- return Optional.of(Expression.literal(((Number) value).byteValue()));
- }
- if (dataType instanceof ShortType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteral(Expression.DType.I16));
- }
- return Optional.of(Expression.literal(((Number) value).shortValue()));
- }
- if (dataType instanceof IntegerType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteral(Expression.DType.I32));
- }
- return Optional.of(Expression.literal(((Number) value).intValue()));
- }
- if (dataType instanceof LongType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteral(Expression.DType.I64));
- }
- return Optional.of(Expression.literal(((Number) value).longValue()));
- }
- if (dataType instanceof FloatType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteral(Expression.DType.F32));
- }
- return Optional.of(Expression.literal(((Number) value).floatValue()));
- }
- if (dataType instanceof DoubleType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteral(Expression.DType.F64));
- }
- return Optional.of(Expression.literal(((Number) value).doubleValue()));
- }
- if (dataType instanceof StringType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteral(Expression.DType.UTF8));
- }
- if (value instanceof UTF8String || value instanceof CharSequence) {
- return Optional.of(Expression.literal(value.toString()));
- }
- }
- if (dataType instanceof BinaryType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteral(Expression.DType.BINARY));
- }
- if (value instanceof byte[]) {
- return Optional.of(Expression.literal((byte[]) value));
- }
- }
- if (dataType instanceof DateType) {
- // Spark stores DateType as a 32-bit int day count since 1970-01-01.
- if (value == null) {
- return Optional.of(Expression.nullLiteralDate(TimeUnit.DAYS));
- }
- return Optional.of(Expression.literalDate(((Number) value).longValue(), TimeUnit.DAYS));
- }
- if (dataType instanceof TimestampType) {
- // Spark stores TimestampType as a 64-bit microseconds-since-epoch in UTC.
- if (value == null) {
- return Optional.of(Expression.nullLiteralTimestamp(TimeUnit.MICROSECONDS, "UTC"));
- }
- return Optional.of(Expression.literalTimestamp(((Number) value).longValue(), TimeUnit.MICROSECONDS, "UTC"));
- }
- if (dataType instanceof TimestampNTZType) {
- if (value == null) {
- return Optional.of(Expression.nullLiteralTimestamp(TimeUnit.MICROSECONDS, null));
- }
- return Optional.of(Expression.literalTimestamp(((Number) value).longValue(), TimeUnit.MICROSECONDS, null));
- }
- if (dataType instanceof DecimalType) {
- DecimalType decimalType = (DecimalType) dataType;
- int precision = decimalType.precision();
- int scale = decimalType.scale();
- if (value == null) {
- return Optional.of(Expression.nullLiteralDecimal(precision, scale));
- }
- BigInteger unscaled = unscaledValueOf(value, scale);
- if (unscaled == null) {
- return Optional.empty();
- }
- return Optional.of(Expression.literalDecimal(unscaled, precision, scale));
- }
- // Some Spark literals (e.g. NullType, GeographyType) have no Vortex representation.
- return Optional.empty();
- }
-
- /** Extract the unscaled integer value of a Spark decimal literal at the supplied {@code scale}. */
- private static BigInteger unscaledValueOf(Object value, int scale) {
- BigDecimal decimal;
- if (value instanceof Decimal) {
- decimal = ((Decimal) value).toJavaBigDecimal();
- } else if (value instanceof BigDecimal) {
- decimal = (BigDecimal) value;
- } else {
- return null;
- }
- try {
- return decimal.setScale(scale).unscaledValue();
- } catch (ArithmeticException ignored) {
- return null;
- }
- }
-}
diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java
deleted file mode 100644
index 198dfd6c77c..00000000000
--- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java
+++ /dev/null
@@ -1,103 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-package dev.vortex.spark.read;
-
-import com.google.common.collect.ImmutableMap;
-import dev.vortex.api.Session;
-import dev.vortex.jni.NativeFiles;
-import dev.vortex.spark.VortexFilePartition;
-import dev.vortex.spark.VortexSparkSession;
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
-import org.apache.spark.sql.connector.catalog.CatalogV2Util;
-import org.apache.spark.sql.connector.catalog.Column;
-import org.apache.spark.sql.connector.expressions.filter.Predicate;
-import org.apache.spark.sql.connector.read.Batch;
-import org.apache.spark.sql.connector.read.InputPartition;
-import org.apache.spark.sql.connector.read.PartitionReaderFactory;
-import org.apache.spark.sql.types.StructType;
-
-/** Execution source for batch scans of Vortex file tables. */
-public final class VortexBatchExec implements Batch {
- private final List paths;
- private final StructType readSchema;
- private final Map formatOptions;
- private final Predicate[] pushedPredicates;
- private List resolvedPaths;
-
- /**
- * Creates a new VortexBatchExec for scanning the specified Vortex files.
- *
- * @param paths the list of file paths to scan
- * @param columns the list of columns to read from the files
- * @param pushedPredicates predicates pushed down by Spark; converted to a single Vortex filter expression at read
- * time
- */
- public VortexBatchExec(
- List