Skip to content

Commit 582bb45

Browse files
sachinnn99sachinnn99
authored
[Java SDK] Infer Beam logical types for JSR-310 and UUID fields (#38194)
Add schema inference for java.time.LocalDate, LocalTime, LocalDateTime, Instant, and UUID as Beam logical types (SqlTypes.DATE, TIME, DATETIME, NanosInstant, SqlTypes.UUID) in both POJO and JavaBean schemas. This enables Beam Rows produced from POJOs with JSR-310 fields to be schema-assignable to rows from external systems (e.g. IcebergIO) that use the same logical types, fixing the incompatibility reported in #37524. The Avro extension overrides the new convertLogicalType hook to preserve existing Joda bridging for java.time.Instant and LocalDate.
1 parent c33bc97 commit 582bb45

10 files changed

Lines changed: 751 additions & 36 deletions

File tree

‎sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/ByteBuddyUtils.java‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@
2828
import java.lang.reflect.Parameter;
2929
import java.lang.reflect.Type;
3030
import java.nio.ByteBuffer;
31+
import java.time.LocalDate;
32+
import java.time.LocalDateTime;
33+
import java.time.LocalTime;
3134
import java.util.ArrayList;
3235
import java.util.Arrays;
3336
import java.util.Collection;
@@ -38,6 +41,7 @@
3841
import java.util.Optional;
3942
import java.util.Set;
4043
import java.util.SortedMap;
44+
import java.util.UUID;
4145
import net.bytebuddy.ByteBuddy;
4246
import net.bytebuddy.NamingStrategy;
4347
import net.bytebuddy.NamingStrategy.SuffixingRandom.BaseNameResolver;
@@ -78,6 +82,9 @@
7882
import org.apache.beam.sdk.schemas.FieldValueHaver;
7983
import org.apache.beam.sdk.schemas.FieldValueSetter;
8084
import org.apache.beam.sdk.schemas.FieldValueTypeInformation;
85+
import org.apache.beam.sdk.schemas.Schema.LogicalType;
86+
import org.apache.beam.sdk.schemas.logicaltypes.NanosInstant;
87+
import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
8188
import org.apache.beam.sdk.util.Preconditions;
8289
import org.apache.beam.sdk.util.common.ReflectHelpers;
8390
import org.apache.beam.sdk.values.TypeDescriptor;
@@ -120,6 +127,65 @@ public class ByteBuddyUtils {
120127
private static final ForLoadedType ENUM_TYPE = new ForLoadedType(Enum.class);
121128
private static final ForLoadedType BYTE_BUDDY_UTILS_TYPE =
122129
new ForLoadedType(ByteBuddyUtils.class);
130+
private static final ForLoadedType LOGICAL_TYPE_TYPE = new ForLoadedType(LogicalType.class);
131+
132+
// Static LogicalType instances used by codegen for JSR-310 and UUID POJO/Bean fields. The
133+
// generated bytecode loads these via FieldAccess and invokes toBaseType / toInputType, so the
134+
// fields must be public so generated classes in user packages can access them.
135+
// See logicalTypeFieldName(...) for the type → field name mapping.
136+
public static final LogicalType<LocalDate, Long> JAVA_LOCAL_DATE_LOGICAL_TYPE = SqlTypes.DATE;
137+
public static final LogicalType<LocalTime, Long> JAVA_LOCAL_TIME_LOGICAL_TYPE = SqlTypes.TIME;
138+
public static final LogicalType<LocalDateTime, org.apache.beam.sdk.values.Row>
139+
JAVA_LOCAL_DATE_TIME_LOGICAL_TYPE = SqlTypes.DATETIME;
140+
public static final LogicalType<java.time.Instant, org.apache.beam.sdk.values.Row>
141+
JAVA_INSTANT_LOGICAL_TYPE = new NanosInstant();
142+
public static final LogicalType<UUID, org.apache.beam.sdk.values.Row> JAVA_UUID_LOGICAL_TYPE =
143+
SqlTypes.UUID;
144+
145+
/**
146+
* Returns the {@link Schema.LogicalType} that {@link StaticSchemaInference} infers for the given
147+
* Java raw type, or {@code null} if no JSR-310 / UUID inference applies.
148+
*/
149+
static @Nullable LogicalType<?, ?> inferredLogicalTypeFor(Class<?> rawType) {
150+
if (LocalDate.class.equals(rawType)) {
151+
return JAVA_LOCAL_DATE_LOGICAL_TYPE;
152+
} else if (LocalTime.class.equals(rawType)) {
153+
return JAVA_LOCAL_TIME_LOGICAL_TYPE;
154+
} else if (LocalDateTime.class.equals(rawType)) {
155+
return JAVA_LOCAL_DATE_TIME_LOGICAL_TYPE;
156+
} else if (java.time.Instant.class.equals(rawType)) {
157+
return JAVA_INSTANT_LOGICAL_TYPE;
158+
} else if (UUID.class.equals(rawType)) {
159+
return JAVA_UUID_LOGICAL_TYPE;
160+
}
161+
return null;
162+
}
163+
164+
/** Maps a Java raw type to the static field name in {@link ByteBuddyUtils} that holds it. */
165+
private static String logicalTypeFieldName(Class<?> rawType) {
166+
if (LocalDate.class.equals(rawType)) {
167+
return "JAVA_LOCAL_DATE_LOGICAL_TYPE";
168+
} else if (LocalTime.class.equals(rawType)) {
169+
return "JAVA_LOCAL_TIME_LOGICAL_TYPE";
170+
} else if (LocalDateTime.class.equals(rawType)) {
171+
return "JAVA_LOCAL_DATE_TIME_LOGICAL_TYPE";
172+
} else if (java.time.Instant.class.equals(rawType)) {
173+
return "JAVA_INSTANT_LOGICAL_TYPE";
174+
} else if (UUID.class.equals(rawType)) {
175+
return "JAVA_UUID_LOGICAL_TYPE";
176+
}
177+
throw new IllegalArgumentException("Not an inferred logical type: " + rawType);
178+
}
179+
180+
/** Stack manipulation that pushes the static {@link LogicalType} for the given Java type. */
181+
private static StackManipulation loadLogicalType(Class<?> rawType) {
182+
return FieldAccess.forField(
183+
BYTE_BUDDY_UTILS_TYPE
184+
.getDeclaredFields()
185+
.filter(ElementMatchers.named(logicalTypeFieldName(rawType)))
186+
.getOnly())
187+
.read();
188+
}
123189

124190
/**
125191
* A naming strategy for ByteBuddy classes.
@@ -286,6 +352,8 @@ public T convert(TypeDescriptor<?> typeDescriptor) {
286352
return convertDateTime(typeDescriptor);
287353
} else if (typeDescriptor.isSubtypeOf(TypeDescriptor.of(ReadablePartial.class))) {
288354
return convertDateTime(typeDescriptor);
355+
} else if (inferredLogicalTypeFor(typeDescriptor.getRawType()) != null) {
356+
return convertLogicalType(typeDescriptor);
289357
} else if (typeDescriptor.isSubtypeOf(TypeDescriptor.of(ByteBuffer.class))) {
290358
return convertByteBuffer(typeDescriptor);
291359
} else if (typeDescriptor.isSubtypeOf(TypeDescriptor.of(CharSequence.class))) {
@@ -324,6 +392,14 @@ protected StackManipulation shortCircuitReturnNull(
324392

325393
protected abstract T convertDateTime(TypeDescriptor<?> type);
326394

395+
/**
396+
* Handles JSR-310 ({@link LocalDate}, {@link LocalTime}, {@link LocalDateTime}, {@link
397+
* java.time.Instant}) and {@link UUID} fields, which {@link StaticSchemaInference} infers as
398+
* Beam {@link LogicalType}s. Subclasses emit code that round-trips through the corresponding
399+
* static {@link LogicalType} instance ({@link #JAVA_LOCAL_DATE_LOGICAL_TYPE} etc.).
400+
*/
401+
protected abstract T convertLogicalType(TypeDescriptor<?> type);
402+
327403
protected abstract T convertByteBuffer(TypeDescriptor<?> type);
328404

329405
protected abstract T convertCharSequence(TypeDescriptor<?> type);
@@ -401,6 +477,15 @@ protected Type convertDateTime(TypeDescriptor<?> type) {
401477
return Instant.class;
402478
}
403479

480+
@Override
481+
protected Type convertLogicalType(TypeDescriptor<?> type) {
482+
// The codegen-generated getter returns the LogicalType's base value (Long for
483+
// Date/Time, Row for DateTime/NanosInstant/UUID). Object.class is a safe upper bound
484+
// for the FieldValueGetter signature; the framework's GetLogicalInputType wrapper
485+
// converts back to the input type before exposing the value to user code.
486+
return Object.class;
487+
}
488+
404489
@Override
405490
protected Type convertByteBuffer(TypeDescriptor<?> type) {
406491
return byte[].class;
@@ -915,6 +1000,26 @@ protected StackManipulation convertDateTime(TypeDescriptor<?> type) {
9151000
return new ShortCircuitReturnNull(readValue, stackManipulation);
9161001
}
9171002

1003+
@Override
1004+
protected StackManipulation convertLogicalType(TypeDescriptor<?> type) {
1005+
// Equivalent code: return STATIC_LOGICAL_TYPE.toBaseType(value);
1006+
// where STATIC_LOGICAL_TYPE is one of the JAVA_*_LOGICAL_TYPE static fields on
1007+
// ByteBuddyUtils. The base type is Long (for LocalDate, LocalTime) or Row (for the
1008+
// others); both are reference types so no boxing/casting is needed beyond the invoke.
1009+
StackManipulation stackManipulation =
1010+
new Compound(
1011+
loadLogicalType(type.getRawType()),
1012+
readValue,
1013+
MethodInvocation.invoke(
1014+
LOGICAL_TYPE_TYPE
1015+
.getDeclaredMethods()
1016+
.filter(
1017+
ElementMatchers.named("toBaseType")
1018+
.and(ElementMatchers.takesArguments(1)))
1019+
.getOnly()));
1020+
return new ShortCircuitReturnNull(readValue, stackManipulation);
1021+
}
1022+
9181023
@Override
9191024
protected StackManipulation convertByteBuffer(TypeDescriptor<?> type) {
9201025
// Generate the following code:
@@ -1361,6 +1466,29 @@ protected StackManipulation convertEnum(TypeDescriptor<?> type) {
13611466
return new ShortCircuitReturnNull(readValue, stackManipulation);
13621467
}
13631468

1469+
@Override
1470+
protected StackManipulation convertLogicalType(TypeDescriptor<?> type) {
1471+
// Equivalent code: return (JavaType) STATIC_LOGICAL_TYPE.toInputType(value);
1472+
// FromRowUsingCreator already converted the row's input-type value (e.g. LocalDate)
1473+
// to the LogicalType's base value (e.g. Long) before invoking the generated creator,
1474+
// so we receive the base type here and need to project back to the POJO field's
1475+
// Java type.
1476+
ForLoadedType loadedType = new ForLoadedType(type.getRawType());
1477+
StackManipulation stackManipulation =
1478+
new Compound(
1479+
loadLogicalType(type.getRawType()),
1480+
readValue,
1481+
MethodInvocation.invoke(
1482+
LOGICAL_TYPE_TYPE
1483+
.getDeclaredMethods()
1484+
.filter(
1485+
ElementMatchers.named("toInputType")
1486+
.and(ElementMatchers.takesArguments(1)))
1487+
.getOnly()),
1488+
TypeCasting.to(loadedType));
1489+
return new ShortCircuitReturnNull(readValue, stackManipulation);
1490+
}
1491+
13641492
@Override
13651493
protected StackManipulation convertDefault(TypeDescriptor<?> type) {
13661494
return readValue;

‎sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/StaticSchemaInference.java‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,17 +22,23 @@
2222
import java.lang.reflect.ParameterizedType;
2323
import java.math.BigDecimal;
2424
import java.nio.ByteBuffer;
25+
import java.time.LocalDate;
26+
import java.time.LocalDateTime;
27+
import java.time.LocalTime;
2528
import java.util.Arrays;
2629
import java.util.Collection;
2730
import java.util.HashMap;
2831
import java.util.List;
2932
import java.util.Map;
33+
import java.util.UUID;
3034
import java.util.function.Function;
3135
import java.util.stream.Collectors;
3236
import org.apache.beam.sdk.schemas.FieldValueTypeInformation;
3337
import org.apache.beam.sdk.schemas.Schema;
3438
import org.apache.beam.sdk.schemas.Schema.FieldType;
3539
import org.apache.beam.sdk.schemas.logicaltypes.EnumerationType;
40+
import org.apache.beam.sdk.schemas.logicaltypes.NanosInstant;
41+
import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
3642
import org.apache.beam.sdk.values.TypeDescriptor;
3743
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
3844
import org.joda.time.ReadableInstant;
@@ -127,7 +133,8 @@ public static Schema.FieldType fieldFromType(
127133
return fieldFromType(type, fieldValueTypeSupplier, new HashMap<>());
128134
}
129135

130-
// TODO(https://github.com/apache/beam/issues/21567): support type inference for logical types
136+
// TODO(https://github.com/apache/beam/issues/21567): support inference for additional/custom
137+
// logical types
131138
private static Schema.FieldType fieldFromType(
132139
TypeDescriptor type,
133140
FieldValueTypeSupplier fieldValueTypeSupplier,
@@ -177,6 +184,16 @@ private static Schema.FieldType fieldFromType(
177184
return FieldType.STRING;
178185
} else if (type.isSubtypeOf(TypeDescriptor.of(ReadableInstant.class))) {
179186
return FieldType.DATETIME;
187+
} else if (type.getRawType().equals(LocalDate.class)) {
188+
return FieldType.logicalType(SqlTypes.DATE);
189+
} else if (type.getRawType().equals(LocalTime.class)) {
190+
return FieldType.logicalType(SqlTypes.TIME);
191+
} else if (type.getRawType().equals(LocalDateTime.class)) {
192+
return FieldType.logicalType(SqlTypes.DATETIME);
193+
} else if (type.getRawType().equals(java.time.Instant.class)) {
194+
return FieldType.logicalType(new NanosInstant());
195+
} else if (type.getRawType().equals(UUID.class)) {
196+
return FieldType.logicalType(SqlTypes.UUID);
180197
} else if (type.isSubtypeOf(TypeDescriptor.of(ByteBuffer.class))) {
181198
return FieldType.BYTES;
182199
} else if (type.isSubtypeOf(TypeDescriptor.of(Iterable.class))) {

‎sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/JavaBeanSchemaTest.java‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.CASE_FORMAT_BEAM_SCHEMA;
2525
import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.FIELD_WITH_DESCRIPTION_BEAN_SCHEMA;
2626
import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.ITERABLE_BEAM_SCHEMA;
27+
import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.JAVA_TIME_BEAN_SCHEMA;
2728
import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.NESTED_ARRAYS_BEAM_SCHEMA;
2829
import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.NESTED_ARRAY_BEAN_SCHEMA;
2930
import static org.apache.beam.sdk.schemas.utils.TestJavaBeans.NESTED_BEAN_SCHEMA;
@@ -46,9 +47,15 @@
4647
import java.math.BigDecimal;
4748
import java.nio.ByteBuffer;
4849
import java.nio.charset.StandardCharsets;
50+
import java.time.LocalDate;
51+
import java.time.LocalDateTime;
52+
import java.time.LocalTime;
4953
import java.util.Arrays;
5054
import java.util.List;
5155
import java.util.Map;
56+
import java.util.UUID;
57+
import org.apache.beam.sdk.schemas.logicaltypes.NanosInstant;
58+
import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
5259
import org.apache.beam.sdk.schemas.utils.SchemaTestUtils;
5360
import org.apache.beam.sdk.schemas.utils.TestJavaBeans;
5461
import org.apache.beam.sdk.schemas.utils.TestJavaBeans.AllNullableBean;
@@ -57,6 +64,7 @@
5764
import org.apache.beam.sdk.schemas.utils.TestJavaBeans.BeanWithNoCreateOption;
5865
import org.apache.beam.sdk.schemas.utils.TestJavaBeans.BeanWithRenamedFieldsAndSetters;
5966
import org.apache.beam.sdk.schemas.utils.TestJavaBeans.IterableBean;
67+
import org.apache.beam.sdk.schemas.utils.TestJavaBeans.JavaTimeBean;
6068
import org.apache.beam.sdk.schemas.utils.TestJavaBeans.MismatchingNullableBean;
6169
import org.apache.beam.sdk.schemas.utils.TestJavaBeans.NestedArrayBean;
6270
import org.apache.beam.sdk.schemas.utils.TestJavaBeans.NestedArraysBean;
@@ -179,6 +187,81 @@ public void testFromRow() throws NoSuchSchemaException {
179187
assertEquals("stringbuilder", bean.getStringBuilder().toString());
180188
}
181189

190+
@Test
191+
public void testJavaTimeSchema() throws NoSuchSchemaException {
192+
SchemaRegistry registry = SchemaRegistry.createDefault();
193+
Schema schema = registry.getSchema(JavaTimeBean.class);
194+
SchemaTestUtils.assertSchemaEquivalent(JAVA_TIME_BEAN_SCHEMA, schema);
195+
Schema icebergStyleSchema =
196+
Schema.builder()
197+
.addLogicalTypeField("localDate", SqlTypes.DATE)
198+
.addLogicalTypeField("localTime", SqlTypes.TIME)
199+
.addLogicalTypeField("localDateTime", SqlTypes.DATETIME)
200+
.addLogicalTypeField("instant", new NanosInstant())
201+
.addLogicalTypeField("uuid", SqlTypes.UUID)
202+
.build();
203+
assertTrue(schema.assignableToIgnoreNullable(icebergStyleSchema));
204+
}
205+
206+
@Test
207+
public void testJavaTimeToRow() throws NoSuchSchemaException {
208+
SchemaRegistry registry = SchemaRegistry.createDefault();
209+
JavaTimeBean bean = new JavaTimeBean();
210+
bean.setLocalDate(LocalDate.of(2024, 1, 15));
211+
bean.setLocalTime(LocalTime.of(10, 30, 45));
212+
bean.setLocalDateTime(LocalDateTime.of(2024, 1, 15, 10, 30, 45));
213+
bean.setInstant(java.time.Instant.ofEpochSecond(1_705_315_845L, 123_456_789L));
214+
bean.setUuid(UUID.fromString("11111111-2222-3333-4444-555555555555"));
215+
216+
Row row = registry.getToRowFunction(JavaTimeBean.class).apply(bean);
217+
218+
assertEquals(5, row.getFieldCount());
219+
assertEquals(bean.getLocalDate(), row.getLogicalTypeValue("localDate", LocalDate.class));
220+
assertEquals(bean.getLocalTime(), row.getLogicalTypeValue("localTime", LocalTime.class));
221+
assertEquals(
222+
bean.getLocalDateTime(), row.getLogicalTypeValue("localDateTime", LocalDateTime.class));
223+
assertEquals(bean.getInstant(), row.getLogicalTypeValue("instant", java.time.Instant.class));
224+
assertEquals(bean.getUuid(), row.getLogicalTypeValue("uuid", UUID.class));
225+
}
226+
227+
@Test
228+
public void testJavaTimeFromRow() throws NoSuchSchemaException {
229+
SchemaRegistry registry = SchemaRegistry.createDefault();
230+
LocalDate localDate = LocalDate.of(2024, 1, 15);
231+
LocalTime localTime = LocalTime.of(10, 30, 45);
232+
LocalDateTime localDateTime = LocalDateTime.of(2024, 1, 15, 10, 30, 45);
233+
java.time.Instant instant = java.time.Instant.ofEpochSecond(1_705_315_845L, 123_456_789L);
234+
UUID uuid = UUID.fromString("11111111-2222-3333-4444-555555555555");
235+
Row row =
236+
Row.withSchema(JAVA_TIME_BEAN_SCHEMA)
237+
.addValues(localDate, localTime, localDateTime, instant, uuid)
238+
.build();
239+
240+
JavaTimeBean bean = registry.getFromRowFunction(JavaTimeBean.class).apply(row);
241+
242+
assertEquals(localDate, bean.getLocalDate());
243+
assertEquals(localTime, bean.getLocalTime());
244+
assertEquals(localDateTime, bean.getLocalDateTime());
245+
assertEquals(instant, bean.getInstant());
246+
assertEquals(uuid, bean.getUuid());
247+
}
248+
249+
@Test
250+
public void testJavaTimeRoundTrip() throws NoSuchSchemaException {
251+
SchemaRegistry registry = SchemaRegistry.createDefault();
252+
JavaTimeBean original = new JavaTimeBean();
253+
original.setLocalDate(LocalDate.of(2024, 1, 15));
254+
original.setLocalTime(LocalTime.of(10, 30, 45));
255+
original.setLocalDateTime(LocalDateTime.of(2024, 1, 15, 10, 30, 45));
256+
original.setInstant(java.time.Instant.ofEpochSecond(1_705_315_845L, 123_456_789L));
257+
original.setUuid(UUID.fromString("11111111-2222-3333-4444-555555555555"));
258+
259+
Row row = registry.getToRowFunction(JavaTimeBean.class).apply(original);
260+
JavaTimeBean roundTripped = registry.getFromRowFunction(JavaTimeBean.class).apply(row);
261+
262+
assertEquals(original, roundTripped);
263+
}
264+
182265
@Test
183266
public void testNullableToRow() throws NoSuchSchemaException {
184267
SchemaRegistry registry = SchemaRegistry.createDefault();

0 commit comments

Comments
 (0)