From a70e7a01894273443f1435d66d115741c96a8e89 Mon Sep 17 00:00:00 2001 From: jackylee Date: Mon, 10 Aug 2026 17:48:49 +0800 Subject: [PATCH] test(spark): characterize partition transform resolution resolveTransforms had no tests. Pin the directory keys it derives, the type guards on the temporal transforms, the bucket count and multi-column resolution, and the rejections for an unknown transform name, a missing numBuckets argument and a column outside the schema. Also cover the Java serialization round trip of ResolvedTransform, which is the reason the resolution happens eagerly on the driver: Spark's Transform objects are Scala case classes that cannot cross to executors. Signed-off-by: jackylee --- .../PartitionedVortexDataWriterTest.java | 169 ++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java 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 new file mode 100644 index 00000000000..7d3318d61b9 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/write/PartitionedVortexDataWriterTest.java @@ -0,0 +1,169 @@ +// 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()); + } +}