Skip to content

Write struct support #15

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stevelordbq wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Write struct support #15

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e8b62c4
Write struct initial version
stevelordbq Feb 13, 2026
69765bf
Write struct add unit test for struct array
stevelordbq Feb 17, 2026
d15abed
Write struct extend unit test for struct array
stevelordbq Feb 17, 2026
547383f
Spanner Write - assume Spark stores strings at UTF8String instead of …
stevelordbq Feb 18, 2026
0e7d5b8
Spanner Write - add types to struct test
stevelordbq Feb 18, 2026
e893f55
Spanner Write - PR comments
stevelordbq Feb 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 89 additions & 37 deletions spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriterUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,16 @@
import com.google.cloud.Date;
import com.google.cloud.Timestamp;
import com.google.cloud.spanner.Mutation;
import com.google.cloud.spanner.Struct;
import com.google.cloud.spanner.Type;
import com.google.cloud.spanner.Value;
import java.math.BigDecimal;
import java.time.LocalDate;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.catalyst.util.ArrayData;
Expand Down Expand Up @@ -55,7 +58,8 @@ private interface FieldConverter {
// String
TYPE_CONVERTERS.put(
DataTypes.StringType,
(row, i, type) -> row.isNullAt(i) ? Value.string(null) : Value.string(row.getString(i)));
(row, i, type) ->
row.isNullAt(i) ? Value.string(null) : Value.string(row.getUTF8String(i).toString()));

// Boolean
TYPE_CONVERTERS.put(
Expand Down Expand Up @@ -98,65 +102,55 @@ private interface FieldConverter {
// Long array
ARRAY_TYPE_CONVERTERS.put(
DataTypes.LongType,
(row, i, type) -> {
return arrayConverter(row, i, ArrayData::getLong, Value::int64Array);
});
(row, i, type) -> arrayConverter(row, i, ArrayData::getLong, Value::int64Array));

// String array
ARRAY_TYPE_CONVERTERS.put(
DataTypes.StringType,
(row, i, type) -> {
return arrayConverter(
row, i, (a, idx) -> a.getUTF8String(idx).toString(), Value::stringArray);
});
(row, i, type) ->
arrayConverter(
row, i, (a, idx) -> a.getUTF8String(idx).toString(), Value::stringArray));

// Boolean
ARRAY_TYPE_CONVERTERS.put(
DataTypes.BooleanType,
(row, i, type) -> {
return arrayConverter(row, i, ArrayData::getBoolean, Value::boolArray);
});
(row, i, type) -> arrayConverter(row, i, ArrayData::getBoolean, Value::boolArray));

// Double
ARRAY_TYPE_CONVERTERS.put(
DataTypes.DoubleType,
(row, i, type) -> {
return arrayConverter(row, i, ArrayData::getDouble, Value::float64Array);
});
(row, i, type) -> arrayConverter(row, i, ArrayData::getDouble, Value::float64Array));

// Binary
ARRAY_TYPE_CONVERTERS.put(
DataTypes.BinaryType,
(row, i, type) -> {
return arrayConverter(
row, i, (a, idx) -> ByteArray.copyFrom(a.getBinary(idx)), Value::bytesArray);
});
(row, i, type) ->
arrayConverter(
row, i, (a, idx) -> ByteArray.copyFrom(a.getBinary(idx)), Value::bytesArray));

// Timestamp
ARRAY_TYPE_CONVERTERS.put(
DataTypes.TimestampType,
(row, i, type) -> {
return arrayConverter(
row,
i,
(a, idx) -> Timestamp.ofTimeMicroseconds(a.getLong(idx)),
Value::timestampArray);
});
(row, i, type) ->
arrayConverter(
row,
i,
(a, idx) -> Timestamp.ofTimeMicroseconds(a.getLong(idx)),
Value::timestampArray));

// Date
ARRAY_TYPE_CONVERTERS.put(
DataTypes.DateType,
(row, i, type) -> {
return arrayConverter(
row,
i,
(a, idx) -> {
LocalDate localDate = LocalDate.ofEpochDay(a.getInt(idx));
return Date.fromYearMonthDay(
localDate.getYear(), localDate.getMonthValue(), localDate.getDayOfMonth());
},
Value::dateArray);
});
(row, i, type) ->
arrayConverter(
row,
i,
(a, idx) -> {
LocalDate localDate = LocalDate.ofEpochDay(a.getInt(idx));
return Date.fromYearMonthDay(
localDate.getYear(), localDate.getMonthValue(), localDate.getDayOfMonth());
},
Value::dateArray));

// Note: DecimalType is handled dynamically in the method below
// because it relies on instanceof checks rather than strict equality.
Expand Down Expand Up @@ -211,6 +205,21 @@ public static Mutation internalRowToMutation(
return builder.build();
}

private static Struct internalRowToStruct(InternalRow record, StructType schema) {
final Struct.Builder builder = Struct.newBuilder();

for (int i = 0; i < schema.length(); i++) {
final StructField field = schema.fields()[i];
final DataType fieldType = field.dataType();
final String fieldName = field.name();

Value spannerValue = convertField(record, i, fieldType);
builder.set(fieldName).to(spannerValue);
}

return builder.build();
}

// Helper method to resolve the strategy
private static Value convertField(InternalRow row, int index, DataType type) {
// Check exact matches
Expand All @@ -229,6 +238,17 @@ private static Value convertField(InternalRow row, int index, DataType type) {
return Value.numeric(bd);
}

if (type instanceof StructType) {
// A Struct instance in Spanner always represents a non-NULL value.
if (row.isNullAt(index)) {
return Value.struct(Struct.newBuilder().build()); // TODO: How should this be handled.
}

StructType structType = (StructType) type;
return Value.struct(
internalRowToStruct(row.getStruct(index, structType.length()), structType));
}

if (type instanceof ArrayType) {
ArrayType arrayType = (ArrayType) type;
DataType elementType = arrayType.elementType();
Expand All @@ -245,8 +265,40 @@ private static Value convertField(InternalRow row, int index, DataType type) {
(ad, i) -> ad.getDecimal(i, dt.precision(), dt.scale()).toJavaBigDecimal(),
Value::numericArray);
}

if (elementType instanceof StructType) {
StructType structType = (StructType) elementType;

if (row.isNullAt(index)) {
return Value.structArray(Type.struct(), null);
}

final ArrayData arrayData = row.getArray(index);
final int numElements = arrayData.numElements();
if (numElements == 0) return Value.structArray(Type.struct(), new ArrayList<>());
List<Struct> convertedList = new ArrayList<>(numElements);

for (int j = 0; j < numElements; j++) {
if (arrayData.isNullAt(j)) {
convertedList.add(null);
} else {
// Uses specialized Spark getters (getLong, getDouble, etc.)
convertedList.add(
internalRowToStruct(arrayData.getStruct(j, structType.length()), structType));
}
}

// Find Spanner schema from first non-null struct array element.
Type itemType =
convertedList.stream()
.filter(Objects::nonNull)
.findFirst()
.map(Struct::getType)
.orElse(Type.struct());
return Value.structArray(itemType, convertedList);
}
}
// TODO handle Struct here.

// unsupported type
throw new UnsupportedOperationException("Unsupported Spark DataType: " + type);
}
Expand Down
91 changes: 86 additions & 5 deletions ...-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerWriterUtilsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
import com.google.cloud.Date;
import com.google.cloud.Timestamp;
import com.google.cloud.spanner.Mutation;
import com.google.cloud.spanner.Struct;
import com.google.cloud.spanner.Type;
import com.google.cloud.spanner.Value;
import java.math.MathContext;
import java.math.RoundingMode;
Expand All @@ -29,6 +31,7 @@
import java.util.Collections;
import java.util.function.BiConsumer;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
import org.apache.spark.sql.catalyst.util.ArrayData;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
Expand All @@ -52,18 +55,69 @@ public class SpannerWriterUtilsTest {
new java.math.BigDecimal("135.791", mc).setScale(9, RoundingMode.HALF_UP);
private static final scala.math.BigDecimal bd = scala.math.BigDecimal$.MODULE$.apply(jbd);
private static final Decimal decimal = new Decimal().set(bd);
private static final byte[] BYTE_DATA = {95, -10, 127};
private static final StructType structType =
new StructType()
.add("long_field", DataTypes.LongType)
.add("str_field", DataTypes.StringType)
.add("bool_field", DataTypes.BooleanType)
.add("double_field", DataTypes.DoubleType)
.add("binary_field", DataTypes.BinaryType)
.add("ts_field", DataTypes.TimestampType)
.add("dt_field", DataTypes.DateType)
.add("decimal_field", new DecimalType(38, 9));
private static final Type typeStruct =
Type.struct(
Type.StructField.of("long_field", Type.int64()),
Type.StructField.of("str_field", Type.string()),
Type.StructField.of("bool_field", Type.bool()),
Type.StructField.of("double_field", Type.float64()),
Type.StructField.of("binary_field", Type.bytes()),
Type.StructField.of("ts_field", Type.timestamp()),
Type.StructField.of("dt_field", Type.date()),
Type.StructField.of("decimal_field", Type.numeric()));
private static final Struct testStruct =
Struct.newBuilder()
.set("long_field")
.to(100L)
.set("str_field")
.to("str_value")
.set("bool_field")
.to(true)
.set("double_field")
.to(95.5)
.set("binary_field")
.to(ByteArray.copyFrom(BYTE_DATA))
.set("ts_field")
.to(Timestamp.ofTimeMicroseconds(1704067200000000L))
.set("dt_field")
.to(Date.fromYearMonthDay(2024, 1, 1))
.set("decimal_field")
.to(jbd)
.build();
private static final Object[] structValues =
new Object[] {
100L,
UTF8String.fromString("str_value"),
true,
95.5,
BYTE_DATA,
1704067200000000L,
19723,
decimal
};
private static final InternalRow testInternalRow = new GenericInternalRow(structValues);

@RunWith(Parameterized.class)
public static class ScalarTests {
private static final byte[] BYTE_DATA = {95, -10, 127};

// Parameters for each test case: [ColumnName, DataType, MockValue, ExpectedSpannerValue]
@Parameters(name = "Testing {0}")
public static Collection<Object[]> data() {
return Arrays.asList(
new Object[][] {
{"long", DataTypes.LongType, 100L, Value.int64(100L)},
{"string", DataTypes.StringType, "Hello", Value.string("Hello")},
{"string", DataTypes.StringType, UTF8String.fromString("Hello"), Value.string("Hello")},
{"boolean", DataTypes.BooleanType, true, Value.bool(true)},
{"double", DataTypes.DoubleType, 95.5, Value.float64(95.5)},
{"binary", DataTypes.BinaryType, BYTE_DATA, Value.bytes(ByteArray.copyFrom(BYTE_DATA))},
Expand All @@ -74,7 +128,8 @@ public static Collection<Object[]> data() {
Value.timestamp(Timestamp.ofTimeMicroseconds(1704067200000000L))
},
{"dt", DataTypes.DateType, 19723, Value.date(Date.fromYearMonthDay(2024, 1, 1))},
{"decimal", new DecimalType(38, 9), decimal, Value.numeric(jbd)}
{"decimal", new DecimalType(38, 9), decimal, Value.numeric(jbd)},
{"struct", structType, testInternalRow, Value.struct(typeStruct, testStruct)}
});
}

Expand Down Expand Up @@ -102,7 +157,7 @@ public void testScalarConversion() {
// Setup specific getter based on type
if (sparkType == DataTypes.LongType) when(row.getLong(0)).thenReturn((Long) mockValue);
else if (sparkType == DataTypes.StringType)
when(row.getString(0)).thenReturn((String) mockValue);
when(row.getUTF8String(0)).thenReturn((UTF8String) mockValue);
else if (sparkType == DataTypes.BooleanType)
when(row.getBoolean(0)).thenReturn((Boolean) mockValue);
else if (sparkType == DataTypes.DoubleType)
Expand All @@ -114,6 +169,9 @@ else if (sparkType == DataTypes.TimestampType)
else if (sparkType == DataTypes.DateType) when(row.getInt(0)).thenReturn((int) mockValue);
else if (sparkType instanceof DecimalType)
when(row.getDecimal(0, 38, 9)).thenReturn((Decimal) mockValue);
else if (sparkType instanceof StructType)
when(row.getStruct(0, ((StructType) sparkType).length()))
.thenReturn((InternalRow) mockValue);

// 3. Execute
com.google.cloud.spanner.Mutation mutation =
Expand Down Expand Up @@ -143,6 +201,7 @@ public static Collection<Object[]> data() {
{"timestamp", DataTypes.TimestampType, Value.timestamp(null)},
{"date", DataTypes.DateType, Value.date(null)},
{"decimal", new DecimalType(38, 9), Value.numeric(null)},
{"struct", new StructType(), Value.struct(Struct.newBuilder().build())}, // {
{
"long_array",
DataTypes.createArrayType(DataTypes.LongType),
Expand Down Expand Up @@ -174,6 +233,11 @@ public static Collection<Object[]> data() {
"decimal_array",
DataTypes.createArrayType(new DecimalType(38, 9)),
Value.numericArray(null)
},
{
"struct_array",
DataTypes.createArrayType(new StructType()),
Value.structArray(Type.struct(), null)
}
});
}
Expand Down Expand Up @@ -341,6 +405,23 @@ public static Collection<Object[]> data() {
ad,
(Decimal[]) d,
(i, val) -> when(ad.getDecimal(i, 38, 9)).thenReturn((Decimal) val))
},

// 9. Struct Array
{
"struct_array",
structType,
new InternalRow[] {testInternalRow}, // Use InternalRow instead of Spanner Struct
Value.structArray(typeStruct, Collections.singletonList(testStruct)),
(BiConsumer<ArrayData, Object>)
(ad, d) ->
setupArrayMock(
ad,
(InternalRow[]) d, // Cast to InternalRow array
(i, val) -> {
when(ad.getStruct(i, structType.length()))
.thenReturn((InternalRow) val);
})
}
});
}
Expand Down Expand Up @@ -368,7 +449,7 @@ public void testArrayConversion() {
when(row.isNullAt(0)).thenReturn(false);
when(row.getArray(0)).thenReturn(arrayData);

// Executes the specific stubbing (e.g., ad.toLongArray() or ad.toIntArray())
// Executes the specific stubbing (e.g., ad.toLongArray(), ad.toIntArray(), ad.to)
mockSetup.accept(arrayData, inputData);

Mutation mutation = SpannerWriterUtils.internalRowToMutation(TABLE_NAME, row, schema);
Expand Down