Skip to content

Commit 203f45c

Browse files
claudevdmclaudevdm
authored
Support beam:logical_type:micros_instant:v1 in SpannerIo. (#36840)
* Support micros isntant in spannerio. * Trigger tests. * Comments. * Fix test.
1 parent 4591122 commit 203f45c

File tree

6 files changed

+261
-38
lines changed

6 files changed

+261
-38
lines changed

‎.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
{
22
"comment": "Modify this file in a trivial way to cause this test suite to run",
3-
"modification": 1
3+
"modification": 14
44
}

‎sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationUtils.java‎

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,13 @@
2828
import com.google.cloud.spanner.Mutation;
2929
import com.google.cloud.spanner.Value;
3030
import java.math.BigDecimal;
31+
import java.time.Instant;
3132
import java.util.HashMap;
3233
import java.util.List;
3334
import java.util.Map;
3435
import java.util.stream.StreamSupport;
3536
import org.apache.beam.sdk.schemas.Schema;
37+
import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
3638
import org.apache.beam.sdk.transforms.SerializableFunction;
3739
import org.apache.beam.sdk.values.Row;
3840
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
@@ -102,6 +104,11 @@ public static Mutation createMutationFromBeamRows(
102104
return mutationBuilder.build();
103105
}
104106

107+
private static Timestamp toSpannerTimestamp(Instant instant) {
108+
long micros = instant.getEpochSecond() * 1_000_000L + instant.getNano() / 1_000L;
109+
return Timestamp.ofTimeMicroseconds(micros);
110+
}
111+
105112
private static void setBeamValueToKey(
106113
Key.Builder keyBuilder, Schema.FieldType field, String columnName, Row row) {
107114
switch (field.getTypeName()) {
@@ -147,6 +154,21 @@ private static void setBeamValueToKey(
147154
keyBuilder.append(row.getDecimal(columnName));
148155
break;
149156
// TODO: Implement logical date and datetime
157+
case LOGICAL_TYPE:
158+
Schema.LogicalType<?, ?> logicalType = checkNotNull(field.getLogicalType());
159+
String identifier = logicalType.getIdentifier();
160+
if (identifier.equals(MicrosInstant.IDENTIFIER)) {
161+
Instant instant = row.getValue(columnName);
162+
if (instant == null) {
163+
keyBuilder.append((Timestamp) null);
164+
} else {
165+
keyBuilder.append(toSpannerTimestamp(instant));
166+
}
167+
} else {
168+
throw new IllegalArgumentException(
169+
String.format("Unsupported logical type in key: %s", identifier));
170+
}
171+
break;
150172
case DATETIME:
151173
@Nullable ReadableDateTime dateTime = row.getDateTime(columnName);
152174
if (dateTime == null) {
@@ -224,7 +246,21 @@ private static void setBeamValueToMutation(
224246
mutationBuilder.set(columnName).to(decimal);
225247
}
226248
break;
227-
// TODO: Implement logical date and datetime
249+
case LOGICAL_TYPE:
250+
Schema.LogicalType<?, ?> logicalType = checkNotNull(fieldType.getLogicalType());
251+
String identifier = logicalType.getIdentifier();
252+
if (identifier.equals(MicrosInstant.IDENTIFIER)) {
253+
@Nullable Instant instant = row.getValue(columnName);
254+
if (instant == null) {
255+
mutationBuilder.set(columnName).to((Timestamp) null);
256+
} else {
257+
mutationBuilder.set(columnName).to(toSpannerTimestamp(instant));
258+
}
259+
} else {
260+
throw new IllegalArgumentException(
261+
String.format("Unsupported logical type: %s", identifier));
262+
}
263+
break;
228264
case DATETIME:
229265
@Nullable ReadableDateTime dateTime = row.getDateTime(columnName);
230266
if (dateTime == null) {
@@ -335,6 +371,27 @@ private static void addIterableToMutationBuilder(
335371
case STRING:
336372
mutationBuilder.set(column).toStringArray((Iterable<String>) ((Object) iterable));
337373
break;
374+
case LOGICAL_TYPE:
375+
String identifier = checkNotNull(beamIterableType.getLogicalType()).getIdentifier();
376+
if (identifier.equals(MicrosInstant.IDENTIFIER)) {
377+
if (iterable == null) {
378+
mutationBuilder.set(column).toTimestampArray(null);
379+
} else {
380+
mutationBuilder
381+
.set(column)
382+
.toTimestampArray(
383+
StreamSupport.stream(iterable.spliterator(), false)
384+
.map(
385+
instant -> {
386+
return toSpannerTimestamp((java.time.Instant) instant);
387+
})
388+
.collect(toList()));
389+
}
390+
} else {
391+
throw new IllegalArgumentException(
392+
String.format("Unsupported logical type in iterable: %s", identifier));
393+
}
394+
break;
338395
case DATETIME:
339396
if (iterable == null) {
340397
mutationBuilder.set(column).toDateArray(null);

‎sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/StructUtils.java‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import java.util.Map;
3232
import java.util.stream.StreamSupport;
3333
import org.apache.beam.sdk.schemas.Schema;
34+
import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
3435
import org.apache.beam.sdk.values.Row;
3536
import org.checkerframework.checker.nullness.qual.Nullable;
3637
import org.joda.time.DateTime;
@@ -352,6 +353,11 @@ private static void addIterableToStructBuilder(
352353
}
353354
}
354355

356+
private static java.time.Instant fromSpannerTimestamp(Timestamp spannerTimestamp) {
357+
long micros = spannerTimestamp.getSeconds() * 1_000_000L + spannerTimestamp.getNanos() / 1_000L;
358+
return java.time.Instant.ofEpochSecond(micros / 1_000_000L, (micros % 1_000_000L) * 1_000L);
359+
}
360+
355361
private static @Nullable Object getStructValue(Struct struct, Schema.Field field) {
356362
String column = field.getName();
357363
Type.Code typeCode = struct.getColumnType(column).getCode();
@@ -365,7 +371,19 @@ private static void addIterableToStructBuilder(
365371
return struct.getBytes(column).toByteArray();
366372
// TODO: implement logical datetime
367373
case TIMESTAMP:
368-
return Instant.ofEpochSecond(struct.getTimestamp(column).getSeconds()).toDateTime();
374+
Timestamp spannerTimestamp = struct.getTimestamp(column);
375+
376+
// Check if the Beam schema expects MicrosInstant logical type
377+
Schema.FieldType fieldType = field.getType();
378+
if (fieldType.getTypeName().isLogicalType()) {
379+
Schema.@Nullable LogicalType<?, ?> logicalType = fieldType.getLogicalType();
380+
if (logicalType != null && logicalType.getIdentifier().equals(MicrosInstant.IDENTIFIER)) {
381+
return fromSpannerTimestamp(spannerTimestamp);
382+
}
383+
}
384+
// Default DATETIME behavior: convert to Joda DateTime
385+
return Instant.ofEpochSecond(spannerTimestamp.getSeconds()).toDateTime();
386+
369387
// TODO: implement logical date
370388
case DATE:
371389
return DateTime.parse(struct.getDate(column).toString());
@@ -407,11 +425,26 @@ private static void addIterableToStructBuilder(
407425
return struct.getBooleanList(column);
408426
case BYTES:
409427
return struct.getBytesList(column);
410-
// TODO: implement logical datetime
411428
case TIMESTAMP:
429+
// Check if expects MicrosInstant in arrays
430+
Schema.@Nullable FieldType elementType = field.getType().getCollectionElementType();
431+
if (elementType != null && elementType.getTypeName().isLogicalType()) {
432+
Schema.@Nullable LogicalType<?, ?> logicalType = elementType.getLogicalType();
433+
if (logicalType != null && logicalType.getIdentifier().equals(MicrosInstant.IDENTIFIER)) {
434+
// Return List<java.time.Instant> for MicrosInstant arrays
435+
return struct.getTimestampList(column).stream()
436+
.map(
437+
timestamp -> {
438+
return fromSpannerTimestamp(timestamp);
439+
})
440+
.collect(toList());
441+
}
442+
}
443+
// Default: return List<DateTime> for DATETIME type
412444
return struct.getTimestampList(column).stream()
413445
.map(timestamp -> Instant.ofEpochSecond(timestamp.getSeconds()).toDateTime())
414446
.collect(toList());
447+
415448
// TODO: implement logical date
416449
case DATE:
417450
return struct.getDateList(column).stream()

‎sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationUtilsTest.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,10 @@
2828
import com.google.cloud.spanner.Struct;
2929
import com.google.cloud.spanner.Type;
3030
import java.math.BigDecimal;
31+
import java.time.Instant;
3132
import java.util.List;
3233
import org.apache.beam.sdk.schemas.Schema;
34+
import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
3335
import org.apache.beam.sdk.values.Row;
3436
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
3537
import org.joda.time.DateTime;
@@ -44,6 +46,7 @@ public class MutationUtilsTest {
4446
private static final Struct EMPTY_STRUCT = Struct.newBuilder().build();
4547
private static final Struct INT64_STRUCT = Struct.newBuilder().set("int64").to(3L).build();
4648
private static final String TABLE = "some_table";
49+
private static final Instant TEST_INSTANT = Instant.parse("2024-01-15T10:30:00.123456Z");
4750

4851
private static final Schema WRITE_ROW_SCHEMA =
4952
Schema.builder()
@@ -71,6 +74,10 @@ public class MutationUtilsTest {
7174
.addNullableField("f_decimal", Schema.FieldType.DECIMAL)
7275
.addNullableField("f_byte", Schema.FieldType.BYTE)
7376
.addNullableField("f_iterable", Schema.FieldType.iterable(Schema.FieldType.INT64))
77+
.addNullableField("f_micros_instant", Schema.FieldType.logicalType(new MicrosInstant()))
78+
.addNullableField(
79+
"f_micros_instant_array",
80+
Schema.FieldType.array(Schema.FieldType.logicalType(new MicrosInstant())))
7481
.build();
7582

7683
private static final Row WRITE_ROW =
@@ -107,6 +114,8 @@ public class MutationUtilsTest {
107114
.withFieldValue("f_decimal", BigDecimal.valueOf(Long.MIN_VALUE))
108115
.withFieldValue("f_byte", Byte.parseByte("127"))
109116
.withFieldValue("f_iterable", ImmutableList.of(2L, 3L))
117+
.withFieldValue("f_micros_instant", TEST_INSTANT)
118+
.withFieldValue("f_micros_instant_array", ImmutableList.of(TEST_INSTANT, TEST_INSTANT))
110119
.build();
111120

112121
private static final Schema WRITE_ROW_SCHEMA_NULLS =
@@ -123,6 +132,10 @@ public class MutationUtilsTest {
123132
.addNullableField("f_array", Schema.FieldType.array(Schema.FieldType.INT64))
124133
.addNullableField(
125134
"f_struct_array", Schema.FieldType.array(Schema.FieldType.row(INT64_SCHEMA)))
135+
.addNullableField("f_micros_instant", Schema.FieldType.logicalType(new MicrosInstant()))
136+
.addNullableField(
137+
"f_micros_instant_array",
138+
Schema.FieldType.array(Schema.FieldType.logicalType(new MicrosInstant())))
126139
.build();
127140

128141
private static final Row WRITE_ROW_NULLS =
@@ -138,6 +151,8 @@ public class MutationUtilsTest {
138151
.addValue(null)
139152
.addValue(null)
140153
.addValue(null)
154+
.addValue(null)
155+
.addValue(null)
141156
.build();
142157

143158
private static final Schema KEY_SCHEMA =
@@ -153,6 +168,7 @@ public class MutationUtilsTest {
153168
.addNullableField("f_int32", Schema.FieldType.INT32)
154169
.addNullableField("f_decimal", Schema.FieldType.DECIMAL)
155170
.addNullableField("f_byte", Schema.FieldType.BYTE)
171+
.addNullableField("f_micros_instant", Schema.FieldType.logicalType(new MicrosInstant()))
156172
.build();
157173

158174
private static final Row KEY_ROW =
@@ -168,6 +184,7 @@ public class MutationUtilsTest {
168184
.withFieldValue("f_int32", 0x7fffffff)
169185
.withFieldValue("f_decimal", BigDecimal.valueOf(Long.MIN_VALUE))
170186
.withFieldValue("f_byte", Byte.parseByte("127"))
187+
.withFieldValue("f_micros_instant", TEST_INSTANT)
171188
.build();
172189

173190
private static final Schema KEY_SCHEMA_NULLS =
@@ -178,6 +195,7 @@ public class MutationUtilsTest {
178195
.addNullableField("f_bytes", Schema.FieldType.BYTES)
179196
.addNullableField("f_date_time", Schema.FieldType.DATETIME)
180197
.addNullableField("f_bool", Schema.FieldType.BOOLEAN)
198+
.addNullableField("f_micros_instant", Schema.FieldType.logicalType(new MicrosInstant()))
181199
.build();
182200

183201
private static final Row KEY_ROW_NULLS =
@@ -188,6 +206,7 @@ public class MutationUtilsTest {
188206
.addValue(null)
189207
.addValue(null)
190208
.addValue(null)
209+
.addValue(null)
191210
.build();
192211

193212
@Test
@@ -264,6 +283,7 @@ public void testCreateDeleteMutationFromRowWithNulls() {
264283
}
265284

266285
private static Mutation createDeleteMutation() {
286+
long micros = TEST_INSTANT.getEpochSecond() * 1_000_000L + TEST_INSTANT.getNano() / 1_000L;
267287
Key key =
268288
Key.newBuilder()
269289
.append(1L)
@@ -277,6 +297,7 @@ private static Mutation createDeleteMutation() {
277297
.append(0x7fffffff)
278298
.append(BigDecimal.valueOf(Long.MIN_VALUE))
279299
.append(Byte.parseByte("127"))
300+
.append(Timestamp.ofTimeMicroseconds(micros))
280301
.build();
281302
return Mutation.delete(TABLE, key);
282303
}
@@ -290,12 +311,14 @@ private static Mutation createDeleteMutationNulls() {
290311
.append((ByteArray) null)
291312
.append((Timestamp) null)
292313
.append((Boolean) null)
314+
.append((Timestamp) null)
293315
.build();
294316
return Mutation.delete(TABLE, key);
295317
}
296318

297319
private static Mutation createMutation(Mutation.Op operation) {
298320
Mutation.WriteBuilder builder = chooseBuilder(operation);
321+
long micros = TEST_INSTANT.getEpochSecond() * 1_000_000L + TEST_INSTANT.getNano() / 1_000L;
299322
return builder
300323
.set("f_int64")
301324
.to(1L)
@@ -353,6 +376,12 @@ private static Mutation createMutation(Mutation.Op operation) {
353376
.to(Byte.parseByte("127"))
354377
.set("f_iterable")
355378
.toInt64Array(ImmutableList.of(2L, 3L))
379+
.set("f_micros_instant")
380+
.to(Timestamp.ofTimeMicroseconds(micros))
381+
.set("f_micros_instant_array")
382+
.toTimestampArray(
383+
ImmutableList.of(
384+
Timestamp.ofTimeMicroseconds(micros), Timestamp.ofTimeMicroseconds(micros)))
356385
.build();
357386
}
358387

@@ -381,6 +410,10 @@ private static Mutation createMutationNulls(Mutation.Op operation) {
381410
.toInt64Array((List<Long>) null)
382411
.set("f_struct_array")
383412
.toStructArray(Type.struct(Type.StructField.of("int64", Type.int64())), null)
413+
.set("f_micros_instant")
414+
.to((Timestamp) null)
415+
.set("f_micros_instant_array")
416+
.toTimestampArray(null)
384417
.build();
385418
}
386419

‎sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/StructUtilsTest.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,10 @@
3333
import com.google.spanner.v1.StructType;
3434
import com.google.spanner.v1.TypeCode;
3535
import java.math.BigDecimal;
36+
import java.time.Instant;
3637
import java.util.List;
3738
import org.apache.beam.sdk.schemas.Schema;
39+
import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
3840
import org.apache.beam.sdk.values.Row;
3941
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
4042
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
@@ -45,6 +47,10 @@
4547
public class StructUtilsTest {
4648
private static final Schema EMPTY_SCHEMA = Schema.builder().build();
4749
private static final Schema INT64_SCHEMA = Schema.builder().addInt64Field("int64").build();
50+
private static final Timestamp TIMESTAMP = Timestamp.ofTimeMicroseconds(1234567890123456L);
51+
private static final Instant INSTANT =
52+
Instant.ofEpochSecond(
53+
1234567890123456L / 1_000_000L, (1234567890123456L % 1_000_000L) * 1_000L);
4854

4955
@Test
5056
public void testStructToBeamRow() {
@@ -286,6 +292,39 @@ public void testStructTypeToBeamRowSchemaFailsTypeNotSupported() {
286292
"Error processing struct to row: Unsupported type 'STRUCT'.", exception.getMessage());
287293
}
288294

295+
@Test
296+
public void testStructToBeamRowWithMicrosInstant() {
297+
Schema schema =
298+
Schema.builder()
299+
.addInt64Field("f_int64")
300+
.addNullableField("f_micros_instant", Schema.FieldType.logicalType(new MicrosInstant()))
301+
.addNullableField(
302+
"f_micros_instant_array",
303+
Schema.FieldType.array(Schema.FieldType.logicalType(new MicrosInstant())))
304+
.build();
305+
306+
Struct struct =
307+
Struct.newBuilder()
308+
.set("f_int64")
309+
.to(42L)
310+
.set("f_micros_instant")
311+
.to(TIMESTAMP)
312+
.set("f_micros_instant_array")
313+
.toTimestampArray(ImmutableList.of(TIMESTAMP, TIMESTAMP))
314+
.build();
315+
316+
Row result = StructUtils.structToBeamRow(struct, schema);
317+
318+
assertEquals(42L, result.getInt64("f_int64").longValue());
319+
320+
assertEquals(INSTANT, result.getValue("f_micros_instant"));
321+
322+
@SuppressWarnings("unchecked")
323+
List<Instant> instants = (List<Instant>) result.getValue("f_micros_instant_array");
324+
assertEquals(2, instants.size());
325+
assertEquals(INSTANT, instants.get(0));
326+
}
327+
289328
private StructType.Field getFieldForTypeCode(String name, TypeCode typeCode) {
290329
return StructType.Field.newBuilder()
291330
.setName(name)

0 commit comments

Comments
 (0)