From a0db09235c31e06924f9b0cec75eb2b5e48cebac Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 24 Aug 2026 13:57:44 +0100 Subject: [PATCH] Vortex is a fileformat not a tableprovider Add Hadoop-backed Java IO layer for Spark Signed-off-by: Robert Kruszewski --- .github/workflows/ci.yml | 6 +- .github/workflows/docs.yml | 2 +- .github/workflows/publish.yml | 2 +- README.md | 2 +- docs/developer-guide/integrations/spark.md | 81 ++- .../developer-guide/internals/architecture.md | 2 +- docs/user-guide/spark.md | 259 +++++---- java/README.md | 6 +- java/build.gradle.kts | 45 +- java/buildSrc/build.gradle.kts | 16 + .../kotlin/vortex-spark-module.gradle.kts | 268 +++++++++ java/gradle/libs.versions.toml | 12 +- java/settings.gradle.kts | 26 +- java/vortex-jni/build.gradle.kts | 7 +- .../main/java/dev/vortex/io/Closeables.java | 34 ++ .../java/dev/vortex/io/PooledReadable.java | 171 ++++++ java/vortex-spark/README.md | 133 ++--- java/vortex-spark/build.gradle.kts | 142 ----- .../spark/bench/BenchmarkSparkSession.java | 40 ++ .../spark/bench/FooterReadBenchmark.java | 98 ++++ .../spark/bench/SparkScanBenchmark.java | 101 ++++ .../java/dev/vortex/spark/ArrowUtils.java | 0 .../java/dev/vortex/spark/VortexOptions.java | 122 +++++ .../vortex/spark/VortexSessionProvider.java | 0 .../dev/vortex/spark/VortexSparkSession.java | 38 +- .../dev/vortex/spark/io/HadoopReadable.java | 82 +++ .../dev/vortex/spark/io/HadoopWritable.java | 51 ++ .../spark/io/SerializableHadoopConf.java | 41 ++ .../java/dev/vortex/spark/io/VortexFile.java | 70 +++ .../java/dev/vortex/spark/io/VortexIo.java | 84 +++ .../spark/read/PartitionColumnVectors.java | 79 +++ .../read/SparkFilterToVortexExpression.java | 401 ++++++++++++++ .../read/VortexAggregateReaderFactory.java | 126 +++++ .../spark/read/VortexArrowColumnVector.java | 0 .../vortex/spark/read/VortexFooterReader.java | 335 ++++++++++++ .../spark/read/VortexPartitionReader.java | 259 +++++++++ .../read/VortexPartitionReaderFactory.java | 68 +++ .../spark/write/SparkToArrowSchema.java | 31 ++ .../spark/write/VortexOutputWriter.java} | 266 ++++----- .../write/VortexOutputWriterFactory.java | 34 ++ ...pache.spark.sql.sources.DataSourceRegister | 0 ...ark.sql.sources.DataSourceRegister.license | 0 .../dev/vortex/spark/VortexDataSourceV2.scala | 45 ++ .../dev/vortex/spark/VortexFileFormat.scala | 220 ++++++++ .../scala/dev/vortex/spark/VortexTable.scala | 89 +++ .../dev/vortex/spark/read/VortexScan.scala | 139 +++++ .../vortex/spark/read/VortexScanBuilder.scala | 104 ++++ .../dev/vortex/spark/write/VortexWrite.scala | 33 ++ .../java/dev/vortex/spark/ArrowUtilsTest.java | 0 .../spark/VortexAggregatePushdownTest.java | 145 +++++ .../spark/VortexDataSourceBasicTest.java | 15 - .../VortexDataSourceInferSchemaTest.java | 10 +- .../spark/VortexDataSourceS3MockTest.java | 14 +- .../spark/VortexDataSourceStatsTest.java | 117 ++++ .../spark/VortexDataSourceWriteTest.java | 82 ++- .../vortex/spark/VortexFileExtensionTest.java | 132 +++++ .../spark/VortexFilterPushdownTest.java | 2 +- .../dev/vortex/spark/VortexOptionsTest.java | 96 ++++ .../vortex/spark/VortexProjectionTest.java | 115 ++++ .../vortex/spark/VortexSchemaMergeTest.java | 210 +++++++ .../java/dev/vortex/spark/VortexSqlTest.java | 42 +- .../vortex/spark/VortexV1FallbackTest.java | 146 +++++ .../vortex/spark/io/HadoopReadableTest.java | 177 ++++++ .../dev/vortex/spark/io/VortexIoTest.java | 78 +++ .../SparkFilterToVortexExpressionTest.java | 101 ++++ .../read/VortexArrowColumnVectorTest.java | 0 .../spark/write/SparkToArrowSchemaTest.java | 0 .../spark/write/VortexOutputWriterTest.java | 60 ++ .../java/dev/vortex/spark/VortexCatalog.java | 129 ----- .../dev/vortex/spark/VortexDataSourceV2.java | 249 --------- .../dev/vortex/spark/VortexFilePartition.java | 32 -- .../vortex/spark/VortexSessionCatalog.java | 85 --- .../java/dev/vortex/spark/VortexTable.java | 127 ----- .../dev/vortex/spark/config/HadoopUtils.java | 77 --- .../spark/config/VortexAzureProperties.java | 48 -- .../spark/config/VortexS3Properties.java | 77 --- .../vortex/spark/read/PartitionPathUtils.java | 110 ---- .../SparkPredicateToVortexExpression.java | 517 ------------------ .../vortex/spark/read/VortexBatchExec.java | 103 ---- .../spark/read/VortexPartitionReader.java | 176 ------ .../read/VortexPartitionReaderFactory.java | 58 -- .../dev/vortex/spark/read/VortexScan.java | 176 ------ .../vortex/spark/read/VortexScanBuilder.java | 196 ------- .../write/PartitionedVortexDataWriter.java | 473 ---------------- .../vortex/spark/write/VortexBatchWrite.java | 168 ------ .../spark/write/VortexDataWriterFactory.java | 85 --- .../spark/write/VortexWriteBuilder.java | 65 --- .../write/VortexWriterCommitMessage.java | 55 -- .../dev/vortex/spark/VortexCatalogTest.java | 114 ---- .../spark/VortexDataSourceStatsTest.java | 241 -------- .../spark/VortexSessionCatalogTest.java | 113 ---- .../dev/vortex/spark/VortexTableTest.java | 143 ----- .../vortex/spark/config/HadoopUtilsTest.java | 150 ----- .../spark/read/PartitionPathUtilsTest.java | 199 ------- .../SparkPredicateToVortexExpressionTest.java | 374 ------------- .../spark/read/VortexBatchExecTest.java | 95 ---- .../spark/read/VortexScanBuilderTest.java | 176 ------ .../PartitionedVortexDataWriterTest.java | 169 ------ java/vortex-spark/v3.5/build.gradle.kts | 6 + java/vortex-spark/v4.0/build.gradle.kts | 6 + 100 files changed, 5090 insertions(+), 5414 deletions(-) create mode 100644 java/buildSrc/build.gradle.kts create mode 100644 java/buildSrc/src/main/kotlin/vortex-spark-module.gradle.kts create mode 100644 java/vortex-jni/src/main/java/dev/vortex/io/Closeables.java create mode 100644 java/vortex-jni/src/main/java/dev/vortex/io/PooledReadable.java delete mode 100644 java/vortex-spark/build.gradle.kts create mode 100644 java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/BenchmarkSparkSession.java create mode 100644 java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/FooterReadBenchmark.java create mode 100644 java/vortex-spark/common/src/jmh/java/dev/vortex/spark/bench/SparkScanBenchmark.java rename java/vortex-spark/{ => common}/src/main/java/dev/vortex/spark/ArrowUtils.java (100%) create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/VortexOptions.java rename java/vortex-spark/{ => common}/src/main/java/dev/vortex/spark/VortexSessionProvider.java (100%) rename java/vortex-spark/{ => common}/src/main/java/dev/vortex/spark/VortexSparkSession.java (69%) create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopReadable.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/io/HadoopWritable.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/io/SerializableHadoopConf.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexFile.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/io/VortexIo.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/read/PartitionColumnVectors.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/read/SparkFilterToVortexExpression.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexAggregateReaderFactory.java rename java/vortex-spark/{ => common}/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java (100%) create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexFooterReader.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java rename java/vortex-spark/{ => common}/src/main/java/dev/vortex/spark/write/SparkToArrowSchema.java (82%) rename java/vortex-spark/{src/main/java/dev/vortex/spark/write/VortexDataWriter.java => common/src/main/java/dev/vortex/spark/write/VortexOutputWriter.java} (65%) create mode 100644 java/vortex-spark/common/src/main/java/dev/vortex/spark/write/VortexOutputWriterFactory.java rename java/vortex-spark/{ => common}/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister (100%) rename java/vortex-spark/{ => common}/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister.license (100%) create mode 100644 java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexDataSourceV2.scala create mode 100644 java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexFileFormat.scala create mode 100644 java/vortex-spark/common/src/main/scala/dev/vortex/spark/VortexTable.scala create mode 100644 java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScan.scala create mode 100644 java/vortex-spark/common/src/main/scala/dev/vortex/spark/read/VortexScanBuilder.scala create mode 100644 java/vortex-spark/common/src/main/scala/dev/vortex/spark/write/VortexWrite.scala rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/ArrowUtilsTest.java (100%) create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexAggregatePushdownTest.java rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/VortexDataSourceBasicTest.java (88%) rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/VortexDataSourceInferSchemaTest.java (84%) rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/VortexDataSourceS3MockTest.java (94%) create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java (88%) create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexFileExtensionTest.java rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/VortexFilterPushdownTest.java (99%) create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexOptionsTest.java create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexProjectionTest.java create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexSchemaMergeTest.java rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/VortexSqlTest.java (76%) create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/VortexV1FallbackTest.java create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/io/HadoopReadableTest.java create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/io/VortexIoTest.java create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/read/SparkFilterToVortexExpressionTest.java rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java (100%) rename java/vortex-spark/{ => common}/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java (100%) create mode 100644 java/vortex-spark/common/src/test/java/dev/vortex/spark/write/VortexOutputWriterTest.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/VortexFilePartition.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/VortexTable.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/config/HadoopUtils.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexAzureProperties.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/config/VortexS3Properties.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/read/PartitionPathUtils.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/read/SparkPredicateToVortexExpression.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/write/PartitionedVortexDataWriter.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriterFactory.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexWriteBuilder.java delete mode 100644 java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexWriterCommitMessage.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/VortexCatalogTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/VortexSessionCatalogTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/VortexTableTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/config/HadoopUtilsTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/read/PartitionPathUtilsTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/read/SparkPredicateToVortexExpressionTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexBatchExecTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java delete mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java create mode 100644 java/vortex-spark/v3.5/build.gradle.kts create mode 100644 java/vortex-spark/v4.0/build.gradle.kts diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b364f0e8aca..fee21976a1d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -549,9 +549,11 @@ jobs: - uses: ./.github/actions/setup-prebuild with: enable-sccache: "true" - - run: ./gradlew javadoc + - name: Generate Java API documentation + run: ./gradlew javadoc working-directory: ./java - - run: ./gradlew check + - name: Check all Spark and Scala variants + run: ./gradlew check working-directory: ./java license-check-and-audit-check: diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index df97e2bb731..66e0649d424 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -46,7 +46,7 @@ jobs: mkdir -p docs/_static/vortex-jni mkdir -p docs/_static/vortex-spark cp -r java/vortex-jni/build/docs/javadoc/* docs/_static/vortex-jni/ - cp -r java/vortex-spark/build/vortex-spark_2.13/docs/javadoc/* docs/_static/vortex-spark/ + cp -r java/vortex-spark/v4.0/build/vortex-spark-4.0_2.13/docs/javadoc/* docs/_static/vortex-spark/ - name: build Python and Rust docs run: | uv run --all-packages make -C docs html diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 2a9d4fbc8d6..3c426de4337 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -104,7 +104,7 @@ jobs: cp ../libvortex_jni_aarch64-apple-darwin.zip/libvortex_jni.dylib ./vortex-jni/src/main/resources/native/darwin-aarch64 cp ../libvortex_jni_aarch64-unknown-linux-gnu.zip/libvortex_jni.so ./vortex-jni/src/main/resources/native/linux-aarch64 cp ../libvortex_jni_x86_64-unknown-linux-gnu.zip/libvortex_jni.so ./vortex-jni/src/main/resources/native/linux-amd64 - - name: Build Java + - name: Build all Java release artifacts run: ./gradlew shadowJar - name: Publish to Maven Central run: ./gradlew -i publishAndReleaseToMavenCentral --no-configuration-cache diff --git a/README.md b/README.md index 24084df9992..39557814993 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ [![CodSpeed Badge](https://img.shields.io/endpoint?url=https://codspeed.io/badge.json)](https://codspeed.io/vortex-data/vortex) [![Crates.io](https://img.shields.io/crates/v/vortex.svg)](https://crates.io/crates/vortex) [![PyPI - Version](https://img.shields.io/pypi/v/vortex-data)](https://pypi.org/project/vortex-data/) -[![Maven - Version](https://img.shields.io/maven-central/v/dev.vortex/vortex-spark_2.13)](https://central.sonatype.com/artifact/dev.vortex/vortex-spark_2.13) +[![Maven - Version](https://img.shields.io/maven-central/v/dev.vortex/vortex-spark-4.0_2.13)](https://central.sonatype.com/artifact/dev.vortex/vortex-spark-4.0_2.13) [![codecov](https://codecov.io/github/vortex-data/vortex/graph/badge.svg)](https://codecov.io/github/vortex-data/vortex) [![Cite](https://img.shields.io/badge/cite-CITATION.cff-blue)](CITATION.cff) diff --git a/docs/developer-guide/integrations/spark.md b/docs/developer-guide/integrations/spark.md index 6758645c0dc..abb7c622528 100644 --- a/docs/developer-guide/integrations/spark.md +++ b/docs/developer-guide/integrations/spark.md @@ -1,57 +1,52 @@ # Spark -The `vortex-spark` connector implements Apache Spark's DataSource V2 API, allowing Spark to read -and write Vortex files as a native data source registered under the format name `vortex`. +The `vortex-spark` connector is built on Spark's file-source DataSource V2 framework. Shared +sources compile into Spark 3.5/Scala 2.12, Spark 3.5/Scala 2.13, and Spark 4.0/Scala 2.13 +artifacts; the Spark 4 artifact is also tested on Spark 4.1. -## Registration +## File-source integration -The connector implements Spark's `TableProvider` and `DataSourceRegister` interfaces. When a -query references the `vortex` format, Spark creates a `VortexTable` that supports both batch -reads (`SupportsRead`) and writes (`SupportsWrite`). Schema inference reads the footer of a -discovered file to extract the Arrow schema and map it to Spark's schema representation. +`VortexDataSourceV2` extends `FileDataSourceV2`, while `VortexTable`, `VortexScan`, and +`VortexWrite` use Spark's `FileTable`, `FileScan`, and `FileWrite` abstractions. Spark therefore +owns file listing, partition discovery and pruning, input bin-packing, output commit, and +overwrite behavior. `VortexFileFormat` supplies the functional V1 fallback used by catalog +tables and direct path queries. -## Multiple Files +Vortex files are not internally split. A Spark file partition may contain several files, and the +reader factory opens each `PartitionedFile` in turn. Hive partition values are appended with +constant column vectors. -Spark's scan builder enumerates Vortex files by scanning the provided path. If the path is a -directory, native code lists all `.vortex` files within it. Each file becomes an independent -input partition, and Spark's task scheduler distributes partitions across executors in the -cluster. +## I/O and JNI -Each partition creates its own file handle and scan state, so there is no shared mutable state -between partitions. This maps naturally to Spark's execution model where each task runs -independently on a separate JVM thread. +Spark lists paths through Hadoop. Content reads go through pooled Hadoop input streams exposed to +native Vortex through the JNI `NativeReadable` interface. Writes expose the committer's Hadoop +task path through `NativeWritable`. Vortex's own object-store clients are not used, so the +connector sees the same schemes and credentials as Spark's file index and commit protocol. -## Threading Model +Native arrays cross into Spark through the Arrow C Data Interface. Each partition reader owns its +native scan, Arrow allocator, and exported batches and closes them at task completion. -The Spark integration crosses the JNI boundary between Java and Rust. Each Spark partition -reader opens a Vortex file and creates a native scan via JNI. The native side manages its own -async runtime and drives I/O internally, returning results to Java as Arrow-compatible columnar -batches. +## Pushdown and statistics -Because each partition reader owns its native resources exclusively, there is no contention -across Spark threads. The JNI boundary is crossed once per batch rather than once per row, -keeping overhead low. A prefetching iterator on the Java side buffers upcoming batches to -overlap I/O with Spark's processing. +Schema inference merges every footer, so the dataset schema is the union of the top-level fields its +files carry. The partition reader projects only the fields the file it opened actually holds and +fills the rest with constant null vectors, and it converts filters against those same fields, so a +filter on a column the file lacks stays a Spark residual rather than reaching the native scan. -## Filter and Projection Pushdown +Spark's required schema becomes the Vortex scan projection. Convertible V1 filters become Vortex +expressions; filters the converter rejects remain Spark residuals. -Projection pushdown is supported through Spark's `SupportsPushDownRequiredColumns` interface. -The scan builder prunes the column list to only those referenced by the query, and the pruned -column set is passed to the native scan via `ScanOptions`. +The scan accepts `COUNT(*)` aggregation when there are no data filters and grouping uses only +partition columns, and only from footers that state their row count exactly. Readers return one +footer count per file and Spark performs the final merge; the scan then reports one row per file as +its statistics rather than reading those footers again on the driver. Footer row counts from files +left after partition pruning are reported through Spark scan statistics when no aggregate is +pushed. MIN/MAX and `COUNT(column)` need read-side column statistics in the JNI API and are +not pushed down. -Filter pushdown is supported through `SupportsPushDownV2Filters`. The scan builder converts each -predicate it recognizes into a Vortex expression and keeps the rest for Spark to evaluate after the -scan; the converted expression reaches native code as the `ScanOptions` filter. +## Source layout -## Data Export - -Native Vortex arrays are exported to Arrow via the C Data Interface, then wrapped in Spark's -columnar batch format using custom `ArrowColumnVector` wrappers. This avoids a copy between -the Rust and JVM heaps -- the Arrow buffers remain in native memory and are accessed from Java -through direct byte buffers. - -## Future Work - -The current integration builds directly on the native file and scan APIs via JNI. Future work -will migrate it to use the [Scan API](/concepts/scanning) `Source` trait, which will provide a -standard interface for file discovery, partitioning, and pushdown. +Shared Java and Scala live in `java/vortex-spark/common`. Thin projects in `v3.5` and `v4.0` +select the corresponding Spark and Scala dependencies. The Scala shims isolate the binary +differences in Spark's Scala APIs; nothing in Java depends on them, so javac compiles the Java +sources and ErrorProne and Nopen keep covering them. diff --git a/docs/developer-guide/internals/architecture.md b/docs/developer-guide/internals/architecture.md index 3354c25cd85..251b828bbd0 100644 --- a/docs/developer-guide/internals/architecture.md +++ b/docs/developer-guide/internals/architecture.md @@ -80,7 +80,7 @@ Query engine integrations allow Vortex files to be queried through existing anal |----------------------------------| ---------- |----------------------------------------------| | `vortex-datafusion/` | DataFusion | `TableProvider` and `FileFormat` integration | | `vortex-duckdb/` | DuckDB | Table function integration | -| `java/vortex-spark_{2.12,2.13}/` | Spark | Spark DataSource V2 connector via JNI | +| `java/vortex-spark/` | Spark | Versioned Spark file-source connector via JNI | | `java/vortex-trino/` | Trino | Trino connector (in development) | ## Other Crates diff --git a/docs/user-guide/spark.md b/docs/user-guide/spark.md index 03f0760a8c4..76573b82d2a 100644 --- a/docs/user-guide/spark.md +++ b/docs/user-guide/spark.md @@ -1,53 +1,48 @@ # Spark -Vortex provides a Spark DataSource V2 connector for reading and writing Vortex files. The -connector is published to Maven Central in two flavors: +Vortex provides a Spark file data source for reading and writing Vortex files. Choose the +artifact matching both the Spark and Scala versions in the application: -- `dev.vortex:vortex-spark_2.13` for Spark 4.x (Scala 2.13) -- `dev.vortex:vortex-spark_2.12` for Spark 3.5.x (Scala 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.78.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 for Linux +(x86_64 and aarch64) and macOS (aarch64), and relocated dependencies that avoid conflicts with +Spark. The unclassified thin JAR is not usable by itself. -## Getting Vortex into Spark +## Installation -For `spark-shell`, `spark-submit`, or `pyspark`, pass the `all` JAR with `--jars`. Spark -accepts either a local path or a URL, so you can point directly at Maven Central: +Pass the classified JAR to `spark-shell`, `spark-submit`, or `pyspark` with `--jars`. For +example, for connector version `VERSION`: ```shell -spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark_2.13/0.78.0/vortex-spark_2.13-0.78.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 equivalently when building a session programmatically, e.g. in PySpark: +Or configure the JAR on a PySpark session: ```python spark = ( SparkSession.builder - .config("spark.jars", "/path/to/vortex-spark_2.13-0.78.0-all.jar") + .config("spark.jars", "/path/to/vortex-spark-4.0_2.13-VERSION-all.jar") .getOrCreate() ) ``` ```{note} -`--packages dev.vortex:vortex-spark_2.13:0.78.0` does not work: `--packages` cannot select -the `all` classifier and resolves the thin JAR, which fails at runtime with -`NoClassDefFoundError: dev/vortex/relocated/...`. +Spark's `--packages` option cannot select the `all` classifier. It resolves the thin JAR, +which fails at runtime because the relocated dependencies are only in the classified JAR. ``` -Once the JAR is on the classpath, the connector registers itself automatically under the -format name `vortex` — no session configuration is required. - -## Installation as a Build Dependency - -To depend on the connector from a JVM project, add the `all` classifier to the dependency: +For a JVM build, specify the classifier explicitly. Gradle (Kotlin): ```kotlin -implementation("dev.vortex:vortex-spark_2.13:0.78.0:all") +implementation("dev.vortex:vortex-spark-4.0_2.13:VERSION:all") ``` Maven: @@ -55,146 +50,186 @@ Maven: ```xml dev.vortex - vortex-spark_2.13 - 0.78.0 + vortex-spark-4.0_2.13 + VERSION all ``` +The connector registers itself as `vortex`; no session extension or catalog configuration is +required. + ## Reading Vortex Files -Paths may be local filesystem paths (`/path/to/data`) or URLs (`file:///path/to/data`, -`s3://bucket/path/to/data`). Use the `vortex` format to read a single file or a directory of -Vortex files: +Read a file, directory, or set of paths with the DataFrame API: ```java Dataset df = spark.read() .format("vortex") - .option("path", "/path/to/data.vortex") - .load(); + .load("/path/to/data"); ``` -When pointed at a directory, the connector discovers all `.vortex` files and creates one read -partition per file. +Spark's file-source framework provides recursive file listing, split bin-packing, Hive-style +partition discovery and pruning, and the standard `pathGlobFilter` and +`recursiveFileLookup` options. As with other Spark file formats, hidden files whose names begin +with `_` or `.` are ignored. Vortex files are not split internally, and only required columns +are read. + +Every file in a Vortex dataset must end with `.vortex`. Writes produce that extension. A dataset +that holds any other file is rejected: schema inference reports that it found no Vortex file, and +a scan names the offending path. Use `pathGlobFilter` to read Vortex files out of a directory that +holds other things too. + +### Schema Inference and Merging -Column pruning is pushed down — only the columns referenced by the query are read from the file. +The schema of a dataset is the merge of every file's footer schema, so a column added by a later +write is part of the dataset and the files written before it read as null. A field that only some +files carry is nullable in the merged schema. + +Merging reads one footer per file on the driver before the job starts. Set `vortex.mergeSchema` to +`false` to read a single footer and let one file's schema stand for the whole dataset, which is +worth doing for a large dataset of uniform files. Passing an explicit `.schema(...)` skips +inference altogether. + +Only top-level columns are merged. A struct column that gained or lost a field cannot be merged, +because the reader projects a struct as the file stores it and cannot widen one file's struct to +match another's; inference fails and names the field. Two files that give the same column +different types fail the same way. + +Supported filters are pushed into the Vortex scan. Nested filter pushdown through the V1 +fallback also requires `vortex` in +`spark.sql.optimizer.nestedPredicatePushdown.supportedFileSources`. A filter reading a column +that some file does not carry is evaluated by Spark above the scan rather than pushed into that +file. + +`COUNT(*)` without data filters is computed from file footers. Spark combines the partial count +from each file, including counts grouped by partition columns. Footer row counts also feed +Spark's scan statistics unless `vortex.stats.rowCount` is disabled. ## Writing Vortex Files ```java df.write() .format("vortex") - .option("path", "/path/to/output") .mode(SaveMode.Overwrite) - .save(); + .save("/path/to/output"); ``` -Each Spark partition produces one output file named `part-{partitionId}-{taskId}.vortex`. - -### Write Options +Spark's commit protocol owns output naming, task retries, append, truncate, and static or dynamic +partition overwrite. Partitioned writes use the normal DataFrame API: -| Option | Default | Description | -|---------------------------|---------|------------------------------------| -| `vortex.write.batch.size` | 2048 | Number of rows per batch (1–65536) | - -### Save Modes +```java +df.write() + .format("vortex") + .partitionBy("date") + .mode(SaveMode.Overwrite) + .save("/path/to/output"); +``` -The connector supports all standard Spark save modes: `Overwrite`, `Append`, `Ignore`, and -`ErrorIfExists`. +| Option | Default | Description | +|---------------------------|---------|-------------| +| `batch.size` | `2048` | Rows buffered per write batch; range 1–65536. | +| `vortex.write.batch.size` | — | Rows buffered per write batch for Vortex alone, overriding `batch.size`. | ## Spark SQL -The connector can also be used from pure SQL. To query existing Vortex files, register them -as a temporary view: +Vortex works as a native file format on Spark 3.5 and 4.x: ```sql CREATE TEMPORARY VIEW people USING vortex OPTIONS (path '/path/to/data'); -SELECT name, age FROM people WHERE age > 30; -``` - -Tables can be created with `USING vortex`, then written to and read back with plain SQL. -With a `LOCATION` clause the table is external, backed by the files at that path; without -one the table is managed, and Spark stores its data under the warehouse directory (and -deletes it on `DROP TABLE`): - -```sql CREATE TABLE student (id INT, name STRING, age INT) USING vortex; INSERT INTO student VALUES (1, 'Alice', 20), (2, 'Bob', 21); - SELECT * FROM student; ``` -`CREATE TABLE ... AS SELECT` works the same way: +With a `LOCATION` clause, the table is external. Without one, Spark stores it below the warehouse +directory and deletes its files on `DROP TABLE`. `CREATE TABLE ... AS SELECT` is also supported. + +Existing files can be queried directly without registering a catalog: ```sql -CREATE TABLE adults -USING vortex -AS SELECT * FROM people WHERE age >= 18; +SELECT * FROM vortex.`/path/to/data`; ``` -```{note} -On Spark 3.5, `CREATE TABLE ... USING vortex` additionally requires replacing the session -catalog, because Spark 3.5's built-in catalog cannot read tables backed by a DataSource -V2-only connector: +## Choosing the V2 or V1 Code Path - spark.sql.catalog.spark_catalog=dev.vortex.spark.VortexSessionCatalog +By default Vortex reads and writes through Spark's DataSource V2 file API. Catalog tables use the V1 +file format instead, and so does every path when `vortex` is listed in +`spark.sql.sources.useV1SourceList`: -The extension delegates everything to the built-in session catalog (including the Hive -metastore, if configured) and only changes how `vortex` tables are resolved; tables of other -providers are untouched. It is not needed on Spark 4, though setting it is harmless. +```python +spark.conf.set("spark.sql.sources.useV1SourceList", "vortex") ``` -## Direct File Queries +Spark lists all of its own file formats there by default. Two features are only available on the V1 +path: -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 in the session configuration under the name `vortex`: +- the `_metadata` column, which Spark's V2 file API does not expose; +- dynamic partition pruning, which Spark applies to V1 relations only. -```shell -spark-sql --conf spark.sql.catalog.vortex=dev.vortex.spark.VortexCatalog -``` +Filter pushdown, column pruning, partition pruning, and partitioned writes work on both paths. +`COUNT(*)` footer pushdown and footer-backed scan statistics are V2 only. -Then query a Vortex file, or a directory of Vortex files, directly by path — no view or -table required: +## I/O and Remote Storage -```sql -SELECT * FROM vortex.`/path/to/data`; +Spark lists files and discovers partitions through Hadoop, and all file content is read and +written through Hadoop streams, so the connector inherits Spark's filesystem implementations, +credentials, and retry behavior. Configure remote storage through Spark's Hadoop configuration +and use the matching Hadoop scheme, for example `s3a://` for S3 and `abfs://` for Azure. Vortex's +own storage clients are not used by the Spark connector. -INSERT INTO vortex.`/path/to/data` VALUES (1, 'Alice', 20); -``` +| Option | Default | Description | +|----------------------------|---------|-------------| +| `vortex.readConcurrency` | `0` | Maximum Hadoop read upcalls per file; `0` uses the native default. | +| `vortex.workerThreads` | `4` | Background threads driving Vortex futures. JVM-wide, and set once per JVM by the first read or write. | +| `vortex.session.provider` | — | Fully-qualified name of a `VortexSessionProvider` supplying a custom session on the driver and on every executor. | +| `vortex.aggregatePushdown` | `true` | Answer `COUNT(*)` from file footers. | +| `vortex.mergeSchema` | `true` | Merge every file's footer schema, rather than reading one footer. | +| `vortex.stats.rowCount` | `true` | Read footer row counts for Spark statistics. | +| `vortex.footerParallelism` | `8` | Maximum concurrent footer reads, for schema merging and statistics. | +| `vortex.stats.maxFiles` | `1000` | Skip footer row counts above this file count; `0` removes the bound. | -## Supported Types +Option names are matched without regard to case, as they are everywhere else in Spark. -| Spark Type | Vortex Type | -|--------------------|----------------------------------------| -| `BooleanType` | Bool | -| `ByteType` | Int8 / UInt8 | -| `ShortType` | Int16 / UInt16 | -| `IntegerType` | Int32 / UInt32 | -| `LongType` | Int64 / UInt64 | -| `FloatType` | Float32 | -| `DoubleType` | Float64 | -| `StringType` | Utf8 | -| `BinaryType` | Binary | -| `DecimalType` | Decimal | -| `DateType` | Date (days) | -| `TimestampType` | Timestamp (microseconds, UTC) | -| `TimestampNTZType` | Timestamp (microseconds, no timezone) | -| `ArrayType` | List | -| `StructType` | Struct | - -## S3 Support - -The connector supports reading and writing to S3 paths: +Scan statistics cost one read per file on the driver, before the job starts. `vortex.stats.maxFiles` +keeps planning bounded on a large dataset by reporting no row count instead. Raise it when exact +row counts matter more than planning time, or set `vortex.stats.rowCount` to `false` to stop +reading footers at all. -```java -Dataset df = spark.read() - .format("vortex") - .option("path", "s3://bucket/path/to/data") - .load(); -``` +## Supported Types + +| Spark Type | Vortex Type | +|--------------------|---------------------------------------| +| `BooleanType` | Bool | +| `ByteType` | Int8 / UInt8 | +| `ShortType` | Int16 / UInt16 | +| `IntegerType` | Int32 / UInt32 | +| `LongType` | Int64 / UInt64 | +| `FloatType` | Float32 | +| `DoubleType` | Float64 | +| `StringType` | Utf8 | +| `BinaryType` | Binary | +| `DecimalType` | Decimal | +| `DateType` | Date (days) | +| `TimestampType` | Timestamp (microseconds, UTC) | +| `TimestampNTZType` | Timestamp (microseconds, no timezone) | +| `ArrayType` | List | +| `StructType` | Struct | +| `MapType` | Map | + +## Migrating from earlier artifacts + +- Replace `vortex-spark_2.12` or `vortex-spark_2.13` with the Spark-versioned artifact listed + above. +- Remove `spark.sql.catalog.spark_catalog=dev.vortex.spark.VortexSessionCatalog` and + `spark.sql.catalog.vortex=dev.vortex.spark.VortexCatalog`. The file-source integration makes + both catalogs unnecessary, and their classes were removed. +- Listing now follows Spark semantics. Hidden files are skipped, and `pathGlobFilter` and + `recursiveFileLookup` are available. +- Remote storage requires a Hadoop filesystem connector. OpenDAL-only schemes such as `cos://` + and `oss://` are not available through Spark's file index. diff --git a/java/README.md b/java/README.md index f45f2474f11..abfaafa75ad 100644 --- a/java/README.md +++ b/java/README.md @@ -9,11 +9,12 @@ We provide two interfaces for working with Vortex from Java: ## Publishing -We publish three artifacts out of this repo at CI time to Maven Central Sonatype: +We publish the JNI artifacts and three Spark connector variants from this repo to Maven Central: * `vortex-jni` JAR containing the JNI code, plus compiled native libraries for all of the following targets: `aarch64-apple-darwin`, `aarch64-unknown-linux-gnu`, `x86_64-unknown-linux-gnu` * `vortex-jni-all` which is the "shadow JAR" containing all of `vortex-jni` as well as all upstream Java dependencies packaged in a single JAR. -* `vortex-spark` which is the runtime JAR needed for the Vortex Spark bindings +* `vortex-spark-3.5_2.12`, `vortex-spark-3.5_2.13`, and `vortex-spark-4.0_2.13`, which are the + Spark-versioned connector artifacts We use the [following GPG key](https://keyserver.ubuntu.com/pks/lookup?search=8745D1A87C0B2159&fingerprint=on&op=index) for publishing: @@ -35,4 +36,3 @@ vWBCujQBRqlcCGIIawcI ``` The private key and passphrase for the publish key are owned by the Vortex Dev Team. - diff --git a/java/build.gradle.kts b/java/build.gradle.kts index 8afda976316..95fc114e22f 100644 --- a/java/build.gradle.kts +++ b/java/build.gradle.kts @@ -8,7 +8,29 @@ plugins { id("com.palantir.git-version") version "5.0.0" id("com.palantir.java-format") version "2.93.0" id("net.ltgt.errorprone") version "5.1.0" apply false - id("com.vanniktech.maven.publish") version "0.36.0" apply false +} + +spotless { + java { + target(fileTree("vortex-spark/common") { include("**/*.java") }) + palantirJavaFormat().formatJavadoc(true) + licenseHeaderFile("${rootProject.projectDir}/.spotless/java-license-header.txt") + removeUnusedImports() + forbidWildcardImports() + importOrder("") + trimTrailingWhitespace() + leadingTabsToSpaces(4) + targetExclude("**/generated/**") + targetExcludeIfContentContains("// spotless:disabled") + } + scala { + target(fileTree("vortex-spark/common") { include("**/*.scala") }) + scalafmt("3.9.10") + licenseHeaderFile( + "${rootProject.projectDir}/.spotless/java-license-header.txt", + "package ", + ) + } } subprojects { @@ -39,6 +61,10 @@ allprojects { spotless { java { + if (project.name.startsWith("vortex-spark-")) { + // Shared sources are formatted by the root project, where they are inside the project directory. + target(project.fileTree("src") { include("**/*.java") }) + } palantirJavaFormat().formatJavadoc(true) licenseHeaderFile("${rootProject.projectDir}/.spotless/java-license-header.txt") removeUnusedImports() @@ -54,6 +80,8 @@ allprojects { tasks.withType { options.errorprone.disable("UnusedVariable") options.errorprone.disableWarningsInGeneratedCode = true + // JMH generates non-final subclasses of every benchmark, which ErrorProne's Nopen check rejects. + options.errorprone.enabled.set(name != "compileBenchmarkJava") options.release = 17 options.compilerArgs.add("-Werror") @@ -70,6 +98,9 @@ allprojects { } tasks["check"].dependsOn("spotlessCheck") + if (project.name == "vortex-spark-4.0_2.13") { + tasks["check"].dependsOn(rootProject.tasks.named("spotlessCheck")) + } } spotless { @@ -78,10 +109,16 @@ allprojects { } } - if (project.name == "vortex-spark_2.12") { - // vortex-spark_2.12 and vortex-spark_2.13 share a projectDir; format from the 2.13 variant only. + if (project.name.startsWith("vortex-spark-") && project.name != "vortex-spark-4.0_2.13") { + // Spark variants share sources. Format them from the root project only. tasks.register("format") { enabled = false } } else { - tasks.register("format").get().dependsOn("spotlessApply") + tasks.register("format").get().dependsOn( + if (project.name == "vortex-spark-4.0_2.13") { + rootProject.tasks.named("spotlessApply") + } else { + tasks.named("spotlessApply") + }, + ) } } diff --git a/java/buildSrc/build.gradle.kts b/java/buildSrc/build.gradle.kts new file mode 100644 index 00000000000..0240c785439 --- /dev/null +++ b/java/buildSrc/build.gradle.kts @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +plugins { + `kotlin-dsl` +} + +repositories { + gradlePluginPortal() + mavenCentral() +} + +dependencies { + implementation("com.gradleup.shadow:shadow-gradle-plugin:9.4.2") + implementation("com.vanniktech.maven.publish:com.vanniktech.maven.publish.gradle.plugin:0.36.0") +} diff --git a/java/buildSrc/src/main/kotlin/vortex-spark-module.gradle.kts b/java/buildSrc/src/main/kotlin/vortex-spark-module.gradle.kts new file mode 100644 index 00000000000..85ff26e1559 --- /dev/null +++ b/java/buildSrc/src/main/kotlin/vortex-spark-module.gradle.kts @@ -0,0 +1,268 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar +import org.gradle.api.tasks.javadoc.Javadoc +import org.gradle.api.tasks.scala.ScalaCompile + +plugins { + scala + `java-library` + `jvm-test-suite` + id("com.gradleup.shadow") + id("com.vanniktech.maven.publish") +} + +val libs = extensions.getByType().named("libs") +val sparkLine = project.name.removePrefix("vortex-spark-").substringBefore('_') +val scalaBinaryVersion = project.name.substringAfterLast('_') +val sparkVersion = + when (sparkLine) { + "3.5" -> libs.findVersion("spark35").get().requiredVersion + "4.0" -> libs.findVersion("spark40").get().requiredVersion + else -> throw GradleException("Unsupported Spark line: $sparkLine") + } +val scalaVersion = + when (sparkLine to scalaBinaryVersion) { + "3.5" to "2.12" -> libs.findVersion("scala212").get().requiredVersion + "3.5" to "2.13" -> libs.findVersion("scala213Spark35").get().requiredVersion + "4.0" to "2.13" -> libs.findVersion("scala213Spark40").get().requiredVersion + else -> throw GradleException("Unsupported Spark/Scala variant: $sparkLine/$scalaBinaryVersion") + } +val spark41Version = libs.findVersion("spark41").get().requiredVersion + +layout.buildDirectory = layout.projectDirectory.dir("build/${project.name}") + +val commonSourceDir = rootProject.file("vortex-spark/common/src") +val versionSourceDir = project.file("src") + +sourceSets { + named("main") { + // Java stays with javac so ErrorProne, Nopen, and -Werror keep covering it. Only Scala + // depends on Java here, and compileScala already sees compileJava's output. + java.setSrcDirs( + listOf( + commonSourceDir.resolve("main/java"), + versionSourceDir.resolve("main/java"), + ), + ) + scala.setSrcDirs( + listOf( + commonSourceDir.resolve("main/scala"), + versionSourceDir.resolve("main/scala"), + ), + ) + resources.setSrcDirs( + listOf( + commonSourceDir.resolve("main/resources"), + versionSourceDir.resolve("main/resources"), + ), + ) + } + named("test") { + java.setSrcDirs( + listOf( + commonSourceDir.resolve("test/java"), + versionSourceDir.resolve("test/java"), + ), + ) + scala.setSrcDirs( + listOf( + commonSourceDir.resolve("test/scala"), + versionSourceDir.resolve("test/scala"), + ), + ) + resources.setSrcDirs( + listOf( + commonSourceDir.resolve("test/resources"), + versionSourceDir.resolve("test/resources"), + ), + ) + } +} + +dependencies { + compileOnly("org.scala-lang:scala-library:$scalaVersion") + compileOnly("org.apache.spark:spark-catalyst_$scalaBinaryVersion:$sparkVersion") + compileOnly("org.apache.spark:spark-sql_$scalaBinaryVersion:$sparkVersion") + api(project(":vortex-jni", configuration = "shadow")) + + implementation(libs.findLibrary("guava").get()) + implementation(libs.findLibrary("slf4j-api").get()) +} + +tasks.withType().configureEach { + scalaCompileOptions.additionalParameters = listOf("-release:17", "-deprecation", "-feature") +} + +testing { + suites { + val test by getting(JvmTestSuite::class) { + useJUnitJupiter() + dependencies { + implementation(libs.findLibrary("junit-jupiter").get()) + implementation("org.apache.spark:spark-core_$scalaBinaryVersion:$sparkVersion") + implementation("org.apache.spark:spark-sql_$scalaBinaryVersion:$sparkVersion") + implementation(libs.findLibrary("s3mock-testcontainers").get()) + implementation(libs.findLibrary("testcontainers-juputer").get()) + runtimeOnly(libs.findLibrary("slf4j-simple").get()) + if (sparkLine == "3.5") { + runtimeOnly("javax.servlet:javax.servlet-api:4.0.1") + } + } + } + + // Benchmarks are compiled by javac so the JMH annotation processor can generate its + // benchmark registry. Joint scalac compilation would skip that step. + register("benchmark") { + sources { + java.setSrcDirs( + listOf( + commonSourceDir.resolve("jmh/java"), + versionSourceDir.resolve("jmh/java"), + ), + ) + resources.setSrcDirs(emptyList()) + } + dependencies { + implementation(project()) + implementation(libs.findLibrary("jmh-core").get()) + annotationProcessor(libs.findLibrary("jmh-generator-annprocess").get()) + implementation("org.apache.spark:spark-core_$scalaBinaryVersion:$sparkVersion") + implementation("org.apache.spark:spark-sql_$scalaBinaryVersion:$sparkVersion") + runtimeOnly(libs.findLibrary("slf4j-simple").get()) + if (sparkLine == "3.5") { + runtimeOnly("javax.servlet:javax.servlet-api:4.0.1") + } + } + // The suite only carries JMH sources; its default test task has nothing to run. + targets.all { testTask.configure { enabled = false } } + } + + if (sparkLine == "4.0") { + register("spark41CompatTest") { + useJUnitJupiter() + sources { + java.setSrcDirs( + listOf( + commonSourceDir.resolve("test/java"), + versionSourceDir.resolve("test/java"), + ), + ) + resources.setSrcDirs( + listOf( + commonSourceDir.resolve("test/resources"), + versionSourceDir.resolve("test/resources"), + ), + ) + } + dependencies { + implementation(project()) + implementation(libs.findLibrary("junit-jupiter").get()) + implementation("org.apache.spark:spark-core_2.13:$spark41Version") + implementation("org.apache.spark:spark-sql_2.13:$spark41Version") + implementation(libs.findLibrary("s3mock-testcontainers").get()) + implementation(libs.findLibrary("testcontainers-juputer").get()) + runtimeOnly(libs.findLibrary("slf4j-simple").get()) + } + } + } + } +} + +if (sparkLine == "4.0") { + tasks.named("check") { + dependsOn("spark41CompatTest") + } +} + +mavenPublishing { + coordinates( + groupId = "dev.vortex", + artifactId = project.name, + version = rootProject.version.toString(), + ) + publishToMavenCentral() + if (!project.hasProperty("skip.signing")) { + signAllPublications() + } + repositories { + mavenCentral() + mavenLocal() + } + pom { + name = project.name + 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" + } + } +} + +tasks.withType().configureEach { + relocate("com.google.common", "dev.vortex.relocated.com.google.common") + relocate("org.apache.arrow", "dev.vortex.relocated.org.apache.arrow") { + exclude("org.apache.arrow.c.jni.JniWrapper") + exclude("org.apache.arrow.c.jni.PrivateData") + exclude("org.apache.arrow.c.jni.CDataJniException") + exclude("org.apache.arrow.c.ArrayStreamExporter\$ExportedArrayStreamPrivateData") + } + relocate("com.fasterxml.jackson", "dev.vortex.relocated.com.fasterxml.jackson") +} + +val sparkJvmArgs = + listOf( + "--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.withType().configureEach { + classpath += project(":vortex-jni").tasks.named("shadowJar").get().outputs.files + jvmArgs(sparkJvmArgs) +} + +tasks.register("jmh") { + description = "Run JMH benchmarks. Pass JMH arguments with -PjmhArgs=\" -f 1 -wi 2 -i 3\"." + group = "verification" + + val benchmarkSourceSet = sourceSets.named("benchmark").get() + dependsOn(benchmarkSourceSet.classesTaskName) + classpath = benchmarkSourceSet.runtimeClasspath + project(":vortex-jni").tasks.named("shadowJar").get().outputs.files + mainClass = "org.openjdk.jmh.Main" + jvmArgs(sparkJvmArgs) + val jmhArgs = project.findProperty("jmhArgs")?.toString() ?: "" + args(jmhArgs.split(' ').filter { it.isNotBlank() }) +} + +tasks.withType().configureEach { + setSource( + files( + fileTree(commonSourceDir.resolve("main/java")) { include("**/*.java") }, + fileTree(versionSourceDir.resolve("main/java")) { include("**/*.java") }, + ), + ) +} + +tasks.named("build") { + dependsOn("shadowJar") +} + +description = "Apache Spark $sparkLine bindings for reading and writing Vortex file datasets" diff --git a/java/gradle/libs.versions.toml b/java/gradle/libs.versions.toml index 0b817553b96..4d9dabaa895 100644 --- a/java/gradle/libs.versions.toml +++ b/java/gradle/libs.versions.toml @@ -6,14 +6,20 @@ arrow = "19.0.0" errorprone = "2.36.0" guava = "33.6.0-jre" immutables = "2.12.2" +jmh = "1.37" junit-jupiter = "6.1.0" logback = "1.5.34" netty = "4.2.15.Final" nopen = "1.0.1" roaringbitmap = "1.6.14" slf4j = "2.0.18" -spark3 = "3.5.9" -spark4 = "4.1.2" +scala212 = "2.12.18" +scala213Spark35 = "2.13.8" +scala213Spark40 = "2.13.16" +scala213Spark41 = "2.13.17" +spark35 = "3.5.9" +spark40 = "4.0.4" +spark41 = "4.1.2" s3mock = "5.0.0" testcontainers-jupiter = "1.21.4" @@ -25,6 +31,8 @@ errorprone-annotations = { module = "com.google.errorprone:error_prone_annotatio errorprone-core = { module = "com.google.errorprone:error_prone_core", version.ref = "errorprone" } guava = { module = "com.google.guava:guava", version.ref = "guava" } immutables-value = { module = "org.immutables:value", version.ref = "immutables" } +jmh-core = { module = "org.openjdk.jmh:jmh-core", version.ref = "jmh" } +jmh-generator-annprocess = { module = "org.openjdk.jmh:jmh-generator-annprocess", version.ref = "jmh" } junit-jupiter = { module = "org.junit.jupiter:junit-jupiter", version.ref = "junit-jupiter" } junit-jupiter-params = { module = "org.junit.jupiter:junit-jupiter-params", version.ref = "junit-jupiter" } logback-classic = { module = "ch.qos.logback:logback-classic", version.ref = "logback" } diff --git a/java/settings.gradle.kts b/java/settings.gradle.kts index a601cfa9488..8c4f8e83895 100644 --- a/java/settings.gradle.kts +++ b/java/settings.gradle.kts @@ -19,8 +19,26 @@ rootProject.name = "vortex-root" // API bindings include("vortex-jni") -include("vortex-spark_2.12") -project(":vortex-spark_2.12").projectDir = file("vortex-spark") -include("vortex-spark_2.13") -project(":vortex-spark_2.13").projectDir = file("vortex-spark") +val sparkModules = + mapOf( + "3.5" to listOf("2.12", "2.13"), + "4.0" to listOf("2.13"), + ) +val requestedSparkVersions = + System + .getProperty("sparkVersions") + ?.split(',') + ?.map(String::trim) + ?.filter(String::isNotEmpty) + ?.toSet() + +sparkModules.forEach { (sparkVersion, scalaVersions) -> + if (requestedSparkVersions == null || sparkVersion in requestedSparkVersions) { + scalaVersions.forEach { scalaVersion -> + val projectName = "vortex-spark-${sparkVersion}_$scalaVersion" + include(projectName) + project(":$projectName").projectDir = file("vortex-spark/v$sparkVersion") + } + } +} diff --git a/java/vortex-jni/build.gradle.kts b/java/vortex-jni/build.gradle.kts index 8f5106a1262..30509863e6a 100644 --- a/java/vortex-jni/build.gradle.kts +++ b/java/vortex-jni/build.gradle.kts @@ -8,7 +8,7 @@ import org.gradle.api.tasks.Exec plugins { `java-library` `jvm-test-suite` - id("com.gradleup.shadow") version "9.4.2" + id("com.gradleup.shadow") } dependencies { @@ -164,6 +164,11 @@ tasks.named("processResources").configure { dependsOn("makeTestFiles") } +// The sources jar packages src/main/resources, where makeTestFiles stages the native library. +tasks.withType().configureEach { + mustRunAfter("makeTestFiles") +} + tasks.withType().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 paths, List columns, Map formatOptions, Predicate[] pushedPredicates) { - this.paths = List.copyOf(paths); - this.readSchema = CatalogV2Util.v2ColumnsToStructType(columns.toArray(new Column[0])); - this.formatOptions = Map.copyOf(formatOptions); - this.pushedPredicates = pushedPredicates == null ? new Predicate[0] : pushedPredicates.clone(); - } - - /** - * Plans the input partitions for this batch scan. - * - *

Directory-like entries are expanded to concrete {@code .vortex} files. Each resolved file becomes its own - * {@link VortexFilePartition}; the partition carries the paths the reader should open, the requested schema, and - * any Hive-style partition values parsed out of the path. - */ - @Override - public InputPartition[] planInputPartitions() { - resolvedPaths = resolvePaths(); - return resolvedPaths.stream() - .map(path -> { - Map partVals = PartitionPathUtils.parsePartitionValues(path); - return new VortexFilePartition( - List.of(path), readSchema, formatOptions, ImmutableMap.copyOf(partVals)); - }) - .toArray(InputPartition[]::new); - } - - @Override - public PartitionReaderFactory createReaderFactory() { - List files = resolvedPaths != null ? resolvedPaths : resolvePaths(); - Set partitionColumns = collectPartitionColumnNames(files); - List dataColumnNames = Arrays.stream(readSchema.fieldNames()) - .filter(name -> !partitionColumns.contains(name)) - .collect(Collectors.toList()); - return new VortexPartitionReaderFactory(dataColumnNames, formatOptions, pushedPredicates); - } - - private List resolvePaths() { - return resolveVortexPaths(VortexSparkSession.get(formatOptions), paths, formatOptions); - } - - /** - * Expands directory-like entries to concrete {@code .vortex} files; entries that already name a {@code .vortex} - * file are kept as-is. Shared with {@link VortexScan#estimateStatistics()} so planning and execution resolve paths - * identically. - */ - static List resolveVortexPaths(Session session, List paths, Map formatOptions) { - return paths.stream() - .flatMap(path -> path.endsWith(".vortex") - ? Stream.of(path) - : NativeFiles.listFiles(session, path, formatOptions).stream()) - .collect(Collectors.toList()); - } - - private static Set collectPartitionColumnNames(List files) { - Set all = new HashSet<>(); - for (String path : files) { - all.addAll(PartitionPathUtils.parsePartitionValues(path).keySet()); - } - return all; - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java deleted file mode 100644 index 46616631d00..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReader.java +++ /dev/null @@ -1,176 +0,0 @@ -// 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.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.spark.VortexFilePartition; -import dev.vortex.spark.VortexSparkSession; -import java.io.IOException; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import org.apache.spark.sql.connector.expressions.filter.Predicate; -import org.apache.spark.sql.connector.read.PartitionReader; -import org.apache.spark.sql.types.StructField; -import org.apache.spark.sql.vectorized.ColumnVector; -import org.apache.spark.sql.vectorized.ColumnarBatch; - -/** - * Per-{@link VortexFilePartition} columnar reader. - * - *

Opens a single Vortex {@link Session}, {@link DataSource} and {@link Scan} spanning all of - * {@link VortexFilePartition#paths()} and streams every Vortex partition's record batches through the - * {@link PartitionReader} interface. - */ -final class VortexPartitionReader implements PartitionReader { - private final VortexFilePartition spark; - private final BufferAllocator allocator; - - // Held so the DataSource/Scan stay reachable even if the JVM-wide singleton is - // ever reset during a task; the actual native session is owned by - // {@link VortexSparkSession} and is not released when this reader closes. - private Session session; - private DataSource dataSource; - private Scan scan; - - private Partition currentPartition; - private ArrowReader currentReader; - private boolean currentBatchLoaded; - private boolean exhausted; - - VortexPartitionReader( - VortexFilePartition spark, - List dataColumnNames, - Map formatOptions, - Predicate[] pushedPredicates) { - this.spark = spark; - this.allocator = ArrowAllocation.rootAllocator(); - - session = VortexSparkSession.get(formatOptions); - dataSource = DataSource.open(session, spark.paths(), formatOptions); - - var options = ScanOptions.builder(); - if (!dataColumnNames.isEmpty()) { - Expression projection = Expression.select(dataColumnNames.toArray(new String[0]), Expression.root()); - options.projection(projection); - } - if (pushedPredicates != null && pushedPredicates.length > 0) { - buildFilterExpression(pushedPredicates).ifPresent(options::filter); - } - scan = dataSource.scan(options.build()); - } - - private static Optional buildFilterExpression(Predicate[] predicates) { - Expression combined = null; - for (Predicate predicate : predicates) { - Optional expr = SparkPredicateToVortexExpression.convert(predicate); - if (expr.isEmpty()) { - continue; - } - combined = combined == null ? expr.get() : Expression.and(combined, expr.get()); - } - return Optional.ofNullable(combined); - } - - @Override - public boolean next() { - if (exhausted) { - return false; - } - while (true) { - if (currentReader != null) { - try { - if (currentReader.loadNextBatch()) { - currentBatchLoaded = 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 (!currentBatchLoaded) { - throw new IllegalStateException("no batch loaded; call next() first"); - } - currentBatchLoaded = false; - - VectorSchemaRoot root; - try { - root = currentReader.getVectorSchemaRoot(); - } catch (IOException e) { - throw failure("read the loaded batch of", e); - } - - int rowCount = root.getRowCount(); - Map partVals = spark.partitionValues(); - if (partVals.isEmpty()) { - ColumnVector[] vectors = new ColumnVector[root.getFieldVectors().size()]; - for (int i = 0; i < vectors.length; i++) { - vectors[i] = new VortexArrowColumnVector(root.getFieldVectors().get(i)); - } - return new ColumnarBatch(vectors, rowCount); - } - - StructField[] fields = spark.readSchema().fields(); - ColumnVector[] combined = new ColumnVector[fields.length]; - int dataIdx = 0; - for (int i = 0; i < fields.length; i++) { - StructField field = fields[i]; - String partValue = partVals.get(field.name()); - if (partValue != null) { - combined[i] = PartitionPathUtils.createConstantVector(rowCount, field.dataType(), partValue); - } else { - combined[i] = new VortexArrowColumnVector(root.getFieldVectors().get(dataIdx++)); - } - } - return new ColumnarBatch(combined, rowCount); - } - - @Override - public void close() { - closeCurrentReader(); - // Scan and DataSource native resources are released by VortexCleaner once - // references are dropped. Session is the JVM-wide singleton and outlives this reader. - scan = null; - dataSource = null; - session = null; - } - - /** Wraps a failure with the paths being read, the one fact a stack trace alone does not give. */ - private RuntimeException failure(String what, IOException cause) { - return new RuntimeException(String.format("Failed to %s %s", what, spark.paths()), cause); - } - - private void closeCurrentReader() { - if (currentReader != null) { - try { - currentReader.close(); - } catch (IOException e) { - throw failure("close the reader over", e); - } - currentReader = null; - } - currentPartition = null; - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java deleted file mode 100644 index e187e4863b1..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexPartitionReaderFactory.java +++ /dev/null @@ -1,58 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.read; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import dev.vortex.jni.NativeRuntime; -import dev.vortex.spark.VortexFilePartition; -import java.io.Serializable; -import java.util.List; -import java.util.Map; -import org.apache.spark.sql.catalyst.InternalRow; -import org.apache.spark.sql.connector.expressions.filter.Predicate; -import org.apache.spark.sql.connector.read.InputPartition; -import org.apache.spark.sql.connector.read.PartitionReader; -import org.apache.spark.sql.connector.read.PartitionReaderFactory; -import org.apache.spark.sql.vectorized.ColumnarBatch; - -/** - * Factory that produces columnar readers for Vortex files. - * - *

The set of paths belongs to each {@link VortexFilePartition} — the factory itself is stateless across partitions. - * For every input partition, {@link VortexPartitionReader} opens a single {@code Session}, {@code DataSource} and - * {@code Scan} spanning that partition's paths and consumes every Vortex partition produced by that scan before - * returning. - */ -public final class VortexPartitionReaderFactory implements PartitionReaderFactory, Serializable { - private static final long serialVersionUID = 1L; - - private final ImmutableList dataColumnNames; - private final ImmutableMap formatOptions; - private final Predicate[] pushedPredicates; - - public VortexPartitionReaderFactory( - List dataColumnNames, Map formatOptions, Predicate[] pushedPredicates) { - this.dataColumnNames = ImmutableList.copyOf(dataColumnNames); - this.formatOptions = ImmutableMap.copyOf(formatOptions); - this.pushedPredicates = pushedPredicates == null ? new Predicate[0] : pushedPredicates.clone(); - } - - @Override - public PartitionReader createReader(InputPartition partition) { - throw new UnsupportedOperationException("row-based reads are not supported"); - } - - @Override - public PartitionReader createColumnarReader(InputPartition partition) { - NativeRuntime.setWorkerThreads(Integer.parseInt(formatOptions.getOrDefault("vortex.workerThreads", "4"))); - VortexFilePartition spark = (VortexFilePartition) partition; - return new VortexPartitionReader(spark, dataColumnNames, formatOptions, pushedPredicates); - } - - @Override - public boolean supportColumnarReads(InputPartition partition) { - return true; - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java deleted file mode 100644 index 02a7563f925..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java +++ /dev/null @@ -1,176 +0,0 @@ -// 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.Session; -import dev.vortex.spark.VortexSparkSession; -import java.util.Arrays; -import java.util.List; -import java.util.Map; -import java.util.OptionalLong; -import org.apache.spark.sql.connector.catalog.CatalogV2Util; -import org.apache.spark.sql.connector.catalog.Column; -import org.apache.spark.sql.connector.expressions.NamedReference; -import org.apache.spark.sql.connector.expressions.filter.Predicate; -import org.apache.spark.sql.connector.read.Batch; -import org.apache.spark.sql.connector.read.Scan; -import org.apache.spark.sql.connector.read.Statistics; -import org.apache.spark.sql.connector.read.SupportsReportStatistics; -import org.apache.spark.sql.connector.read.colstats.ColumnStatistics; -import org.apache.spark.sql.internal.SQLConf; -import org.apache.spark.sql.types.StructType; - -/** - * Spark V2 {@link Scan} over a table of Vortex files. - * - *

Implements {@link SupportsReportStatistics} to surface both the row count Vortex records in each file footer and a - * Spark scan-size estimate. The byte estimate starts from the on-storage file sizes collected by - * {@code MultiFileDataSource}, then follows Spark's file scan convention by applying the SQL file-compression factor - * and scaling by the pushed read schema's default size relative to the full table schema's default size. When the - * listing did not return a size for one or more files the file-byte total is extrapolated before Spark scaling is - * applied. - */ -public final class VortexScan implements Scan, SupportsReportStatistics { - - private final List paths; - private final List tableColumns; - private final List readColumns; - private final Map formatOptions; - private final Predicate[] pushedPredicates; - - private volatile Statistics cachedStatistics; - - /** - * Creates a new VortexScan for the specified file paths and columns. The caller is responsible for passing - * immutable collections; the constructor does not copy. - * - * @param paths the list of Vortex file paths to scan - * @param tableColumns the full table columns before projection pushdown - * @param readColumns the list of columns to read from the files - * @param pushedPredicates predicates pushed down by Spark; {@code null} or empty means no pushdown - */ - public VortexScan( - List paths, - List tableColumns, - List readColumns, - Predicate[] pushedPredicates, - Map formatOptions) { - this.paths = paths; - this.tableColumns = tableColumns; - this.readColumns = readColumns; - this.formatOptions = formatOptions; - this.pushedPredicates = pushedPredicates == null ? new Predicate[0] : pushedPredicates.clone(); - } - - /** - * Returns the schema for the data that will be read by this scan. - * - *

The schema is constructed from the read columns that were specified when this scan was created. - * - * @return the StructType representing the schema of the read data - */ - @Override - public StructType readSchema() { - return CatalogV2Util.v2ColumnsToStructType(readColumns.toArray(new Column[0])); - } - - /** Logging-friendly readable description of the scan source. */ - @Override - public String description() { - return String.format( - "VortexScan{paths=%s, columns=%s, pushedPredicates=%s}", - paths, readColumns, Arrays.toString(pushedPredicates)); - } - - /** - * Converts this scan to a Batch for execution. - * - *

Creates a VortexBatchExec that will handle the actual reading of the specified files and columns. - * - * @return a Batch implementation for executing this scan - */ - @Override - public Batch toBatch() { - return new VortexBatchExec(paths, readColumns, formatOptions, pushedPredicates); - } - - /** - * Returns the columnar support mode for this scan. - * - *

Vortex always provides columnar data access, so this method always returns SUPPORTED. - * - * @return ColumnarSupportMode.SUPPORTED - */ - @Override - public ColumnarSupportMode columnarSupportMode() { - return ColumnarSupportMode.SUPPORTED; - } - - /** - * Returns statistics for this scan. - * - *

Opens the Vortex {@link DataSource} on first invocation and caches the result. The row count is taken from the - * data source (sum of file-footer row counts; extrapolated from the first opened file when other files are - * deferred). {@link Statistics#sizeInBytes()} is derived from the per-file sizes reported by the filesystem - * listing, then adjusted by Spark's compression factor and the ratio between the pushed read schema and the full - * table schema. When a listing did not return a size for some file the file-byte total is extrapolated. When no - * file size is known at all the value is left empty so Spark falls back to its default heuristic. - * - * @return statistics with row-count and Spark scan-size estimates - */ - @Override - public Statistics estimateStatistics() { - Statistics local = cachedStatistics; - if (local != null) { - return local; - } - synchronized (this) { - if (cachedStatistics == null) { - cachedStatistics = computeStatistics(); - } - return cachedStatistics; - } - } - - private Statistics computeStatistics() { - Session session = VortexSparkSession.get(formatOptions); - List resolvedPaths = VortexBatchExec.resolveVortexPaths(session, paths, formatOptions); - if (resolvedPaths.isEmpty()) { - return new VortexStatistics(OptionalLong.empty(), OptionalLong.empty()); - } - - DataSource source = DataSource.open(session, resolvedPaths, formatOptions); - return new VortexStatistics( - source.rowCount().asOptional(), - scaleSizeInBytes(source.byteSize().asOptional())); - } - - private OptionalLong scaleSizeInBytes(OptionalLong fileBytes) { - if (fileBytes.isEmpty()) { - return OptionalLong.empty(); - } - - StructType tableSchema = CatalogV2Util.v2ColumnsToStructType(tableColumns.toArray(new Column[0])); - StructType readSchema = readSchema(); - int tableDefaultSize = tableSchema.defaultSize(); - if (tableDefaultSize <= 0) { - return fileBytes; - } - - double scaled = SQLConf.get().fileCompressionFactor() - * fileBytes.getAsLong() - / tableDefaultSize - * readSchema.defaultSize(); - return OptionalLong.of((long) scaled); - } - - private record VortexStatistics(OptionalLong numRows, OptionalLong sizeInBytes) implements Statistics { - - @Override - public Map columnStats() { - return Map.of(); - } - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java deleted file mode 100644 index 62c8085aa0f..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScanBuilder.java +++ /dev/null @@ -1,196 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.read; - -import static com.google.common.base.Preconditions.checkState; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.Maps; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import org.apache.spark.sql.connector.catalog.CatalogV2Util; -import org.apache.spark.sql.connector.catalog.Column; -import org.apache.spark.sql.connector.expressions.NamedReference; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.connector.expressions.filter.Predicate; -import org.apache.spark.sql.connector.read.Scan; -import org.apache.spark.sql.connector.read.ScanBuilder; -import org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns; -import org.apache.spark.sql.connector.read.SupportsPushDownV2Filters; -import org.apache.spark.sql.types.DataType; -import org.apache.spark.sql.types.StructType; - -/** Spark V2 {@link ScanBuilder} for table scans over Vortex files. */ -public final class VortexScanBuilder - implements ScanBuilder, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters { - private final ImmutableList.Builder paths; - private final List tableColumns; - private final List readColumns; - private final Map formatOptions; - private final Set partitionColumnNames; - private Predicate[] pushedPredicates = new Predicate[0]; - - /** Creates a new VortexScanBuilder with empty paths and columns. */ - public VortexScanBuilder(Map formatOptions) { - this(formatOptions, new Transform[0]); - } - - /** - * Creates a new VortexScanBuilder with empty paths and columns and the supplied partition transforms. Filters that - * reference partition columns are not pushed down, since the partition columns are not stored inside the Vortex - * files. - */ - public VortexScanBuilder(Map formatOptions, Transform[] partitionTransforms) { - this.paths = ImmutableList.builder(); - Map options = Maps.newHashMap(); - options.put("vortex.workerThreads", "4"); - options.putAll(formatOptions); - this.tableColumns = new ArrayList<>(); - this.readColumns = new ArrayList<>(); - this.formatOptions = options; - this.partitionColumnNames = collectPartitionColumnNames(partitionTransforms); - } - - /** - * Adds a file path to scan. - * - * @param path the file path to add - * @return this builder for method chaining - */ - public VortexScanBuilder addPath(String path) { - this.paths.add(path); - return this; - } - - /** - * Adds a column to read. - * - * @param column the column to add - * @return this builder for method chaining - */ - public VortexScanBuilder addColumn(Column column) { - this.tableColumns.add(column); - this.readColumns.add(column); - return this; - } - - /** - * Adds multiple file paths to scan. - * - * @param paths the iterable of file paths to add - * @return this builder for method chaining - */ - public VortexScanBuilder addAllPaths(Iterable paths) { - this.paths.addAll(paths); - return this; - } - - /** - * Adds multiple columns to read. - * - * @param columns the iterable of columns to add - * @return this builder for method chaining - */ - public VortexScanBuilder addAllColumns(Iterable columns) { - for (Column column : columns) { - addColumn(column); - } - return this; - } - - /** - * Builds a VortexScan with the configured paths and columns. - * - *

An empty column list is allowed: aggregates such as {@code count()} need no column data, and the scan then - * reads the minimal schema. - * - * @return a new VortexScan instance - * @throws IllegalStateException if no paths have been added - */ - @Override - public Scan build() { - var paths = this.paths.build(); - - checkState(!paths.isEmpty(), "paths cannot be empty"); - - return new VortexScan( - paths, - List.copyOf(this.tableColumns), - List.copyOf(this.readColumns), - pushedPredicates, - this.formatOptions); - } - - /** - * Prunes the columns to only include those specified in the required schema. - * - *

This method clears the current column list and replaces it with columns derived from the required schema. - * Currently only supports top-level schema pruning - deeply nested schema pruning is not yet implemented. - * - * @param requiredSchema the schema specifying which columns are required - */ - @Override - public void pruneColumns(StructType requiredSchema) { - readColumns.clear(); - readColumns.addAll(Arrays.asList(CatalogV2Util.structTypeToV2Columns(requiredSchema))); - } - - /** - * Splits the supplied predicates into pushed and not-pushed sets. - * - *

A predicate is pushed when it references only data columns (not partition columns) and uses operators and - * literal types that {@link SparkPredicateToVortexExpression} can map to Vortex expressions. Predicates that - * reference partition columns or use unsupported features are returned to Spark for post-scan evaluation. - * - * @return the predicates that Spark must still evaluate - */ - @Override - public Predicate[] pushPredicates(Predicate[] predicates) { - Map dataColumnTypes = new HashMap<>(); - for (Column column : readColumns) { - if (!partitionColumnNames.contains(column.name())) { - dataColumnTypes.put(column.name(), column.dataType()); - } - } - List pushed = new ArrayList<>(); - List postScan = new ArrayList<>(); - for (Predicate predicate : predicates) { - if (SparkPredicateToVortexExpression.isPushable(predicate, dataColumnTypes)) { - pushed.add(predicate); - } else { - postScan.add(predicate); - } - } - this.pushedPredicates = pushed.toArray(new Predicate[0]); - return postScan.toArray(new Predicate[0]); - } - - /** Returns the predicates this scan promises to apply. */ - @Override - public Predicate[] pushedPredicates() { - return Arrays.copyOf(pushedPredicates, pushedPredicates.length); - } - - private static Set collectPartitionColumnNames(Transform[] transforms) { - if (transforms == null || transforms.length == 0) { - return Collections.emptySet(); - } - Set names = new HashSet<>(); - for (Transform transform : transforms) { - for (NamedReference ref : transform.references()) { - String[] parts = ref.fieldNames(); - if (parts.length == 1) { - names.add(parts[0]); - } - } - } - return names; - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/PartitionedVortexDataWriter.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/PartitionedVortexDataWriter.java deleted file mode 100644 index 01ed570e4fc..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/PartitionedVortexDataWriter.java +++ /dev/null @@ -1,473 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.write; - -import com.google.common.collect.ImmutableList; -import com.google.common.primitives.ImmutableIntArray; -import java.io.IOException; -import java.io.Serializable; -import java.net.URLEncoder; -import java.nio.charset.StandardCharsets; -import java.time.Instant; -import java.time.LocalDate; -import java.time.LocalDateTime; -import java.time.ZoneOffset; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.stream.Collectors; -import org.apache.hadoop.shaded.com.google.common.collect.Streams; -import org.apache.spark.sql.catalyst.InternalRow; -import org.apache.spark.sql.catalyst.expressions.BoundReference; -import org.apache.spark.sql.catalyst.expressions.UnsafeProjection; -import org.apache.spark.sql.connector.expressions.Expression; -import org.apache.spark.sql.connector.expressions.Literal; -import org.apache.spark.sql.connector.expressions.NamedReference; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.connector.write.DataWriter; -import org.apache.spark.sql.connector.write.WriterCommitMessage; -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.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.util.CaseInsensitiveStringMap; -import org.apache.spark.unsafe.types.UTF8String; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Writes Spark InternalRow data to Vortex files organized in Hive-style partition directories. - * - *

Supports the standard Spark partition transforms: {@code identity}, {@code years}, {@code months}, {@code days}, - * {@code hours}, and {@code bucket}. For each unique combination of evaluated transform values, a separate subdirectory - * is created and a dedicated {@link VortexDataWriter} writes data within it. - */ -public final class PartitionedVortexDataWriter implements DataWriter, AutoCloseable { - private static final Logger logger = LoggerFactory.getLogger(PartitionedVortexDataWriter.class); - private static final String HIVE_DEFAULT_PARTITION = "__HIVE_DEFAULT_PARTITION__"; - - private final String baseOutputUri; - private final StructType dataSchema; - private final UnsafeProjection dataProjection; - private final CaseInsensitiveStringMap options; - private final ResolvedTransform[] resolvedTransforms; - private final int partitionId; - private final long taskId; - - private final Map writers = new HashMap<>(); - private boolean closed = false; - - /** - * Creates a new PartitionedVortexDataWriter. - * - * @param baseOutputUri the base output path - * @param schema the full schema of the data - * @param options write options - * @param resolvedTransforms pre-resolved partition transforms - * @param partitionId the Spark partition ID - * @param taskId the Spark task ID - */ - PartitionedVortexDataWriter( - String baseOutputUri, - StructType schema, - CaseInsensitiveStringMap options, - ResolvedTransform[] resolvedTransforms, - int partitionId, - long taskId) { - this.baseOutputUri = baseOutputUri.endsWith("/") ? baseOutputUri : baseOutputUri + "/"; - this.options = options; - this.partitionId = partitionId; - this.taskId = taskId; - this.resolvedTransforms = resolvedTransforms; - - // Compute the data schema by removing identity partition columns. - // Only identity transforms correspond to columns that should be stripped from the data, - // since temporal/bucket transforms derive values from the source column. - Set identityPartitionIndices = new HashSet<>(); - for (ResolvedTransform rt : resolvedTransforms) { - if ("identity".equals(rt.transformName())) { - identityPartitionIndices.add(rt.columnIndices().get(0)); - } - } - - StructField[] fields = schema.fields(); - List dataFields = new ArrayList<>(); - List projExprs = new ArrayList<>(); - for (int i = 0; i < fields.length; i++) { - if (!identityPartitionIndices.contains(i)) { - dataFields.add(fields[i]); - projExprs.add(new BoundReference(i, fields[i].dataType(), fields[i].nullable())); - } - } - this.dataSchema = new StructType(dataFields.toArray(new StructField[0])); - this.dataProjection = UnsafeProjection.create(asScalaSeq(projExprs)); - } - - @SuppressWarnings("deprecation") // JavaConverters is deprecated in Scala 2.13 but works in both 2.12 and 2.13 - private static scala.collection.immutable.Seq asScalaSeq(List list) { - return scala.collection.JavaConverters.asScalaBufferConverter(list) - .asScala() - .toList(); - } - - @Override - public void write(InternalRow row) throws IOException { - String partitionPath = getPartitionPath(row); - VortexDataWriter writer = writers.get(partitionPath); - if (writer == null) { - writer = createWriterForPartition(partitionPath); - writers.put(partitionPath, writer); - } - writer.write(dataProjection.apply(row)); - } - - @Override - public WriterCommitMessage commit() throws IOException { - if (closed) { - return new PartitionedWriterCommitMessage(List.of()); - } - - List messages = new ArrayList<>(); - IOException firstException = null; - - for (Map.Entry entry : writers.entrySet()) { - try { - WriterCommitMessage msg = entry.getValue().commit(); - if (msg instanceof VortexWriterCommitMessage) { - messages.add((VortexWriterCommitMessage) msg); - } - } catch (IOException e) { - if (firstException == null) { - firstException = e; - } else { - firstException.addSuppressed(e); - } - } - } - - closed = true; - - if (firstException != null) { - throw firstException; - } - - logger.info("Committed {} partition writers", messages.size()); - return new PartitionedWriterCommitMessage(messages); - } - - @Override - public void abort() throws IOException { - if (closed) { - return; - } - - for (VortexDataWriter writer : writers.values()) { - try { - writer.abort(); - } catch (IOException e) { - logger.error("Error aborting partition writer", e); - } - } - closed = true; - } - - @Override - public void close() throws IOException { - if (!closed) { - logger.warn("PartitionedVortexDataWriter.close() called without commit() or abort() - cleaning up"); - try { - abort(); - } catch (IOException e) { - logger.error("Error during cleanup in close()", e); - } - } - } - - private VortexDataWriter createWriterForPartition(String partitionPath) { - String fileName = String.format("part-%05d-%d.vortex", partitionId, taskId); - String fileUri = baseOutputUri + partitionPath + "/" + fileName; - logger.debug("Creating writer for partition path: {}", fileUri); - return new VortexDataWriter(fileUri, dataSchema, options); - } - - private String getPartitionPath(InternalRow row) { - StringBuilder sb = new StringBuilder(); - for (int i = 0; i < resolvedTransforms.length; i++) { - if (i > 0) { - sb.append("/"); - } - ResolvedTransform rt = resolvedTransforms[i]; - sb.append(URLEncoder.encode(rt.directoryKey, StandardCharsets.UTF_8)); - sb.append("="); - - String value = evaluateTransform(rt, row); - sb.append(URLEncoder.encode(value, StandardCharsets.UTF_8)); - } - return sb.toString(); - } - - // ------------------------------------------------------------------ - // Transform resolution: converts Transform[] into ResolvedTransform[] - // ------------------------------------------------------------------ - - static ResolvedTransform[] resolveTransforms(Transform[] transforms, StructType schema) { - return Arrays.stream(transforms) - .map(transform -> resolveOne(transform, schema)) - .toArray(ResolvedTransform[]::new); - } - - private static ResolvedTransform resolveOne(Transform transform, StructType schema) { - String transformName = transform.name(); - NamedReference[] refs = transform.references(); - - if (refs.length == 0) { - throw new IllegalArgumentException("Partition transform has no column references: " + transform); - } - - // Primary column (all single-column transforms use this) - String primaryColName = String.join(".", refs[0].fieldNames()); - int primaryColIdx = schema.fieldIndex(primaryColName); - DataType primaryType = schema.fields()[primaryColIdx].dataType(); - - switch (transformName) { - case "identity": - return new ResolvedTransform(primaryColName, transformName, primaryColIdx, primaryType); - - case "years": - requireTemporalType(primaryType, transformName); - return new ResolvedTransform(primaryColName + "_year", transformName, primaryColIdx, primaryType); - - case "months": - requireTemporalType(primaryType, transformName); - return new ResolvedTransform(primaryColName + "_month", transformName, primaryColIdx, primaryType); - - case "days": - requireTemporalType(primaryType, transformName); - return new ResolvedTransform(primaryColName + "_day", transformName, primaryColIdx, primaryType); - - case "hours": - requireTimestampType(primaryType, transformName); - return new ResolvedTransform(primaryColName + "_hour", transformName, primaryColIdx, primaryType); - - case "bucket": { - int bucketCount = extractBucketCount(transform); - String colNames = Arrays.stream(refs) - .map(r -> String.join(".", r.fieldNames())) - .collect(Collectors.joining("_")); - - // Resolve all referenced columns for multi-column bucket - ImmutableIntArray.Builder allIndices = ImmutableIntArray.builder(refs.length); - ImmutableList.Builder allTypes = ImmutableList.builderWithExpectedSize(refs.length); - for (NamedReference ref : refs) { - String colName = String.join(".", ref.fieldNames()); - int idx = schema.fieldIndex(colName); - allIndices.add(idx); - allTypes.add(schema.fields()[idx].dataType()); - } - return new ResolvedTransform( - colNames + "_bucket", transformName, allIndices.build(), allTypes.build(), bucketCount); - } - - default: - throw new IllegalArgumentException("Unsupported partition transform: " + transformName); - } - } - - private static int extractBucketCount(Transform transform) { - for (Expression arg : transform.arguments()) { - if (arg instanceof Literal) { - Object value = ((Literal) arg).value(); - if (value instanceof Integer) { - return (Integer) value; - } - } - } - throw new IllegalArgumentException("bucket transform missing integer numBuckets argument"); - } - - private static void requireTemporalType(DataType type, String transformName) { - if (!(type instanceof DateType || type instanceof TimestampType || type instanceof TimestampNTZType)) { - throw new IllegalArgumentException( - transformName + " transform requires a date or timestamp column, got: " + type); - } - } - - private static void requireTimestampType(DataType type, String transformName) { - if (!(type instanceof TimestampType || type instanceof TimestampNTZType)) { - throw new IllegalArgumentException(transformName + " transform requires a timestamp column, got: " + type); - } - } - - // ------------------------------------------------------------------ - // Transform evaluation: produces partition values from rows - // ------------------------------------------------------------------ - - private static String evaluateTransform(ResolvedTransform rt, InternalRow row) { - int colIdx = rt.columnIndices.get(0); - - if (row.isNullAt(colIdx)) { - return HIVE_DEFAULT_PARTITION; - } - - return switch (rt.transformName) { - case "identity" -> extractIdentityValue(row, colIdx, rt.columnTypes.get(0)); - case "years" -> extractYearValue(row, colIdx, rt.columnTypes.get(0)); - case "months" -> extractMonthValue(row, colIdx, rt.columnTypes.get(0)); - case "days" -> extractDayValue(row, colIdx, rt.columnTypes.get(0)); - case "hours" -> extractHourValue(row, colIdx, rt.columnTypes.get(0)); - case "bucket" -> extractBucketValue(row, rt); - default -> throw new IllegalArgumentException("Unsupported transform: " + rt.transformName); - }; - } - - private static String extractIdentityValue(InternalRow row, int ordinal, DataType dataType) { - if (dataType instanceof BooleanType) { - return String.valueOf(row.getBoolean(ordinal)); - } else if (dataType instanceof ByteType) { - return String.valueOf(row.getByte(ordinal)); - } else if (dataType instanceof ShortType) { - return String.valueOf(row.getShort(ordinal)); - } else if (dataType instanceof IntegerType) { - return String.valueOf(row.getInt(ordinal)); - } else if (dataType instanceof LongType) { - return String.valueOf(row.getLong(ordinal)); - } else if (dataType instanceof FloatType) { - return String.valueOf(row.getFloat(ordinal)); - } else if (dataType instanceof DoubleType) { - return String.valueOf(row.getDouble(ordinal)); - } else if (dataType instanceof StringType) { - UTF8String str = row.getUTF8String(ordinal); - return str != null ? str.toString() : HIVE_DEFAULT_PARTITION; - } else if (dataType instanceof DateType) { - return String.valueOf(row.getInt(ordinal)); - } else if (dataType instanceof TimestampType || dataType instanceof TimestampNTZType) { - return String.valueOf(row.getLong(ordinal)); - } else { - throw new IllegalArgumentException("Unsupported partition column type: " + dataType); - } - } - - private static String extractYearValue(InternalRow row, int colIdx, DataType type) { - if (type instanceof DateType) { - return String.valueOf(LocalDate.ofEpochDay(row.getInt(colIdx)).getYear()); - } else { - return String.valueOf(microsToDateTime(row.getLong(colIdx)).getYear()); - } - } - - private static String extractMonthValue(InternalRow row, int colIdx, DataType type) { - LocalDate date; - if (type instanceof DateType) { - date = LocalDate.ofEpochDay(row.getInt(colIdx)); - } else { - date = microsToDateTime(row.getLong(colIdx)).toLocalDate(); - } - return String.format("%04d-%02d", date.getYear(), date.getMonthValue()); - } - - private static String extractDayValue(InternalRow row, int colIdx, DataType type) { - LocalDate date; - if (type instanceof DateType) { - date = LocalDate.ofEpochDay(row.getInt(colIdx)); - } else { - date = microsToDateTime(row.getLong(colIdx)).toLocalDate(); - } - return date.toString(); // YYYY-MM-DD - } - - private static String extractHourValue(InternalRow row, int colIdx, DataType type) { - LocalDateTime dt = microsToDateTime(row.getLong(colIdx)); - return String.format("%s-%02d", dt.toLocalDate(), dt.getHour()); - } - - /** - * Computes the bucket value matching Spark's {@code InMemoryBaseTable} reference implementation: per-column values - * are converted to longs (hashed for strings/binary), summed, then {@code Math.floorMod(sum, numBuckets)}. - */ - private static String extractBucketValue(InternalRow row, ResolvedTransform rt) { - long hash = Streams.zip( - rt.columnIndices.stream().boxed(), - rt.columnTypes.stream(), - (Integer idx, DataType dt) -> columnHashValue(row, idx, dt)) - .reduce(0L, Long::sum); - int bucket = Math.floorMod(hash, rt.bucketCount); - return String.valueOf(bucket); - } - - private static long columnHashValue(InternalRow row, int ordinal, DataType dataType) { - if (dataType instanceof ByteType) { - return row.getByte(ordinal); - } else if (dataType instanceof ShortType) { - return row.getShort(ordinal); - } else if (dataType instanceof IntegerType) { - return row.getInt(ordinal); - } else if (dataType instanceof LongType - || dataType instanceof TimestampType - || dataType instanceof TimestampNTZType) { - return row.getLong(ordinal); - } else if (dataType instanceof StringType) { - return row.getUTF8String(ordinal).hashCode(); - } else if (dataType instanceof BinaryType) { - return java.util.Arrays.hashCode(row.getBinary(ordinal)); - } else { - throw new IllegalArgumentException("Unsupported bucket column type: " + dataType); - } - } - - private static LocalDateTime microsToDateTime(long micros) { - long epochSecond = Math.floorDiv(micros, 1_000_000); - int nanoOfSecond = (int) (Math.floorMod(micros, 1_000_000) * 1000); - return LocalDateTime.ofInstant(Instant.ofEpochSecond(epochSecond, nanoOfSecond), ZoneOffset.UTC); - } - - // ------------------------------------------------------------------ - // Internal types - // ------------------------------------------------------------------ - - /** - * Pre-resolved representation of a partition transform, ready for per-row evaluation. - * - * @param bucketCount -1 if not a bucket transform - */ - record ResolvedTransform( - String directoryKey, - String transformName, - ImmutableIntArray columnIndices, - List columnTypes, - int bucketCount) - implements Serializable { - ResolvedTransform(String directoryKey, String transformName, int columnIndex, DataType columnType) { - this(directoryKey, transformName, ImmutableIntArray.of(columnIndex), List.of(columnType), -1); - } - } - - /** Commit message that aggregates results from multiple partition writers. */ - public static final class PartitionedWriterCommitMessage implements WriterCommitMessage, Serializable { - private final List partitionMessages; - - PartitionedWriterCommitMessage(List partitionMessages) { - this.partitionMessages = partitionMessages; - } - - /** Returns the commit messages from each individual partition writer. */ - public List getPartitionMessages() { - return partitionMessages; - } - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java deleted file mode 100644 index c9be4159a59..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java +++ /dev/null @@ -1,168 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.write; - -import dev.vortex.jni.NativeFiles; -import dev.vortex.spark.VortexSparkSession; -import java.io.Serializable; -import java.util.Arrays; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; -import java.util.stream.Stream; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.connector.write.BatchWrite; -import org.apache.spark.sql.connector.write.DataWriterFactory; -import org.apache.spark.sql.connector.write.PhysicalWriteInfo; -import org.apache.spark.sql.connector.write.Write; -import org.apache.spark.sql.connector.write.WriterCommitMessage; -import org.apache.spark.sql.types.StructType; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Manages the batch write operation for creating Vortex files. - * - *

This class coordinates the distributed write operation across Spark executors, handling the creation of data - * writers and managing commits/aborts. - */ -public final class VortexBatchWrite implements Write, BatchWrite, Serializable { - - private static final Logger log = LoggerFactory.getLogger(VortexBatchWrite.class); - private final String outputPath; - private final StructType schema; - private final Map options; - private final boolean overwrite; - // Resolved eagerly so that Spark Transform objects (Scala case classes that are not - // Java-serializable) never reach the DataWriterFactory serialization boundary. - private final PartitionedVortexDataWriter.ResolvedTransform[] resolvedTransforms; - - /** - * Creates a new VortexBatchWrite. - * - * @param outputPath the base path where Vortex files will be written - * @param schema the schema of the data to write - * @param options additional write options - * @param overwrite whether to overwrite existing files - * @param partitionTransforms partition transforms (may be empty) - */ - VortexBatchWrite( - String outputPath, - StructType schema, - Map options, - boolean overwrite, - Transform[] partitionTransforms) { - this.outputPath = outputPath; - this.schema = schema; - this.options = options; - this.overwrite = overwrite; - this.resolvedTransforms = PartitionedVortexDataWriter.resolveTransforms(partitionTransforms, schema); - } - - /** - * Returns this object as a BatchWrite. - * - *

This method is required by the Write interface to support batch writes. - * - * @return this object - */ - @Override - public BatchWrite toBatch() { - return this; - } - - /** - * Creates a DataWriterFactory for producing data writers on executors. - * - *

This method is called once at the start of the write operation, making it the right place to handle overwrite - * cleanup. - * - * @return a new VortexDataWriterFactory - */ - @Override - public DataWriterFactory createBatchWriterFactory(PhysicalWriteInfo info) { - // Handle overwrite cleanup BEFORE writing starts - if (overwrite) { - var session = VortexSparkSession.get(options); - var uris = NativeFiles.listFiles(session, outputPath, options); - // Deleting the existing files is destructive and happens before the new data is written: - // if the subsequent write fails, abort() only removes the newly written files and cannot - // restore what was deleted here. Log loudly so operators can see what was removed. - log.warn( - "Deleting {} existing file(s) under {} because of overwrite, before writing new data; " - + "this cannot be undone if the subsequent write fails", - uris.size(), - outputPath); - NativeFiles.delete(session, uris.toArray(new String[0]), options); - } - - return new VortexDataWriterFactory(outputPath, schema, options, resolvedTransforms); - } - - /** - * Called when a single data writer task completes successfully. - * - *

This is called for each successful task but individual file commits are handled in the data writer itself. - * - * @param message commit message from a successful data writer task - */ - @Override - public void onDataWriterCommit(WriterCommitMessage message) { - // Individual file commits are handled in the data writer - // This is called for each successful task - log.debug("Committing DataWriter"); - } - - /** - * Commits the entire write job after all tasks complete successfully. - * - *

This finalizes the write operation and ensures all Vortex files are properly written. - * - * @param messages commit messages from all successful write tasks - */ - @Override - public void commit(WriterCommitMessage[] messages) { - List writtenFiles = extractFilePaths(messages); - - if (!writtenFiles.isEmpty()) { - log.info("Successfully wrote {} Vortex files to {}", writtenFiles.size(), outputPath); - } - } - - /** - * Aborts the write job due to failures. - * - *

Deletes the files the tasks reported, through the same native filesystem layer the writers used, so that URL - * paths ({@code file://}, {@code s3://}) resolve the same way on cleanup as they did on write. - * - * @param messages commit messages from write tasks (may include failures) - */ - @Override - public void abort(WriterCommitMessage[] messages) { - List filePaths = extractFilePaths(messages); - if (filePaths.isEmpty()) { - return; - } - log.warn("Deleting {} file(s) written before the job failed, under {}", filePaths.size(), outputPath); - try { - NativeFiles.delete(VortexSparkSession.get(options), filePaths.toArray(new String[0]), options); - } catch (RuntimeException e) { - log.error("Failed to clean up {} file(s) under {}", filePaths.size(), outputPath, e); - } - } - - private static List extractFilePaths(WriterCommitMessage[] messages) { - return Arrays.stream(messages) - .flatMap(msg -> { - if (msg instanceof VortexWriterCommitMessage) { - return Stream.of(((VortexWriterCommitMessage) msg).filePath()); - } else if (msg instanceof PartitionedVortexDataWriter.PartitionedWriterCommitMessage) { - return ((PartitionedVortexDataWriter.PartitionedWriterCommitMessage) msg) - .getPartitionMessages().stream().map(VortexWriterCommitMessage::filePath); - } - return Stream.empty(); - }) - .collect(Collectors.toList()); - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriterFactory.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriterFactory.java deleted file mode 100644 index e8237bde309..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriterFactory.java +++ /dev/null @@ -1,85 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.write; - -import java.io.Serializable; -import java.util.Map; -import org.apache.spark.sql.catalyst.InternalRow; -import org.apache.spark.sql.connector.write.DataWriter; -import org.apache.spark.sql.connector.write.DataWriterFactory; -import org.apache.spark.sql.types.StructType; -import org.apache.spark.sql.util.CaseInsensitiveStringMap; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Factory for creating VortexDataWriter instances on Spark executors. - * - *

This factory is serialized and sent to executors where it creates data writers for each task. When partition - * transforms are specified, it creates partitioned writers that organize output into Hive-style partition directories. - */ -public final class VortexDataWriterFactory implements DataWriterFactory, Serializable { - - private static final Logger log = LoggerFactory.getLogger(VortexDataWriterFactory.class); - - private final String outputUri; - private final StructType schema; - // Store options as a serializable Map instead of CaseInsensitiveStringMap - private final Map options; - private final PartitionedVortexDataWriter.ResolvedTransform[] resolvedTransforms; - - /** - * Creates a new VortexDataWriterFactory. - * - * @param outputUri the base path where Vortex files will be written - * @param schema the schema of the data to write - * @param options additional write options - * @param resolvedTransforms pre-resolved partition transforms (may be empty) - */ - VortexDataWriterFactory( - String outputUri, - StructType schema, - Map options, - PartitionedVortexDataWriter.ResolvedTransform[] resolvedTransforms) { - this.outputUri = outputUri; - this.schema = schema; - this.options = options; - this.resolvedTransforms = resolvedTransforms; - } - - /** - * Creates a new data writer for a specific partition and task. - * - *

Each task writes its data to a separate Vortex file to avoid conflicts. When partition transforms are - * configured, returns a {@link PartitionedVortexDataWriter} that creates Hive-style partition directories. - * - * @param partitionId the partition ID - * @param taskId the task ID - * @return a new DataWriter instance - */ - @Override - public DataWriter createWriter(int partitionId, long taskId) { - log.debug("Creating writer for partition={} task={}", partitionId, taskId); - - CaseInsensitiveStringMap optionsMap = new CaseInsensitiveStringMap(options); - - if (resolvedTransforms.length > 0) { - log.debug("Creating partitioned writer with {} transforms", resolvedTransforms.length); - return new PartitionedVortexDataWriter( - outputUri, schema, optionsMap, resolvedTransforms, partitionId, taskId); - } - - // Non-partitioned write: single file per task - String fileName = String.format("part-%05d-%d.vortex", partitionId, taskId); - String fileUri; - if (outputUri.endsWith("/")) { - fileUri = outputUri + fileName; - } else { - fileUri = outputUri + "/" + fileName; - } - - log.debug("Output file: {}", fileUri); - return new VortexDataWriter(fileUri, schema, optionsMap); - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexWriteBuilder.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexWriteBuilder.java deleted file mode 100644 index 921e586a910..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexWriteBuilder.java +++ /dev/null @@ -1,65 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.write; - -import java.util.Map; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.connector.write.LogicalWriteInfo; -import org.apache.spark.sql.connector.write.SupportsTruncate; -import org.apache.spark.sql.connector.write.Write; -import org.apache.spark.sql.connector.write.WriteBuilder; - -/** - * Builder for configuring Vortex write operations. - * - *

This class is responsible for creating BatchWrite instances that execute the actual write operations to create - * Vortex files from Spark DataFrames. - */ -public final class VortexWriteBuilder implements WriteBuilder, SupportsTruncate { - - private final String paths; - private final LogicalWriteInfo writeInfo; - private final Map options; - private final Transform[] partitionTransforms; - private boolean truncate = false; - - /** - * Creates a new VortexWriteBuilder. - * - * @param paths root path for write - * @param writeInfo logical information about the write operation - * @param options additional write options - * @param partitionTransforms partition transforms (may be empty) - */ - public VortexWriteBuilder( - String paths, LogicalWriteInfo writeInfo, Map options, Transform[] partitionTransforms) { - this.paths = paths; - this.writeInfo = writeInfo; - this.options = options; - this.partitionTransforms = partitionTransforms; - } - - /** - * Builds a Write for executing the write operation. - * - * @return a new VortexBatchWrite configured with this builder's settings - */ - @Override - public Write build() { - return new VortexBatchWrite(paths, writeInfo.schema(), options, truncate, partitionTransforms); - } - - /** - * Configures the write operation to truncate existing data. - * - *

When truncate is enabled, existing Vortex files at the output path will be removed before writing new data. - * - * @return this builder for method chaining - */ - @Override - public WriteBuilder truncate() { - this.truncate = true; - return this; - } -} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexWriterCommitMessage.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexWriterCommitMessage.java deleted file mode 100644 index c081d261dc3..00000000000 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexWriterCommitMessage.java +++ /dev/null @@ -1,55 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.write; - -import java.io.Serializable; -import org.apache.spark.sql.connector.write.WriterCommitMessage; - -/** - * Commit message containing information about a successfully written Vortex file. - * - *

This message is passed from executors back to the driver to coordinate the commit phase of the write operation. - */ -public record VortexWriterCommitMessage(String filePath, long recordCount, long bytesWritten) - implements WriterCommitMessage, Serializable { - - /** - * Creates a new commit message for a written Vortex file. - * - * @param filePath the path to the written file - * @param recordCount the number of records written - * @param bytesWritten the number of bytes written - */ - public VortexWriterCommitMessage {} - - /** - * Gets the path to the written Vortex file. - * - * @return the file path - */ - @Override - public String filePath() { - return filePath; - } - - /** - * Gets the number of records written to the file. - * - * @return the record count - */ - @Override - public long recordCount() { - return recordCount; - } - - /** - * Gets the number of bytes written to the file. - * - * @return the byte count - */ - @Override - public long bytesWritten() { - return bytesWritten; - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCatalogTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCatalogTest.java deleted file mode 100644 index d08455d376c..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexCatalogTest.java +++ /dev/null @@ -1,114 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark; - -import static org.junit.jupiter.api.Assertions.assertArrayEquals; -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 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.TableChange; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.types.DataTypes; -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 org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link VortexCatalog}, the path-based catalog behind {@code SELECT * FROM vortex.`/path`}. - * - *

Characterizes which identifiers the catalog accepts as paths and which it rejects, plus the DDL surface it - * deliberately does not implement. End-to-end querying through a real session is covered by {@code VortexSqlTest}; - * these tests pin the identifier contract without starting Spark. - */ -final class VortexCatalogTest { - - private static final StructType SCHEMA = - new StructType(new StructField[] {new StructField("id", DataTypes.IntegerType, false, Metadata.empty())}); - - private static VortexCatalog catalog() { - VortexCatalog catalog = new VortexCatalog(); - catalog.initialize("vortex", new CaseInsensitiveStringMap(Map.of())); - return catalog; - } - - @Test - @DisplayName("Takes its name from the session config that registered it") - void nameComesFromInitialize() { - VortexCatalog catalog = new VortexCatalog(); - catalog.initialize("my_vortex", new CaseInsensitiveStringMap(Map.of())); - - assertEquals("my_vortex", catalog.name()); - } - - @Test - @DisplayName("Defaults to the name \"vortex\" before initialize is called") - void nameDefaultsToVortex() { - assertEquals("vortex", new VortexCatalog().name()); - } - - @Test - @DisplayName("Lists no tables: the catalog holds no state, tables are addressed by path") - void listTablesIsAlwaysEmpty() { - assertArrayEquals(new Identifier[0], catalog().listTables(new String[0])); - assertArrayEquals(new Identifier[0], catalog().listTables(new String[] {"any", "namespace"})); - } - - @Test - @DisplayName("An identifier without a path separator is not a table") - void identifierWithoutSlashIsRejected() { - assertThrows(NoSuchTableException.class, () -> catalog().loadTable(Identifier.of(new String[0], "not_a_path"))); - } - - @Test - @DisplayName("A namespaced identifier is not a table, even when the name looks like a path") - void namespacedIdentifierIsRejected() { - assertThrows(NoSuchTableException.class, () -> catalog() - .loadTable(Identifier.of(new String[] {"db"}, "/data/a.vortex"))); - } - - @Test - @DisplayName("A path-shaped identifier that cannot be read surfaces as table not found") - void unreadablePathIsReportedAsMissingTable() { - // Schema inference throws for a path with no Vortex files; the catalog translates that into - // NoSuchTableException so SQL users see "table not found" rather than an internal error. - assertThrows(NoSuchTableException.class, () -> catalog() - .loadTable(Identifier.of(new String[0], "/nonexistent/vortex/path"))); - } - - @Test - @DisplayName("Dropping a table never deletes data, it reports that nothing was dropped") - void dropTableReturnsFalse() { - assertFalse(catalog().dropTable(Identifier.of(new String[0], "/data/a.vortex"))); - } - - @Test - @DisplayName("CREATE TABLE is unsupported: write to the path instead") - void createTableIsUnsupported() { - assertThrows(UnsupportedOperationException.class, () -> catalog() - .createTable(Identifier.of(new String[0], "/data/a.vortex"), SCHEMA, new Transform[0], Map.of())); - } - - @Test - @DisplayName("ALTER TABLE is unsupported: there is no metadata to alter") - void alterTableIsUnsupported() { - assertThrows(UnsupportedOperationException.class, () -> catalog() - .alterTable(Identifier.of(new String[0], "/data/a.vortex"), new TableChange[0])); - } - - @Test - @DisplayName("RENAME TABLE is unsupported: there is no metadata to rename") - void renameTableIsUnsupported() { - assertThrows(UnsupportedOperationException.class, () -> catalog() - .renameTable( - Identifier.of(new String[0], "/data/a.vortex"), - Identifier.of(new String[0], "/data/b.vortex"))); - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java deleted file mode 100644 index 0595349a49a..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java +++ /dev/null @@ -1,241 +0,0 @@ -// 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.assertInstanceOf; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import dev.vortex.spark.read.VortexScan; -import dev.vortex.spark.read.VortexScanBuilder; -import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.Comparator; -import java.util.Map; -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.apache.spark.sql.connector.catalog.Column; -import org.apache.spark.sql.connector.read.Statistics; -import org.apache.spark.sql.types.StructField; -import org.apache.spark.sql.types.StructType; -import org.junit.jupiter.api.AfterAll; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInstance; -import org.junit.jupiter.api.io.TempDir; - -/** - * Integration tests for {@link VortexScan#estimateStatistics()}. - * - *

Verifies that the Spark V2 scan surfaces both the row count Vortex stores in each file footer and the sum of the - * on-storage file sizes reported by the filesystem listing. - */ -@TestInstance(TestInstance.Lifecycle.PER_CLASS) -public final class VortexDataSourceStatsTest { - private static final String FILE_COMPRESSION_FACTOR_KEY = "spark.sql.sources.fileCompressionFactor"; - - 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.shuffle.partitions", "2") - .config("spark.sql.adaptive.enabled", "false") - .config("spark.ui.enabled", "false") - .getOrCreate(); - } - - @AfterAll - public void tearDown() { - if (spark != null) { - spark.stop(); - } - } - - @Test - @DisplayName("VortexScan reports exact row count for single-file scans") - public void testEstimateStatisticsReportsRowCount() throws IOException { - int numRows = 250; - Path outputPath = writeRows(numRows, "single_file"); - - VortexScan scan = buildScan(outputPath); - Statistics stats = scan.estimateStatistics(); - - assertTrue( - stats.numRows().isPresent(), - "VortexScan should report a row count for a Vortex dataset with a populated footer"); - assertEquals(numRows, stats.numRows().getAsLong(), "Row count should match the rows we wrote"); - } - - @Test - @DisplayName("VortexScan reports aggregate row count across multi-file scans") - public void testEstimateStatisticsAcrossMultipleFiles() throws IOException { - int numRows = 400; - Path outputPath = writeRows(numRows, "multi_file", 4); - - VortexScan scan = buildScan(outputPath); - Statistics stats = scan.estimateStatistics(); - - assertTrue(stats.numRows().isPresent(), "Row count should be reported for multi-file Vortex datasets"); - assertEquals(numRows, stats.numRows().getAsLong(), "Row count should sum across all files"); - } - - @Test - @DisplayName("VortexScan reports sizeInBytes equal to the sum of on-storage file sizes") - public void testEstimateStatisticsReportsSizeInBytes() throws IOException { - Path outputPath = writeRows(120, "with_size", 3); - - long fileBytes = totalVortexFileBytes(outputPath); - assertTrue(fileBytes > 0, "Test setup should produce at least one non-empty .vortex file"); - - VortexScan scan = buildScan(outputPath); - Statistics stats = scan.estimateStatistics(); - - assertTrue( - stats.sizeInBytes().isPresent(), - "VortexScan should surface a sizeInBytes when the filesystem listing reports file sizes"); - // Mirror the scan's Spark-convention scaling (factor 1.0, unpruned schema), which divides and - // re-multiplies by the schema default size in double arithmetic before truncating; asserting - // against the raw byte sum would be sensitive to the floating-point round trip. - StructType schema = spark.read() - .format("vortex") - .option("path", outputPath.toUri().toString()) - .load() - .schema(); - long expectedSize = (long) (1.0 * fileBytes / schema.defaultSize() * schema.defaultSize()); - assertEquals( - expectedSize, - stats.sizeInBytes().getAsLong(), - "sizeInBytes should equal the sum of on-storage .vortex file sizes"); - } - - @Test - @DisplayName("VortexScan scales sizeInBytes by the pushed read schema") - public void testEstimateStatisticsScalesSizeInBytesForProjection() throws IOException { - Path outputPath = writeRows(120, "projected_size", 3); - long fileBytes = totalVortexFileBytes(outputPath); - - StructType fullSchema = spark.read() - .format("vortex") - .option("path", outputPath.toUri().toString()) - .load() - .schema(); - StructType idOnlySchema = new StructType(new StructField[] {fullSchema.fields()[0]}); - - String previousCompressionFactor = spark.conf().get(FILE_COMPRESSION_FACTOR_KEY); - spark.conf().set(FILE_COMPRESSION_FACTOR_KEY, "0.5"); - try { - VortexScan scan = buildScan(outputPath, idOnlySchema); - Statistics stats = scan.estimateStatistics(); - - long expectedSize = (long) (0.5 * fileBytes / fullSchema.defaultSize() * idOnlySchema.defaultSize()); - assertTrue(stats.sizeInBytes().isPresent(), "Projected scans should still surface sizeInBytes"); - assertEquals( - expectedSize, - stats.sizeInBytes().getAsLong(), - "sizeInBytes should follow Spark FileScan's compression and schema-width scaling"); - assertTrue( - stats.sizeInBytes().getAsLong() < fileBytes, - "Projected scan stats should be smaller than full file bytes"); - } finally { - spark.conf().set(FILE_COMPRESSION_FACTOR_KEY, previousCompressionFactor); - } - } - - @Test - @DisplayName("VortexScan caches statistics across repeated calls") - public void testEstimateStatisticsIsCached() throws IOException { - Path outputPath = writeRows(50, "cached", 1); - - VortexScan scan = buildScan(outputPath); - Statistics first = scan.estimateStatistics(); - Statistics second = scan.estimateStatistics(); - - // Same instance returned -- the second call hits the cached value. - assertEquals(first, second, "estimateStatistics() should return the same Statistics object on repeat calls"); - assertInstanceOf(Statistics.class, first); - } - - private VortexScan buildScan(Path outputPath) { - return buildScan(outputPath, null); - } - - private VortexScan buildScan(Path outputPath, StructType requiredSchema) { - Dataset readDf = spark.read() - .format("vortex") - .option("path", outputPath.toUri().toString()) - .load(); - StructType readSchema = readDf.schema(); - - VortexScanBuilder builder = new VortexScanBuilder(Map.of()); - builder.addPath(outputPath.toUri().toString()); - for (StructField field : readSchema.fields()) { - builder.addColumn(Column.create(field.name(), field.dataType())); - } - if (requiredSchema != null) { - builder.pruneColumns(requiredSchema); - } - return (VortexScan) builder.build(); - } - - private Path writeRows(int numRows, String name) throws IOException { - return writeRows(numRows, name, 1); - } - - private long totalVortexFileBytes(Path outputPath) throws IOException { - try (Stream paths = Files.walk(outputPath)) { - return paths.filter(Files::isRegularFile) - .filter(path -> path.getFileName().toString().endsWith(".vortex")) - .mapToLong(path -> { - try { - return Files.size(path); - } catch (IOException e) { - throw new RuntimeException(e); - } - }) - .sum(); - } - } - - private Path writeRows(int numRows, String name, int partitions) throws IOException { - Path outputPath = tempDir.resolve(name); - Dataset df = spark.range(0, numRows) - .selectExpr("cast(id as int) as id", "concat('value_', cast(id as string)) as value"); - - df.repartition(partitions) - .write() - .format("vortex") - .option("path", outputPath.toUri().toString()) - .mode(SaveMode.Overwrite) - .save(); - return outputPath; - } - - @AfterEach - public void cleanupTempFiles() throws IOException { - if (tempDir != null && Files.exists(tempDir)) { - try (Stream paths = Files.walk(tempDir)) { - paths.sorted(Comparator.reverseOrder()).forEach(path -> { - try { - Files.deleteIfExists(path); - } catch (IOException e) { - System.err.println("Failed to delete: " + path); - } - }); - } - } - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSessionCatalogTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSessionCatalogTest.java deleted file mode 100644 index b31d674d884..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSessionCatalogTest.java +++ /dev/null @@ -1,113 +0,0 @@ -// 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.Files; -import java.nio.file.Path; -import java.util.List; -import org.apache.spark.sql.Row; -import org.apache.spark.sql.SparkSession; -import org.junit.jupiter.api.AfterAll; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInstance; -import org.junit.jupiter.api.io.TempDir; - -/** - * Integration tests for {@link VortexSessionCatalog}, the session catalog extension that makes {@code CREATE TABLE ... - * USING vortex} work on Spark 3.5 as well as Spark 4. - */ -@TestInstance(TestInstance.Lifecycle.PER_CLASS) -public final class VortexSessionCatalogTest { - - private SparkSession spark; - private Path warehouseDir; - - @TempDir - Path tempDir; - - @BeforeAll - public void setUp() throws IOException { - warehouseDir = Files.createTempDirectory("vortex-warehouse"); - spark = SparkSession.builder() - .appName("VortexSessionCatalogTest") - .master("local[2]") - .config("spark.driver.host", "127.0.0.1") - .config("spark.sql.warehouse.dir", warehouseDir.toUri().toString()) - .config("spark.sql.catalog.spark_catalog", VortexSessionCatalog.class.getName()) - .config("spark.ui.enabled", "false") - .getOrCreate(); - } - - @AfterAll - public void tearDown() { - if (spark != null) { - spark.stop(); - } - } - - @Test - @DisplayName("Managed table lifecycle: CREATE, SELECT while empty, INSERT, INSERT OVERWRITE, DROP") - public void testManagedTableLifecycle() { - 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"); - - spark.sql("INSERT INTO managed_students VALUES (1, 'Alice', 20), (2, 'Bob', 21)"); - List rows = spark.sql("SELECT name FROM managed_students WHERE age > 20 ORDER BY name") - .collectAsList(); - assertEquals(1, rows.size()); - assertEquals("Bob", rows.get(0).getString(0)); - - spark.sql("INSERT OVERWRITE managed_students VALUES (3, 'Carol', 22)"); - assertEquals(1, spark.sql("SELECT * FROM managed_students").count(), "Overwrite should replace all rows"); - - Path tableDir = warehouseDir.resolve("managed_students"); - assertTrue(Files.exists(tableDir), "Managed table data should live under the warehouse dir"); - spark.sql("DROP TABLE managed_students"); - assertFalse(Files.exists(tableDir), "Dropping a managed table should remove its data"); - } - - @Test - @DisplayName("CREATE TABLE AS SELECT without a LOCATION clause") - public void testCreateManagedTableAsSelect() { - spark.sql("CREATE TABLE ctas_source (id INT, name STRING) USING vortex"); - spark.sql("INSERT INTO ctas_source VALUES (1, 'Alice'), (2, 'Bob')"); - - spark.sql("CREATE TABLE ctas_target USING vortex AS SELECT * FROM ctas_source WHERE id > 1"); - List rows = spark.sql("SELECT name FROM ctas_target").collectAsList(); - assertEquals(1, rows.size()); - assertEquals("Bob", rows.get(0).getString(0)); - - spark.sql("DROP TABLE ctas_target"); - spark.sql("DROP TABLE ctas_source"); - } - - @Test - @DisplayName("External table with a LOCATION clause") - public void testExternalTable() { - Path location = tempDir.resolve("ext_students"); - spark.sql( - String.format("CREATE TABLE ext_students (id INT, name STRING) USING vortex LOCATION '%s'", location)); - spark.sql("INSERT INTO ext_students VALUES (1, 'Alice')"); - assertEquals(1, spark.sql("SELECT * FROM ext_students").count()); - spark.sql("DROP TABLE ext_students"); - } - - @Test - @DisplayName("Tables of other providers pass through the extension untouched") - public void testOtherProviderPassthrough() { - spark.sql("CREATE TABLE pq_table (id INT) USING parquet"); - spark.sql("INSERT INTO pq_table VALUES (7)"); - assertEquals( - 7, spark.sql("SELECT * FROM pq_table").collectAsList().get(0).getInt(0)); - spark.sql("DROP TABLE pq_table"); - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexTableTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexTableTest.java deleted file mode 100644 index 8c23d6d874b..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexTableTest.java +++ /dev/null @@ -1,143 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark; - -import static org.junit.jupiter.api.Assertions.assertArrayEquals; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertInstanceOf; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import com.google.common.collect.ImmutableList; -import dev.vortex.spark.read.VortexScanBuilder; -import dev.vortex.spark.write.VortexWriteBuilder; -import java.util.Map; -import java.util.NoSuchElementException; -import org.apache.spark.sql.connector.catalog.TableCapability; -import org.apache.spark.sql.connector.expressions.Expressions; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.connector.read.Scan; -import org.apache.spark.sql.connector.write.LogicalWriteInfo; -import org.apache.spark.sql.types.DataTypes; -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 org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link VortexTable}, the Spark V2 table of Vortex files. - * - *

Characterizes the table-level contract Spark relies on: which capabilities are advertised, how the table name is - * rendered from its paths, and the read/write builders it hands back. Notably, writes accept exactly one path, so - * {@link VortexTable#newWriteBuilder} rejects tables built from several paths. - */ -final class VortexTableTest { - - private static final StructType SCHEMA = new StructType(new StructField[] { - new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), - new StructField("name", DataTypes.StringType, true, Metadata.empty()) - }); - - private static VortexTable tableFor(String... paths) { - return new VortexTable(ImmutableList.copyOf(paths), SCHEMA, Map.of(), new Transform[0]); - } - - @Test - @DisplayName("Advertises exactly batch read, batch write and truncate") - void capabilitiesAreBatchReadWriteAndTruncate() { - assertEquals( - java.util.Set.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE), - tableFor("/data/a.vortex").capabilities()); - } - - @Test - @DisplayName("Table name joins every path with commas under the vortex prefix") - void nameJoinsAllPaths() { - assertEquals("vortex.\"/data/a.vortex\"", tableFor("/data/a.vortex").name()); - assertEquals( - "vortex.\"/data/a.vortex,/data/b.vortex\"", - tableFor("/data/a.vortex", "/data/b.vortex").name()); - } - - @Test - @DisplayName("Schema and partitioning are returned as supplied") - void schemaAndPartitioningRoundTrip() { - Transform[] transforms = new Transform[] {Expressions.identity("year")}; - VortexTable table = new VortexTable(ImmutableList.of("/tbl"), SCHEMA, Map.of(), transforms); - - assertEquals(SCHEMA, table.schema()); - assertArrayEquals(transforms, table.partitioning()); - } - - @Test - @DisplayName("Scan builder projects the full table schema by default") - void scanBuilderStartsFromTableSchema() { - VortexTable table = tableFor("/data/a.vortex"); - - var builder = table.newScanBuilder(new CaseInsensitiveStringMap(Map.of())); - - assertInstanceOf(VortexScanBuilder.class, builder); - Scan scan = builder.build(); - assertEquals(SCHEMA, scan.readSchema()); - } - - @Test - @DisplayName("Scan description carries every table path") - void scanDescriptionCarriesPaths() { - VortexTable table = tableFor("/data/a.vortex", "/data/b.vortex"); - - Scan scan = table.newScanBuilder(new CaseInsensitiveStringMap(Map.of())).build(); - - String description = scan.description(); - assertTrue(description.startsWith("VortexScan{"), description); - assertTrue(description.contains("/data/a.vortex"), description); - assertTrue(description.contains("/data/b.vortex"), description); - assertTrue(description.contains("pushedPredicates=[]"), description); - } - - @Test - @DisplayName("Write builder is created for a single-path table") - void writeBuilderAcceptsSinglePath() { - VortexTable table = tableFor("/data/out"); - - assertInstanceOf(VortexWriteBuilder.class, table.newWriteBuilder(writeInfo())); - } - - @Test - @DisplayName("Writing a table of several paths is rejected: there is no single output path") - void writeBuilderRejectsMultiplePaths() { - VortexTable table = tableFor("/data/a.vortex", "/data/b.vortex"); - - assertThrows(IllegalArgumentException.class, () -> table.newWriteBuilder(writeInfo())); - } - - @Test - @DisplayName("Writing a table with no path is rejected") - void writeBuilderRejectsNoPaths() { - VortexTable table = tableFor(); - - assertThrows(NoSuchElementException.class, () -> table.newWriteBuilder(writeInfo())); - } - - private static LogicalWriteInfo writeInfo() { - return new LogicalWriteInfo() { - @Override - public String queryId() { - return "query-1"; - } - - @Override - public StructType schema() { - return SCHEMA; - } - - @Override - public CaseInsensitiveStringMap options() { - return new CaseInsensitiveStringMap(Map.of()); - } - }; - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/config/HadoopUtilsTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/config/HadoopUtilsTest.java deleted file mode 100644 index 234db1a5ed6..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/config/HadoopUtilsTest.java +++ /dev/null @@ -1,150 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.config; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.Map; -import org.apache.hadoop.conf.Configuration; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link HadoopUtils}, which forwards S3 and Azure credentials from a Hadoop {@link Configuration} to - * the Vortex object-store property keys. - * - *

Characterizes the {@code fs.s3a.*} to {@code aws_*} key mapping (including https-qualification of bare endpoints), - * the prefix-based Azure account key and SAS token extraction, and the skip-signature fallback for anonymous Azure - * access. - */ -final class HadoopUtilsTest { - /** A Configuration that does not load the default resources, so tests see only the keys they set. */ - private static Configuration emptyConf() { - return new Configuration(/* loadDefaults= */ false); - } - - // --- s3PropertiesFromHadoopConf --- - - @Test - @DisplayName("Maps fs.s3a credentials, session token, and region to the aws_* property keys") - void s3MapsCredentialKeys() { - Configuration conf = emptyConf(); - conf.set(HadoopUtils.FS_S3A_ACCESS_KEY, "AKIAEXAMPLE"); - conf.set(HadoopUtils.FS_S3A_SECRET_KEY, "secret"); - conf.set(HadoopUtils.FS_S3A_SESSION_TOKEN, "token"); - conf.set(HadoopUtils.FS_S3A_ENDPOINT_REGION, "us-west-2"); - - Map properties = HadoopUtils.s3PropertiesFromHadoopConf(conf); - - assertEquals("AKIAEXAMPLE", properties.get("aws_access_key_id")); - assertEquals("secret", properties.get("aws_secret_access_key")); - assertEquals("token", properties.get("aws_session_token")); - assertEquals("us-west-2", properties.get("aws_region")); - } - - @Test - @DisplayName("Qualifies a bare fs.s3a.endpoint with https://") - void s3QualifiesBareEndpoint() { - Configuration conf = emptyConf(); - conf.set(HadoopUtils.FS_S3A_ENDPOINT, "s3.us-west-2.amazonaws.com"); - - Map properties = HadoopUtils.s3PropertiesFromHadoopConf(conf); - - assertEquals("https://s3.us-west-2.amazonaws.com", properties.get("aws_endpoint")); - } - - @Test - @DisplayName("Preserves an fs.s3a.endpoint that already has an http or https scheme") - void s3PreservesSchemedEndpoint() { - Configuration httpConf = emptyConf(); - httpConf.set(HadoopUtils.FS_S3A_ENDPOINT, "http://localhost:9000"); - Map httpProperties = HadoopUtils.s3PropertiesFromHadoopConf(httpConf); - assertEquals("http://localhost:9000", httpProperties.get("aws_endpoint")); - assertEquals("true", httpProperties.get("aws_allow_http")); - - Configuration httpsConf = emptyConf(); - httpsConf.set(HadoopUtils.FS_S3A_ENDPOINT, "https://s3.example.com"); - Map httpsProperties = HadoopUtils.s3PropertiesFromHadoopConf(httpsConf); - assertEquals("https://s3.example.com", httpsProperties.get("aws_endpoint")); - assertNull(httpsProperties.get("aws_allow_http")); - } - - @Test - @DisplayName("Ignores Hadoop keys that are not S3-relevant") - void s3IgnoresUnrelatedKeys() { - Configuration conf = emptyConf(); - conf.set("fs.defaultFS", "hdfs://namenode:8020"); - conf.set("fs.s3a.connection.maximum", "64"); - - assertTrue(HadoopUtils.s3PropertiesFromHadoopConf(conf).isEmpty()); - } - - @Test - @DisplayName("Returns no properties for an empty configuration") - void s3EmptyConfYieldsNoProperties() { - assertTrue(HadoopUtils.s3PropertiesFromHadoopConf(emptyConf()).isEmpty()); - } - - // --- azurePropertiesFromHadoopConf --- - - @Test - @DisplayName("Extracts the storage account key from any fs.azure.account.key-prefixed entry") - void azureExtractsAccountKey() { - Configuration conf = emptyConf(); - conf.set(HadoopUtils.ACCESS_KEY_PREFIX + ".myaccount.dfs.core.windows.net", "account-key"); - - Map properties = HadoopUtils.azurePropertiesFromHadoopConf(conf); - - assertEquals("account-key", properties.get("azure_storage_account_key")); - assertFalse(properties.containsKey("azure_skip_signature")); - } - - @Test - @DisplayName("Extracts the SAS token from any fs.azure.sas.fixed.token-prefixed entry") - void azureExtractsSasToken() { - Configuration conf = emptyConf(); - conf.set(HadoopUtils.FIXED_TOKEN_PREFIX + "myaccount.dfs.core.windows.net", "sas-token"); - - Map properties = HadoopUtils.azurePropertiesFromHadoopConf(conf); - - assertEquals("sas-token", properties.get("azure_storage_sas_key")); - } - - @Test - @DisplayName("Falls back to skipping signatures when no account key is configured") - void azureSkipsSignatureWithoutAccountKey() { - Map properties = HadoopUtils.azurePropertiesFromHadoopConf(emptyConf()); - - assertEquals("true", properties.get("azure_skip_signature")); - assertFalse(properties.containsKey("azure_storage_account_key")); - } - - @Test - @DisplayName("Skips signatures even when only a SAS token is configured") - void azureSasOnlyStillSkipsSignature() { - Configuration conf = emptyConf(); - conf.set(HadoopUtils.FIXED_TOKEN_PREFIX + "myaccount.dfs.core.windows.net", "sas-token"); - - Map properties = HadoopUtils.azurePropertiesFromHadoopConf(conf); - - assertEquals("sas-token", properties.get("azure_storage_sas_key")); - assertEquals("true", properties.get("azure_skip_signature")); - } - - @Test - @DisplayName("Ignores Hadoop keys that are not Azure-relevant") - void azureIgnoresUnrelatedKeys() { - Configuration conf = emptyConf(); - conf.set("fs.azure.io.retry.max.retries", "5"); - conf.set("fs.defaultFS", "abfss://container@account.dfs.core.windows.net/"); - - Map properties = HadoopUtils.azurePropertiesFromHadoopConf(conf); - - assertFalse(properties.containsKey("azure_storage_account_key")); - assertFalse(properties.containsKey("azure_storage_sas_key")); - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/PartitionPathUtilsTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/PartitionPathUtilsTest.java deleted file mode 100644 index 445fad2f56c..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/read/PartitionPathUtilsTest.java +++ /dev/null @@ -1,199 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.read; - -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.util.List; -import java.util.Map; -import org.apache.spark.sql.execution.vectorized.ConstantColumnVector; -import org.apache.spark.sql.types.DataTypes; -import org.apache.spark.unsafe.types.UTF8String; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link PartitionPathUtils}, which discovers and materializes Hive-style partition columns from file - * paths. - * - *

Characterizes partition-value parsing (including URL decoding and the {@code __HIVE_DEFAULT_PARTITION__} - * sentinel), partition column type inference, and constant-vector materialization for each supported type. - */ -final class PartitionPathUtilsTest { - // --- parsePartitionValues --- - - @Test - @DisplayName("Parses key=value segments from a Hive-style partition path") - void parsesKeyValueSegments() { - Map values = - PartitionPathUtils.parsePartitionValues("/data/warehouse/year=2024/month=12/part-0.vortex"); - assertEquals(Map.of("year", "2024", "month", "12"), values); - } - - @Test - @DisplayName("Preserves the order of partition segments in the path") - void preservesSegmentOrder() { - Map values = PartitionPathUtils.parsePartitionValues("/tbl/b=2/a=1/c=3/file.vortex"); - assertEquals(List.of("b", "a", "c"), List.copyOf(values.keySet())); - } - - @Test - @DisplayName("Ignores segments without key=value shape") - void ignoresNonPartitionSegments() { - Map values = PartitionPathUtils.parsePartitionValues("/data/plain/dir/file.vortex"); - assertTrue(values.isEmpty()); - } - - @Test - @DisplayName("Ignores segments with an empty key or empty value") - void ignoresEmptyKeyOrValue() { - assertTrue(PartitionPathUtils.parsePartitionValues("/tbl/=v/file").isEmpty()); - assertTrue(PartitionPathUtils.parsePartitionValues("/tbl/k=/file").isEmpty()); - } - - @Test - @DisplayName("URL-decodes both keys and values") - void urlDecodesKeysAndValues() { - Map values = PartitionPathUtils.parsePartitionValues("/tbl/city=New%20York/file.vortex"); - assertEquals(Map.of("city", "New York"), values); - } - - @Test - @DisplayName("Splits on the first '=' so values may contain '='") - void splitsOnFirstEquals() { - Map values = PartitionPathUtils.parsePartitionValues("/tbl/k=a=b/file.vortex"); - assertEquals(Map.of("k", "a=b"), values); - } - - // --- inferPartitionColumnType --- - - @Test - @DisplayName("Infers IntegerType for values that fit in an int") - void infersInteger() { - assertEquals(DataTypes.IntegerType, PartitionPathUtils.inferPartitionColumnType("42")); - assertEquals(DataTypes.IntegerType, PartitionPathUtils.inferPartitionColumnType("-7")); - } - - @Test - @DisplayName("Infers LongType for integral values wider than an int") - void infersLong() { - assertEquals(DataTypes.LongType, PartitionPathUtils.inferPartitionColumnType("9999999999")); - } - - @Test - @DisplayName("Infers DoubleType for decimal-point values") - void infersDouble() { - assertEquals(DataTypes.DoubleType, PartitionPathUtils.inferPartitionColumnType("3.14")); - } - - @Test - @DisplayName("Infers BooleanType for true/false in any case") - void infersBoolean() { - assertEquals(DataTypes.BooleanType, PartitionPathUtils.inferPartitionColumnType("true")); - assertEquals(DataTypes.BooleanType, PartitionPathUtils.inferPartitionColumnType("FALSE")); - } - - @Test - @DisplayName("Falls back to StringType for everything else") - void fallsBackToString() { - assertEquals(DataTypes.StringType, PartitionPathUtils.inferPartitionColumnType("hello")); - assertEquals(DataTypes.StringType, PartitionPathUtils.inferPartitionColumnType("2024-12-01")); - } - - @Test - @DisplayName("Null and __HIVE_DEFAULT_PARTITION__ infer StringType") - void nullAndDefaultPartitionInferString() { - assertEquals(DataTypes.StringType, PartitionPathUtils.inferPartitionColumnType(null)); - assertEquals(DataTypes.StringType, PartitionPathUtils.inferPartitionColumnType("__HIVE_DEFAULT_PARTITION__")); - } - - // --- createConstantVector --- - - @Test - @DisplayName("Null and __HIVE_DEFAULT_PARTITION__ materialize as null vectors") - void nullValuesMaterializeAsNull() { - ConstantColumnVector fromNull = PartitionPathUtils.createConstantVector(3, DataTypes.StringType, null); - assertTrue(fromNull.isNullAt(0)); - - ConstantColumnVector fromSentinel = - PartitionPathUtils.createConstantVector(3, DataTypes.StringType, "__HIVE_DEFAULT_PARTITION__"); - assertTrue(fromSentinel.isNullAt(0)); - } - - @Test - @DisplayName("String values materialize as UTF8 strings") - void stringMaterializes() { - ConstantColumnVector vec = PartitionPathUtils.createConstantVector(2, DataTypes.StringType, "us-east"); - assertFalse(vec.isNullAt(0)); - assertEquals(UTF8String.fromString("us-east"), vec.getUTF8String(0)); - } - - @Test - @DisplayName("Integral values materialize with the width of the target type") - void integralTypesMaterialize() { - assertEquals( - 42, - PartitionPathUtils.createConstantVector(1, DataTypes.IntegerType, "42") - .getInt(0)); - assertEquals( - 9999999999L, - PartitionPathUtils.createConstantVector(1, DataTypes.LongType, "9999999999") - .getLong(0)); - assertEquals( - (short) 7, - PartitionPathUtils.createConstantVector(1, DataTypes.ShortType, "7") - .getShort(0)); - assertEquals( - (byte) 3, - PartitionPathUtils.createConstantVector(1, DataTypes.ByteType, "3") - .getByte(0)); - } - - @Test - @DisplayName("DateType parses the value as days since epoch (int)") - void dateMaterializesAsInt() { - assertEquals( - 19000, - PartitionPathUtils.createConstantVector(1, DataTypes.DateType, "19000") - .getInt(0)); - } - - @Test - @DisplayName("Timestamp types parse the value as micros since epoch (long)") - void timestampMaterializesAsLong() { - assertEquals( - 1700000000000000L, - PartitionPathUtils.createConstantVector(1, DataTypes.TimestampType, "1700000000000000") - .getLong(0)); - assertEquals( - 1700000000000000L, - PartitionPathUtils.createConstantVector(1, DataTypes.TimestampNTZType, "1700000000000000") - .getLong(0)); - } - - @Test - @DisplayName("Boolean, float, and double values materialize with their target types") - void booleanAndFloatingPointMaterialize() { - assertTrue(PartitionPathUtils.createConstantVector(1, DataTypes.BooleanType, "true") - .getBoolean(0)); - assertEquals( - 1.5f, - PartitionPathUtils.createConstantVector(1, DataTypes.FloatType, "1.5") - .getFloat(0)); - assertEquals( - 2.25, - PartitionPathUtils.createConstantVector(1, DataTypes.DoubleType, "2.25") - .getDouble(0)); - } - - @Test - @DisplayName("Unrecognized target types fall back to UTF8 strings") - void unrecognizedTypeFallsBackToString() { - ConstantColumnVector vec = - PartitionPathUtils.createConstantVector(1, DataTypes.CalendarIntervalType, "whatever"); - assertEquals(UTF8String.fromString("whatever"), vec.getUTF8String(0)); - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/SparkPredicateToVortexExpressionTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/SparkPredicateToVortexExpressionTest.java deleted file mode 100644 index 7be565f3553..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/read/SparkPredicateToVortexExpressionTest.java +++ /dev/null @@ -1,374 +0,0 @@ -// 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 dev.vortex.jni.NativeLoader; -import java.math.BigDecimal; -import java.util.List; -import java.util.Map; -import org.apache.spark.sql.connector.expressions.Expression; -import org.apache.spark.sql.connector.expressions.LiteralValue; -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.DataType; -import org.apache.spark.sql.types.DataTypes; -import org.apache.spark.sql.types.Decimal; -import org.apache.spark.sql.types.StructType; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link SparkPredicateToVortexExpression#isPushable(Predicate, Map)} and - * {@link SparkPredicateToVortexExpression#convert(Predicate)}. - * - *

{@code isPushable} decides which predicates {@code VortexScanBuilder.pushPredicates} lets Spark drop, while - * {@code convert} builds the filter {@code VortexPartitionReader} actually pushes down. A predicate that passes the - * first but fails the second is silently skipped by the reader, so the scan returns rows the query excluded. The tests - * below pin that {@code isPushable} implies {@code convert().isPresent()} across every accepted shape. - */ -final class SparkPredicateToVortexExpressionTest { - - private static final StructType ADDRESS = DataTypes.createStructType(new org.apache.spark.sql.types.StructField[] { - DataTypes.createStructField("city", DataTypes.StringType, true), - DataTypes.createStructField("zip", DataTypes.IntegerType, true) - }); - - private static final StructType PROFILE = DataTypes.createStructType(new org.apache.spark.sql.types.StructField[] { - DataTypes.createStructField("email", DataTypes.StringType, true), - DataTypes.createStructField("address", ADDRESS, true) - }); - - private static final DataType DECIMAL = DataTypes.createDecimalType(10, 2); - - /** One data column per literal type {@code convertLiteral} maps, so every literal meets a same-typed column. */ - private static final Map SCHEMA = Map.ofEntries( - Map.entry("id", DataTypes.IntegerType), - Map.entry("name", DataTypes.StringType), - Map.entry("active", DataTypes.BooleanType), - Map.entry("tiny", DataTypes.ByteType), - Map.entry("small", DataTypes.ShortType), - Map.entry("big", DataTypes.LongType), - Map.entry("ratio", DataTypes.FloatType), - Map.entry("weight", DataTypes.DoubleType), - Map.entry("payload", DataTypes.BinaryType), - Map.entry("birthday", DataTypes.DateType), - Map.entry("createdAt", DataTypes.TimestampType), - Map.entry("createdLocal", DataTypes.TimestampNTZType), - Map.entry("amount", DECIMAL), - Map.entry("profile", PROFILE)); - - private static final List COMPARISON_OPERATORS = List.of("=", "<>", "!=", ">", ">=", "<", "<="); - - /** The columns whose type {@code isPushableLiteral} accepts a {@code null} value for. */ - private static final List NULLABLE_LITERAL_COLUMNS = List.of( - "active", - "tiny", - "small", - "id", - "big", - "ratio", - "weight", - "name", - "payload", - "birthday", - "createdAt", - "createdLocal", - "amount"); - - @BeforeAll - static void loadNativeLibrary() { - // `convert` allocates native expressions; `isPushable` does not. - NativeLoader.loadJni(); - } - - @Test - @DisplayName("Top-level column reference is pushable when present in the schema") - void topLevelColumnIsPushable() { - Predicate equality = equality(ref("id"), literal(42)); - assertTrue(SparkPredicateToVortexExpression.isPushable(equality, SCHEMA)); - } - - @Test - @DisplayName("Top-level column reference is not pushable when absent from the schema") - void unknownTopLevelColumnIsNotPushable() { - Predicate equality = equality(ref("missing"), literal(0)); - assertFalse(SparkPredicateToVortexExpression.isPushable(equality, SCHEMA)); - } - - @Test - @DisplayName("Nested field reference is pushable when every part resolves under struct types") - void nestedFieldThatExistsIsPushable() { - Predicate equality = equality(ref("profile", "email"), literal("a@b.com")); - assertTrue(SparkPredicateToVortexExpression.isPushable(equality, SCHEMA)); - } - - @Test - @DisplayName("Doubly nested field reference resolves through multiple struct levels") - void doublyNestedFieldIsPushable() { - Predicate equality = equality(ref("profile", "address", "zip"), literal(12345)); - assertTrue(SparkPredicateToVortexExpression.isPushable(equality, SCHEMA)); - } - - @Test - @DisplayName("Nested field that does not exist in the struct is not pushable") - void nestedFieldThatDoesNotExistIsNotPushable() { - Predicate equality = equality(ref("profile", "phone"), literal("555")); - assertFalse(SparkPredicateToVortexExpression.isPushable(equality, SCHEMA)); - } - - @Test - @DisplayName("Descending past a leaf (non-struct) field is not pushable") - void descendingPastLeafFieldIsNotPushable() { - // `name` is a String, not a struct — `name.first` cannot resolve. - Predicate equality = equality(ref("name", "first"), literal("alice")); - assertFalse(SparkPredicateToVortexExpression.isPushable(equality, SCHEMA)); - } - - @Test - @DisplayName("Empty named reference is not pushable") - void emptyReferenceIsNotPushable() { - Predicate equality = equality(ref(), literal(1)); - assertFalse(SparkPredicateToVortexExpression.isPushable(equality, SCHEMA)); - } - - @Test - @DisplayName("Every accepted comparison operator converts, with the column on either side") - void everyComparisonOperatorConverts() { - for (String op : COMPARISON_OPERATORS) { - assertPushableAndConvertible(predicate(op, ref("id"), literal(42)), op + " with column on the left"); - // Spark's V2 builder sometimes commutes; `convertComparison` swaps the operator back. - assertPushableAndConvertible(predicate(op, literal(42), ref("id")), op + " with column on the right"); - } - } - - @Test - @DisplayName("Comparison between two columns converts") - void columnToColumnComparisonConverts() { - assertPushableAndConvertible(equality(ref("id"), ref("profile", "address", "zip"))); - } - - @Test - @DisplayName("Comparison between two literals is rejected by both sides") - void literalToLiteralComparisonIsRejected() { - assertNotPushableAndNotConvertible(equality(literal(1), literal(2))); - } - - @Test - @DisplayName("Comparison with the wrong number of children is rejected by both sides") - void comparisonWithWrongArityIsRejected() { - assertNotPushableAndNotConvertible(predicate("=", ref("id"))); - assertNotPushableAndNotConvertible(predicate("=", ref("id"), literal(1), literal(2))); - } - - @Test - @DisplayName("Nested column comparison converts") - void nestedColumnComparisonConverts() { - assertPushableAndConvertible(equality(ref("profile", "email"), literal("a@b.com"))); - } - - @Test - @DisplayName("IS_NULL and IS_NOT_NULL convert for top-level and nested columns") - void nullChecksConvert() { - assertPushableAndConvertible(predicate("IS_NULL", ref("name"))); - assertPushableAndConvertible(predicate("IS_NOT_NULL", ref("name"))); - assertPushableAndConvertible(predicate("IS_NULL", ref("profile", "address", "city"))); - } - - @Test - @DisplayName("IN converts for a single literal and for many literals") - void inConverts() { - assertPushableAndConvertible(predicate("IN", ref("id"), literal(1))); - assertPushableAndConvertible(predicate("IN", ref("id"), literal(1), literal(2), literal(3))); - } - - @Test - @DisplayName("IN with no literals is rejected by both sides") - void inWithoutLiteralsIsRejected() { - assertNotPushableAndNotConvertible(predicate("IN", ref("id"))); - } - - @Test - @DisplayName("String matching predicates convert, including LIKE meta-characters in the needle") - void stringMatchesConvert() { - for (String name : List.of("STARTS_WITH", "ENDS_WITH", "CONTAINS")) { - assertPushableAndConvertible(predicate(name, ref("name"), literal("ali")), name); - // `buildLikePattern` escapes `%`, `_` and `\` so the match stays an exact substring. - assertPushableAndConvertible(predicate(name, ref("name"), literal("100%_a\\b")), name + " with escapes"); - } - } - - @Test - @DisplayName("String matching against a non-string literal is rejected by both sides") - void stringMatchAgainstNonStringLiteralIsRejected() { - assertNotPushableAndNotConvertible(predicate("STARTS_WITH", ref("name"), literal(1))); - } - - @Test - @DisplayName("BOOLEAN_EXPRESSION over a column reference converts") - void bareBooleanColumnConverts() { - assertPushableAndConvertible(predicate("BOOLEAN_EXPRESSION", ref("active"))); - } - - @Test - @DisplayName("An unrecognised predicate name is rejected by both sides") - void unknownPredicateNameIsRejected() { - assertNotPushableAndNotConvertible(predicate("BLOOM_FILTER", ref("id"), literal(1))); - } - - @Test - @DisplayName("AlwaysTrue and AlwaysFalse convert to boolean literals") - void constantPredicatesConvert() { - assertPushableAndConvertible(new AlwaysTrue()); - assertPushableAndConvertible(new AlwaysFalse()); - } - - @Test - @DisplayName("AND, OR and NOT convert when every leaf converts") - void compoundPredicatesConvert() { - Predicate left = equality(ref("id"), literal(1)); - Predicate right = predicate("IS_NOT_NULL", ref("name")); - assertPushableAndConvertible(new And(left, right)); - assertPushableAndConvertible(new Or(left, right)); - assertPushableAndConvertible(new Not(left)); - assertPushableAndConvertible(new Not(new And(left, new Or(right, new AlwaysFalse())))); - } - - @Test - @DisplayName("A compound predicate with one unconvertible leaf is rejected by both sides") - void compoundPredicateWithBadLeafIsRejected() { - Predicate good = equality(ref("id"), literal(1)); - Predicate bad = predicate("BLOOM_FILTER", ref("id"), literal(1)); - assertNotPushableAndNotConvertible(new And(good, bad)); - assertNotPushableAndNotConvertible(new Or(bad, good)); - assertNotPushableAndNotConvertible(new Not(bad)); - } - - @Test - @DisplayName("Every literal type that accepts a null value converts") - void nullLiteralsConvert() { - for (String column : NULLABLE_LITERAL_COLUMNS) { - assertPushableAndConvertible( - equality(ref(column), new LiteralValue<>(null, SCHEMA.get(column))), "null literal for " + column); - } - } - - @Test - @DisplayName("Every non-null literal type the translator maps converts") - void nonNullLiteralsConvert() { - assertPushableAndConvertible(equality(ref("active"), new LiteralValue<>(true, DataTypes.BooleanType))); - assertPushableAndConvertible(equality(ref("tiny"), new LiteralValue<>((byte) 1, DataTypes.ByteType))); - assertPushableAndConvertible(equality(ref("small"), new LiteralValue<>((short) 1, DataTypes.ShortType))); - assertPushableAndConvertible(equality(ref("id"), literal(42))); - assertPushableAndConvertible(equality(ref("big"), new LiteralValue<>(1L, DataTypes.LongType))); - assertPushableAndConvertible(equality(ref("ratio"), new LiteralValue<>(1.5f, DataTypes.FloatType))); - assertPushableAndConvertible(equality(ref("weight"), new LiteralValue<>(1.5d, DataTypes.DoubleType))); - assertPushableAndConvertible(equality(ref("name"), literal("alice"))); - assertPushableAndConvertible( - equality(ref("payload"), new LiteralValue<>(new byte[] {1, 2, 3}, DataTypes.BinaryType))); - // Spark encodes DateType as an epoch-day int and both timestamp types as epoch micros. - assertPushableAndConvertible(equality(ref("birthday"), new LiteralValue<>(19_000, DataTypes.DateType))); - assertPushableAndConvertible( - equality(ref("createdAt"), new LiteralValue<>(1_700_000_000L, DataTypes.TimestampType))); - assertPushableAndConvertible( - equality(ref("createdLocal"), new LiteralValue<>(1_700_000_000L, DataTypes.TimestampNTZType))); - assertPushableAndConvertible(equality(ref("amount"), decimalLiteral("12.34"))); - } - - @Test - @DisplayName("A literal type with no Vortex representation is rejected by both sides") - void unrepresentableLiteralIsRejected() { - assertNotPushableAndNotConvertible(equality(ref("id"), new LiteralValue<>(null, DataTypes.NullType))); - } - - @Test - @DisplayName("A decimal literal that does not fit the declared scale is rejected by both sides") - void decimalThatDoesNotFitTheScaleIsRejected() { - // `unscaledValueOf` calls `setScale(2)` without a rounding mode, so 12.345 throws and - // `isPushableLiteral` falls through to `literalOf`, which is empty. - assertNotPushableAndNotConvertible(equality(ref("amount"), decimalLiteral("12.345"))); - } - - @Test - @DisplayName("An empty named reference is rejected by convert as well as by isPushable") - void emptyReferenceIsAlsoNotConvertible() { - // `isFieldRefExpr` on the convert path only checks `instanceof NamedReference`, so the - // zero-part guard lives in `columnOf`. Without it a pushable-looking predicate would reach - // the reader and be silently dropped. - assertNotPushableAndNotConvertible(equality(ref(), literal(1))); - assertNotPushableAndNotConvertible(predicate("IS_NULL", ref())); - assertNotPushableAndNotConvertible(predicate("IN", ref(), literal(1))); - assertNotPushableAndNotConvertible(predicate("STARTS_WITH", ref(), literal("a"))); - assertNotPushableAndNotConvertible(predicate("BOOLEAN_EXPRESSION", ref())); - } - - /** - * Asserts the invariant documented on {@link SparkPredicateToVortexExpression#isPushable(Predicate, Map)}: a - * predicate Spark is allowed to drop must produce a Vortex expression. - */ - private static void assertPushableAndConvertible(Predicate predicate) { - assertPushableAndConvertible(predicate, predicate.name()); - } - - private static void assertPushableAndConvertible(Predicate predicate, String what) { - assertTrue(SparkPredicateToVortexExpression.isPushable(predicate, SCHEMA), () -> "not pushable: " + what); - assertTrue( - SparkPredicateToVortexExpression.convert(predicate).isPresent(), - () -> "pushable but not convertible: " + what); - } - - private static void assertNotPushableAndNotConvertible(Predicate predicate) { - assertNotPushableAndNotConvertible(predicate, predicate.name()); - } - - private static void assertNotPushableAndNotConvertible(Predicate predicate, String what) { - assertFalse(SparkPredicateToVortexExpression.isPushable(predicate, SCHEMA), () -> "pushable: " + what); - assertFalse(SparkPredicateToVortexExpression.convert(predicate).isPresent(), () -> "convertible: " + what); - } - - private static Predicate predicate(String name, Expression... children) { - return new Predicate(name, children); - } - - private static LiteralValue decimalLiteral(String value) { - return new LiteralValue<>(Decimal.apply(new BigDecimal(value)), DECIMAL); - } - - private static Predicate equality(Expression left, Expression right) { - return new Predicate("=", new Expression[] {left, right}); - } - - private static NamedReference ref(String... parts) { - return new TestNamedReference(parts); - } - - private static LiteralValue literal(int value) { - return new LiteralValue<>(value, DataTypes.IntegerType); - } - - private static LiteralValue literal(String value) { - return new LiteralValue<>(org.apache.spark.unsafe.types.UTF8String.fromString(value), DataTypes.StringType); - } - - private static final class TestNamedReference implements NamedReference { - private final String[] fieldNames; - - TestNamedReference(String[] fieldNames) { - this.fieldNames = fieldNames; - } - - @Override - public String[] fieldNames() { - return fieldNames; - } - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexBatchExecTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexBatchExecTest.java deleted file mode 100644 index 967940567b2..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexBatchExecTest.java +++ /dev/null @@ -1,95 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.read; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertInstanceOf; - -import dev.vortex.spark.VortexFilePartition; -import java.util.List; -import java.util.Map; -import org.apache.spark.sql.connector.catalog.Column; -import org.apache.spark.sql.connector.expressions.filter.Predicate; -import org.apache.spark.sql.connector.read.InputPartition; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link VortexBatchExec} input-partition planning. - * - *

Characterizes how explicit {@code .vortex} paths are planned: one {@link VortexFilePartition} per file, each - * carrying the requested read schema and the Hive-style partition values parsed from its own path. All paths used here - * end in {@code .vortex} so planning never lists directories through native I/O. - */ -final class VortexBatchExecTest { - - private static final List COLUMNS = List.of( - Column.create("id", org.apache.spark.sql.types.DataTypes.IntegerType), - Column.create("name", org.apache.spark.sql.types.DataTypes.StringType)); - - private static VortexBatchExec execFor(List paths) { - return new VortexBatchExec(paths, COLUMNS, Map.of(), new Predicate[0]); - } - - @Test - @DisplayName("Plans one input partition per .vortex file") - void plansOnePartitionPerFile() { - VortexBatchExec exec = execFor(List.of("/data/a.vortex", "/data/b.vortex", "/data/c.vortex")); - - InputPartition[] partitions = exec.planInputPartitions(); - - assertEquals(3, partitions.length); - for (InputPartition partition : partitions) { - assertInstanceOf(VortexFilePartition.class, partition); - assertEquals(1, ((VortexFilePartition) partition).paths().size()); - } - } - - @Test - @DisplayName("Each partition carries exactly its own file path, in input order") - void partitionsKeepInputOrder() { - VortexBatchExec exec = execFor(List.of("/data/first.vortex", "/data/second.vortex")); - - InputPartition[] partitions = exec.planInputPartitions(); - - assertEquals(List.of("/data/first.vortex"), ((VortexFilePartition) partitions[0]).paths()); - assertEquals(List.of("/data/second.vortex"), ((VortexFilePartition) partitions[1]).paths()); - } - - @Test - @DisplayName("Partitions carry the requested read schema") - void partitionsCarryReadSchema() { - VortexBatchExec exec = execFor(List.of("/data/a.vortex")); - - VortexFilePartition partition = (VortexFilePartition) exec.planInputPartitions()[0]; - - assertEquals(2, partition.readSchema().size()); - assertEquals("id", partition.readSchema().fields()[0].name()); - assertEquals("name", partition.readSchema().fields()[1].name()); - } - - @Test - @DisplayName("Hive-style partition values are parsed from each file's own path") - void partitionValuesParsedPerFile() { - VortexBatchExec exec = execFor( - List.of("/tbl/year=2024/month=01/a.vortex", "/tbl/year=2025/month=02/b.vortex", "/tbl/plain.vortex")); - - InputPartition[] partitions = exec.planInputPartitions(); - - assertEquals(Map.of("year", "2024", "month", "01"), ((VortexFilePartition) partitions[0]).partitionValues()); - assertEquals(Map.of("year", "2025", "month", "02"), ((VortexFilePartition) partitions[1]).partitionValues()); - assertEquals(Map.of(), ((VortexFilePartition) partitions[2]).partitionValues()); - } - - @Test - @DisplayName("Format options are propagated to every partition") - void formatOptionsPropagated() { - Map options = Map.of("vortex.workerThreads", "8"); - VortexBatchExec exec = new VortexBatchExec(List.of("/data/a.vortex"), COLUMNS, options, new Predicate[0]); - - VortexFilePartition partition = (VortexFilePartition) exec.planInputPartitions()[0]; - - assertEquals("8", partition.formatOptions().get("vortex.workerThreads")); - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java deleted file mode 100644 index 9cc841b631b..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexScanBuilderTest.java +++ /dev/null @@ -1,176 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -package dev.vortex.spark.read; - -import static org.junit.jupiter.api.Assertions.assertArrayEquals; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNotSame; -import static org.junit.jupiter.api.Assertions.assertThrows; - -import java.util.Map; -import org.apache.spark.sql.connector.catalog.Column; -import org.apache.spark.sql.connector.expressions.Expression; -import org.apache.spark.sql.connector.expressions.Expressions; -import org.apache.spark.sql.connector.expressions.LiteralValue; -import org.apache.spark.sql.connector.expressions.NamedReference; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.connector.expressions.filter.Predicate; -import org.apache.spark.sql.types.DataTypes; -import org.apache.spark.unsafe.types.UTF8String; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link VortexScanBuilder}, focusing on how {@link VortexScanBuilder#pushPredicates(Predicate[])} - * splits predicates between Vortex pushdown and Spark post-scan evaluation. - * - *

Characterizes the partition-column exclusion rule (partition columns live in directory paths, not in the Vortex - * files, so predicates referencing them must stay with Spark), the handling of unsupported operators, and the defensive - * copy returned by {@link VortexScanBuilder#pushedPredicates()}. - */ -final class VortexScanBuilderTest { - - private static VortexScanBuilder builderWithColumns(Transform[] partitionTransforms, Column... columns) { - VortexScanBuilder builder = new VortexScanBuilder(Map.of(), partitionTransforms); - builder.addPath("/tmp/data.vortex"); - for (Column column : columns) { - builder.addColumn(column); - } - return builder; - } - - private static Predicate equality(String column, int value) { - return new Predicate("=", new Expression[] {ref(column), new LiteralValue<>(value, DataTypes.IntegerType)}); - } - - private static NamedReference ref(String name) { - return Expressions.column(name); - } - - // --- pushPredicates: data columns --- - - @Test - @DisplayName("Predicate on a data column is pushed and removed from post-scan predicates") - void dataColumnPredicateIsPushed() { - VortexScanBuilder builder = builderWithColumns( - new Transform[0], - Column.create("id", DataTypes.IntegerType), - Column.create("name", DataTypes.StringType)); - - Predicate predicate = equality("id", 42); - Predicate[] postScan = builder.pushPredicates(new Predicate[] {predicate}); - - assertEquals(0, postScan.length); - assertArrayEquals(new Predicate[] {predicate}, builder.pushedPredicates()); - } - - @Test - @DisplayName("Predicate on a column missing from the read schema is left to Spark") - void unknownColumnPredicateIsNotPushed() { - VortexScanBuilder builder = builderWithColumns(new Transform[0], Column.create("id", DataTypes.IntegerType)); - - Predicate predicate = equality("missing", 1); - Predicate[] postScan = builder.pushPredicates(new Predicate[] {predicate}); - - assertArrayEquals(new Predicate[] {predicate}, postScan); - assertEquals(0, builder.pushedPredicates().length); - } - - @Test - @DisplayName("Predicate with an unsupported operator is left to Spark") - void unsupportedOperatorIsNotPushed() { - VortexScanBuilder builder = builderWithColumns(new Transform[0], Column.create("name", DataTypes.StringType)); - - // LIKE with arbitrary user pattern is not a translatable V2 predicate name. - Predicate predicate = new Predicate( - "LIKE", - new Expression[] {ref("name"), new LiteralValue<>(UTF8String.fromString("%a%"), DataTypes.StringType)}); - Predicate[] postScan = builder.pushPredicates(new Predicate[] {predicate}); - - assertArrayEquals(new Predicate[] {predicate}, postScan); - assertEquals(0, builder.pushedPredicates().length); - } - - // --- pushPredicates: partition columns --- - - @Test - @DisplayName("Predicate on a partition column is left to Spark") - void partitionColumnPredicateIsNotPushed() { - Transform[] transforms = new Transform[] {Expressions.identity("year")}; - VortexScanBuilder builder = builderWithColumns( - transforms, Column.create("id", DataTypes.IntegerType), Column.create("year", DataTypes.IntegerType)); - - Predicate predicate = equality("year", 2024); - Predicate[] postScan = builder.pushPredicates(new Predicate[] {predicate}); - - assertArrayEquals(new Predicate[] {predicate}, postScan); - assertEquals(0, builder.pushedPredicates().length); - } - - @Test - @DisplayName("Mixed predicates split into pushed (data) and post-scan (partition)") - void mixedPredicatesAreSplit() { - Transform[] transforms = new Transform[] {Expressions.identity("year")}; - VortexScanBuilder builder = builderWithColumns( - transforms, Column.create("id", DataTypes.IntegerType), Column.create("year", DataTypes.IntegerType)); - - Predicate onData = equality("id", 7); - Predicate onPartition = equality("year", 2024); - Predicate[] postScan = builder.pushPredicates(new Predicate[] {onData, onPartition}); - - assertArrayEquals(new Predicate[] {onPartition}, postScan); - assertArrayEquals(new Predicate[] {onData}, builder.pushedPredicates()); - } - - @Test - @DisplayName("AND spanning a data column and a partition column is left to Spark as a whole") - void conjunctionSpanningPartitionColumnIsNotPushed() { - Transform[] transforms = new Transform[] {Expressions.identity("year")}; - VortexScanBuilder builder = builderWithColumns( - transforms, Column.create("id", DataTypes.IntegerType), Column.create("year", DataTypes.IntegerType)); - - Predicate and = - new org.apache.spark.sql.connector.expressions.filter.And(equality("id", 7), equality("year", 2024)); - Predicate[] postScan = builder.pushPredicates(new Predicate[] {and}); - - assertArrayEquals(new Predicate[] {and}, postScan); - assertEquals(0, builder.pushedPredicates().length); - } - - // --- pushedPredicates --- - - @Test - @DisplayName("pushedPredicates is empty before any pushPredicates call") - void pushedPredicatesEmptyByDefault() { - VortexScanBuilder builder = builderWithColumns(new Transform[0], Column.create("id", DataTypes.IntegerType)); - assertEquals(0, builder.pushedPredicates().length); - } - - @Test - @DisplayName("pushedPredicates returns a defensive copy") - void pushedPredicatesReturnsCopy() { - VortexScanBuilder builder = builderWithColumns(new Transform[0], Column.create("id", DataTypes.IntegerType)); - builder.pushPredicates(new Predicate[] {equality("id", 1)}); - - Predicate[] first = builder.pushedPredicates(); - Predicate[] second = builder.pushedPredicates(); - assertNotSame(first, second); - - // Mutating one returned array must not leak into subsequent calls. - first[0] = null; - assertEquals(1, builder.pushedPredicates().length); - assertNotNull(builder.pushedPredicates()[0]); - } - - // --- build --- - - @Test - @DisplayName("build without paths fails fast") - void buildWithoutPathsThrows() { - VortexScanBuilder builder = new VortexScanBuilder(Map.of()); - builder.addColumn(Column.create("id", DataTypes.IntegerType)); - assertThrows(IllegalStateException.class, builder::build); - } -} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java deleted file mode 100644 index 7d3318d61b9..00000000000 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java +++ /dev/null @@ -1,169 +0,0 @@ -// 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.write.PartitionedVortexDataWriter.ResolvedTransform; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.ObjectInputStream; -import java.io.ObjectOutputStream; -import java.util.List; -import org.apache.spark.sql.connector.expressions.Expressions; -import org.apache.spark.sql.connector.expressions.Transform; -import org.apache.spark.sql.types.DataTypes; -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.DisplayName; -import org.junit.jupiter.api.Test; - -/** - * Unit tests for {@link PartitionedVortexDataWriter#resolveTransforms}, which turns Spark partition transforms into the - * directory keys and column indices a writer uses. - * - *

Resolution happens eagerly on the driver because Spark's {@code Transform} objects are Scala case classes that are - * not Java-serializable, so only the resolved form may cross to the executors; the serialization round trip below is - * what that eager step exists for. - */ -final class PartitionedVortexDataWriterTest { - - private static final StructType SCHEMA = new StructType(new StructField[] { - new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), - new StructField("name", DataTypes.StringType, true, Metadata.empty()), - new StructField("event_date", DataTypes.DateType, true, Metadata.empty()), - new StructField("event_ts", DataTypes.TimestampType, true, Metadata.empty()) - }); - - private static ResolvedTransform resolveOne(Transform transform) { - ResolvedTransform[] resolved = - PartitionedVortexDataWriter.resolveTransforms(new Transform[] {transform}, SCHEMA); - assertEquals(1, resolved.length); - return resolved[0]; - } - - @Test - @DisplayName("No transforms resolve to no partitioning") - void emptyInput() { - assertEquals(0, PartitionedVortexDataWriter.resolveTransforms(new Transform[0], SCHEMA).length); - } - - @Test - @DisplayName("An identity transform keeps the column name as its directory key") - void identity() { - ResolvedTransform resolved = resolveOne(Expressions.identity("name")); - - assertEquals("name", resolved.directoryKey()); - assertEquals("identity", resolved.transformName()); - assertEquals(1, resolved.columnIndices().get(0)); - assertEquals(DataTypes.StringType, resolved.columnTypes().get(0)); - assertEquals(-1, resolved.bucketCount(), "only bucket transforms carry a bucket count"); - } - - @Test - @DisplayName("Temporal transforms suffix the directory key with the unit they truncate to") - void temporalDirectoryKeys() { - assertEquals( - "event_date_year", resolveOne(Expressions.years("event_date")).directoryKey()); - assertEquals( - "event_date_month", resolveOne(Expressions.months("event_date")).directoryKey()); - assertEquals( - "event_date_day", resolveOne(Expressions.days("event_date")).directoryKey()); - assertEquals("event_ts_hour", resolveOne(Expressions.hours("event_ts")).directoryKey()); - } - - @Test - @DisplayName("Temporal transforms reject a non-temporal column, naming the transform and the type") - void temporalTransformsRejectNonTemporalColumns() { - IllegalArgumentException thrown = - assertThrows(IllegalArgumentException.class, () -> resolveOne(Expressions.years("id"))); - - assertTrue(thrown.getMessage().contains("years"), thrown.getMessage()); - assertTrue(thrown.getMessage().contains("IntegerType"), thrown.getMessage()); - } - - @Test - @DisplayName("The hours transform additionally rejects a date column, which has no hour to truncate to") - void hoursRejectsDateColumn() { - IllegalArgumentException thrown = - assertThrows(IllegalArgumentException.class, () -> resolveOne(Expressions.hours("event_date"))); - - assertTrue(thrown.getMessage().contains("hours"), thrown.getMessage()); - } - - @Test - @DisplayName("A bucket transform carries its bucket count and joins its columns with underscores") - void bucket() { - ResolvedTransform single = resolveOne(Expressions.bucket(8, "id")); - assertEquals("id_bucket", single.directoryKey()); - assertEquals(8, single.bucketCount()); - - ResolvedTransform multi = resolveOne(Expressions.bucket(4, "id", "name")); - assertEquals("id_name_bucket", multi.directoryKey()); - assertEquals(4, multi.bucketCount()); - assertEquals(0, multi.columnIndices().get(0)); - assertEquals(1, multi.columnIndices().get(1)); - assertEquals(List.of(DataTypes.IntegerType, DataTypes.StringType), multi.columnTypes()); - } - - @Test - @DisplayName("A bucket transform without a numBuckets argument is rejected") - void bucketWithoutCount() { - Transform noCount = Expressions.apply("bucket", Expressions.column("id")); - - IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> resolveOne(noCount)); - - assertTrue(thrown.getMessage().contains("numBuckets"), thrown.getMessage()); - } - - @Test - @DisplayName("An unsupported transform name is rejected") - void unsupportedTransform() { - IllegalArgumentException thrown = assertThrows( - IllegalArgumentException.class, - () -> resolveOne(Expressions.apply("truncate", Expressions.column("name")))); - - assertTrue(thrown.getMessage().contains("truncate"), thrown.getMessage()); - } - - @Test - @DisplayName("A transform that references no column is rejected") - void transformWithoutReferences() { - IllegalArgumentException thrown = - assertThrows(IllegalArgumentException.class, () -> resolveOne(Expressions.apply("identity"))); - - assertTrue(thrown.getMessage().contains("no column references"), thrown.getMessage()); - } - - @Test - @DisplayName("A transform on a column outside the schema is rejected") - void transformOnUnknownColumn() { - assertThrows(IllegalArgumentException.class, () -> resolveOne(Expressions.identity("missing"))); - } - - @Test - @DisplayName("Resolved transforms survive Java serialization, which is why they are resolved eagerly") - void resolvedTransformsAreSerializable() throws IOException, ClassNotFoundException { - ResolvedTransform original = resolveOne(Expressions.bucket(4, "id", "name")); - - ByteArrayOutputStream bytes = new ByteArrayOutputStream(); - try (ObjectOutputStream out = new ObjectOutputStream(bytes)) { - out.writeObject(original); - } - ResolvedTransform restored; - try (ObjectInputStream in = new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) { - restored = (ResolvedTransform) in.readObject(); - } - - assertEquals(original.directoryKey(), restored.directoryKey()); - assertEquals(original.transformName(), restored.transformName()); - assertEquals(original.bucketCount(), restored.bucketCount()); - assertEquals(original.columnTypes(), restored.columnTypes()); - assertEquals(original.columnIndices().asList(), restored.columnIndices().asList()); - } -} diff --git a/java/vortex-spark/v3.5/build.gradle.kts b/java/vortex-spark/v3.5/build.gradle.kts new file mode 100644 index 00000000000..3afde2ce9d5 --- /dev/null +++ b/java/vortex-spark/v3.5/build.gradle.kts @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +plugins { + id("vortex-spark-module") +} diff --git a/java/vortex-spark/v4.0/build.gradle.kts b/java/vortex-spark/v4.0/build.gradle.kts new file mode 100644 index 00000000000..3afde2ce9d5 --- /dev/null +++ b/java/vortex-spark/v4.0/build.gradle.kts @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +plugins { + id("vortex-spark-module") +}