Skip to content

Commit fb19e1d

Browse files
committed
test(spark): characterize SparkToArrowSchema Spark-to-Arrow conversion
Signed-off-by: jackylee <qcsd2011@gmail.com>
1 parent 159207e commit fb19e1d

1 file changed

Lines changed: 176 additions & 0 deletions

File tree

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
package dev.vortex.spark.write;
5+
6+
import static org.junit.jupiter.api.Assertions.assertEquals;
7+
import static org.junit.jupiter.api.Assertions.assertFalse;
8+
import static org.junit.jupiter.api.Assertions.assertThrows;
9+
import static org.junit.jupiter.api.Assertions.assertTrue;
10+
11+
import dev.vortex.relocated.org.apache.arrow.vector.types.DateUnit;
12+
import dev.vortex.relocated.org.apache.arrow.vector.types.FloatingPointPrecision;
13+
import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit;
14+
import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType;
15+
import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Field;
16+
import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Schema;
17+
import java.util.List;
18+
import org.apache.spark.sql.types.DataTypes;
19+
import org.apache.spark.sql.types.Metadata;
20+
import org.apache.spark.sql.types.StructField;
21+
import org.apache.spark.sql.types.StructType;
22+
import org.junit.jupiter.api.DisplayName;
23+
import org.junit.jupiter.api.Test;
24+
25+
/**
26+
* Unit tests for {@link SparkToArrowSchema}, which converts Spark SQL schemas to Arrow schemas on the write path.
27+
*
28+
* <p>Characterizes the supported conversions (including field-level nullability and nested types) and the Spark types
29+
* that are explicitly rejected, mirroring the read-path coverage in {@code ArrowUtilsTest}.
30+
*/
31+
final class SparkToArrowSchemaTest {
32+
private static Field single(String name, org.apache.spark.sql.types.DataType type, boolean nullable) {
33+
StructType schema = new StructType(new StructField[] {new StructField(name, type, nullable, Metadata.empty())});
34+
Schema arrowSchema = SparkToArrowSchema.convert(schema);
35+
assertEquals(1, arrowSchema.getFields().size());
36+
return arrowSchema.getFields().get(0);
37+
}
38+
39+
@Test
40+
@DisplayName("BooleanType converts to Arrow Bool")
41+
void booleanConvertsToBool() {
42+
assertEquals(
43+
new ArrowType.Bool(), single("b", DataTypes.BooleanType, true).getType());
44+
}
45+
46+
@Test
47+
@DisplayName("Integral types convert to signed Arrow Ints of matching width")
48+
void integralTypesConvertBySignedWidth() {
49+
assertEquals(
50+
new ArrowType.Int(8, true),
51+
single("v", DataTypes.ByteType, true).getType());
52+
assertEquals(
53+
new ArrowType.Int(16, true),
54+
single("v", DataTypes.ShortType, true).getType());
55+
assertEquals(
56+
new ArrowType.Int(32, true),
57+
single("v", DataTypes.IntegerType, true).getType());
58+
assertEquals(
59+
new ArrowType.Int(64, true),
60+
single("v", DataTypes.LongType, true).getType());
61+
}
62+
63+
@Test
64+
@DisplayName("Float/Double convert to single/double precision Arrow FloatingPoint")
65+
void floatingPointConvertsByPrecision() {
66+
assertEquals(
67+
new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE),
68+
single("v", DataTypes.FloatType, true).getType());
69+
assertEquals(
70+
new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE),
71+
single("v", DataTypes.DoubleType, true).getType());
72+
}
73+
74+
@Test
75+
@DisplayName("String and Binary convert to Utf8 and Binary")
76+
void stringAndBinaryConvert() {
77+
assertEquals(
78+
new ArrowType.Utf8(), single("v", DataTypes.StringType, true).getType());
79+
assertEquals(
80+
new ArrowType.Binary(), single("v", DataTypes.BinaryType, true).getType());
81+
}
82+
83+
@Test
84+
@DisplayName("DecimalType preserves precision and scale as 128-bit Arrow Decimal")
85+
void decimalPreservesPrecisionAndScale() {
86+
assertEquals(
87+
new ArrowType.Decimal(20, 4, 128),
88+
single("v", DataTypes.createDecimalType(20, 4), true).getType());
89+
}
90+
91+
@Test
92+
@DisplayName("DateType converts to Date(DAY)")
93+
void dateConvertsToDayUnit() {
94+
assertEquals(
95+
new ArrowType.Date(DateUnit.DAY),
96+
single("v", DataTypes.DateType, true).getType());
97+
}
98+
99+
@Test
100+
@DisplayName("Timestamp converts to microsecond precision, with UTC tz only for the tz-aware type")
101+
void timestampConvertsByTimezoneAwareness() {
102+
assertEquals(
103+
new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"),
104+
single("v", DataTypes.TimestampType, true).getType());
105+
assertEquals(
106+
new ArrowType.Timestamp(TimeUnit.MICROSECOND, null),
107+
single("v", DataTypes.TimestampNTZType, true).getType());
108+
}
109+
110+
@Test
111+
@DisplayName("Field nullability is carried over to the Arrow field")
112+
void nullabilityIsPreserved() {
113+
assertTrue(single("v", DataTypes.IntegerType, true).isNullable());
114+
assertFalse(single("v", DataTypes.IntegerType, false).isNullable());
115+
}
116+
117+
@Test
118+
@DisplayName("Field names and ordering are preserved")
119+
void fieldNamesAndOrderArePreserved() {
120+
StructType schema = new StructType()
121+
.add("first", DataTypes.IntegerType)
122+
.add("second", DataTypes.StringType)
123+
.add("third", DataTypes.DoubleType);
124+
125+
Schema arrowSchema = SparkToArrowSchema.convert(schema);
126+
assertEquals(
127+
List.of("first", "second", "third"),
128+
arrowSchema.getFields().stream().map(Field::getName).toList());
129+
}
130+
131+
@Test
132+
@DisplayName("Nested StructType converts to a Struct field with converted children")
133+
void structConvertsWithChildren() {
134+
StructType inner =
135+
new StructType().add("a", DataTypes.IntegerType, true).add("b", DataTypes.StringType, false);
136+
Field structField = single("s", inner, true);
137+
138+
assertEquals(new ArrowType.Struct(), structField.getType());
139+
assertEquals(2, structField.getChildren().size());
140+
141+
Field a = structField.getChildren().get(0);
142+
assertEquals("a", a.getName());
143+
assertEquals(new ArrowType.Int(32, true), a.getType());
144+
assertTrue(a.isNullable());
145+
146+
Field b = structField.getChildren().get(1);
147+
assertEquals("b", b.getName());
148+
assertEquals(new ArrowType.Utf8(), b.getType());
149+
assertFalse(b.isNullable());
150+
}
151+
152+
@Test
153+
@DisplayName("ArrayType converts to a List field whose element carries containsNull")
154+
void arrayConvertsWithElementNullability() {
155+
Field withNulls = single("l", DataTypes.createArrayType(DataTypes.IntegerType, true), true);
156+
assertEquals(new ArrowType.List(), withNulls.getType());
157+
assertEquals(1, withNulls.getChildren().size());
158+
159+
Field element = withNulls.getChildren().get(0);
160+
assertEquals("element", element.getName());
161+
assertEquals(new ArrowType.Int(32, true), element.getType());
162+
assertTrue(element.isNullable());
163+
164+
Field withoutNulls = single("l", DataTypes.createArrayType(DataTypes.StringType, false), true);
165+
assertFalse(withoutNulls.getChildren().get(0).isNullable());
166+
}
167+
168+
@Test
169+
@DisplayName("Unsupported Spark types raise UnsupportedOperationException naming the type")
170+
void unsupportedTypeIsRejected() {
171+
StructType schema = new StructType().add("v", DataTypes.CalendarIntervalType);
172+
UnsupportedOperationException e =
173+
assertThrows(UnsupportedOperationException.class, () -> SparkToArrowSchema.convert(schema));
174+
assertTrue(e.getMessage().contains("CalendarIntervalType"));
175+
}
176+
}

0 commit comments

Comments
 (0)