Skip to content

Commit 9c561e2

Browse files
ahmedabu98Ahmed AbualsaudAhmed Abualsaud
authored
[Iceberg] Make timestamptz return new Timestamp.MICROS logical type (#39344)
* switch to new timestamp logical type * changes and trigger ITs * address comments * format changes * add link * tighten BQ conversion logic * spotless * negative test; cleanup * nuance timestamp conversion; use if-block * nuance timestamp conversion; use if-block * add Timestamp to test and trigger ITs --------- Co-authored-by: Ahmed Abualsaud <ahmedabualsaud@MacBook-Pro-2.local> Co-authored-by: Ahmed Abualsaud <ahmedabualsaud@mac.attlocal.net>
1 parent 980c114 commit 9c561e2

18 files changed

Lines changed: 248 additions & 44 deletions

File tree

‎.github/trigger_files/IO_Iceberg_Integration_Tests.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": 2
3+
"modification": 3
44
}

‎.github/trigger_files/beam_PostCommit_Python_Xlang_IO_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": 2
44
}

‎CHANGES.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,10 @@
8181
## Breaking Changes
8282

8383
* (Python) Removed `google-perftools` from the SDK container images. Users who wish to use `--profiler_agent=tcmalloc` should install google-perftools APT package in their custom container images separately ([#39323](https://github.com/apache/beam/issues/39323)).
84+
* [IcebergIO] Reading a `timestamptz` column will now return a `Timestamp.MICROS` Beam logical type to preserve
85+
microseconds (the old Beam `Schema.FieldType#DATETIME` primitive type truncates past milliseconds). This may break
86+
existing streaming read pipelines. It also breaks Python reads when a `timestamptz` column is present. Use pipeline
87+
option `--updateCompatibilityVersion=2.75.0` (or any older version) to keep the old behavior ([#39344](https://github.com/apache/beam/issues/39344)).
8488
* `DoFn.process` returning a `str`, `bytes`, or `dict` (instead of an iterable wrapping one) now raises a `TypeError` rather than silently iterating per-character/byte/key (Python) ([#18712](https://github.com/apache/beam/issues/18712)).
8589
* (Java) Added `DRAINING` and `DRAINED` states to `PipelineResult`, including runner state mappings and Dataflow update handling ([#39020](https://github.com/apache/beam/issues/39020)).
8690
* (Python) Typehints of dataclass fields are honored during type inferences. To restore the behavior of fallback-to-any,

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -932,6 +932,12 @@ public static Row toBeamRow(Schema rowSchema, TableSchema bqSchema, TableRow jso
932932
return java.time.Instant.parse(jsonBQString);
933933
}
934934
} else if (fieldType.isLogicalType(Timestamp.IDENTIFIER)) {
935+
if (!jsonBQString.contains("UTC")) {
936+
BigDecimal bd = new BigDecimal(jsonBQString);
937+
long seconds = bd.longValue();
938+
long nanos = bd.subtract(BigDecimal.valueOf(seconds)).movePointRight(9).longValue();
939+
return java.time.Instant.ofEpochSecond(seconds, nanos);
940+
}
935941
return VAR_PRECISION_FORMATTER.parse(jsonBQString, java.time.Instant::from);
936942
}
937943
}

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

Lines changed: 32 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import java.time.LocalDate;
4242
import java.time.LocalDateTime;
4343
import java.time.LocalTime;
44+
import java.time.OffsetDateTime;
4445
import java.util.Arrays;
4546
import java.util.Base64;
4647
import java.util.Collections;
@@ -1444,22 +1445,43 @@ public void testToBeamRow_timestampNanos_utcSuffix() {
14441445

14451446
@Test
14461447
@SuppressWarnings("JavaInstantGetSecondsGetNano")
1447-
public void testToBeamRow_timestampMicros_utcSuffix() {
1448+
public void testToBeamRow_timestampMicros() {
14481449
Schema schema = Schema.builder().addLogicalTypeField("ts", Timestamp.MICROS).build();
14491450

14501451
// BigQuery format with " UTC" suffix
14511452
String timestamp = "2024-08-10 16:52:07.123456 UTC";
1453+
String parsableTimestamp = "2024-08-10T16:52:07.123456Z";
1454+
String negativeTimestamp = "1960-08-10T16:52:07.000123Z";
14521455

1453-
Row beamRow = BigQueryUtils.toBeamRow(schema, new TableRow().set("ts", timestamp));
1456+
java.time.Instant instant = OffsetDateTime.parse(parsableTimestamp).toInstant();
1457+
String value = instant.getEpochSecond() + "." + instant.getNano() / 1000;
1458+
java.time.Instant negInstant = OffsetDateTime.parse(negativeTimestamp).toInstant();
1459+
String negValue =
1460+
BigDecimal.valueOf(negInstant.getEpochSecond())
1461+
.add(BigDecimal.valueOf(negInstant.getNano(), 9))
1462+
.toPlainString();
14541463

1455-
java.time.Instant actual = (java.time.Instant) beamRow.getValue("ts");
1456-
assertEquals(2024, actual.atZone(java.time.ZoneOffset.UTC).getYear());
1457-
assertEquals(8, actual.atZone(java.time.ZoneOffset.UTC).getMonthValue());
1458-
assertEquals(10, actual.atZone(java.time.ZoneOffset.UTC).getDayOfMonth());
1459-
assertEquals(16, actual.atZone(java.time.ZoneOffset.UTC).getHour());
1460-
assertEquals(52, actual.atZone(java.time.ZoneOffset.UTC).getMinute());
1461-
assertEquals(7, actual.atZone(java.time.ZoneOffset.UTC).getSecond());
1462-
assertEquals(123456000, actual.getNano());
1464+
List<TableRow> testRows =
1465+
Arrays.asList(
1466+
new TableRow().set("ts", timestamp),
1467+
new TableRow().set("ts", value),
1468+
new TableRow().set("negative", true).set("ts", negValue));
1469+
1470+
for (TableRow row : testRows) {
1471+
Row beamRow = BigQueryUtils.toBeamRow(schema, row);
1472+
1473+
java.time.Instant actual = (java.time.Instant) beamRow.getValue("ts");
1474+
1475+
assertEquals(
1476+
row.get("negative") == null ? 2024 : 1960,
1477+
actual.atZone(java.time.ZoneOffset.UTC).getYear());
1478+
assertEquals(8, actual.atZone(java.time.ZoneOffset.UTC).getMonthValue());
1479+
assertEquals(10, actual.atZone(java.time.ZoneOffset.UTC).getDayOfMonth());
1480+
assertEquals(16, actual.atZone(java.time.ZoneOffset.UTC).getHour());
1481+
assertEquals(52, actual.atZone(java.time.ZoneOffset.UTC).getMinute());
1482+
assertEquals(7, actual.atZone(java.time.ZoneOffset.UTC).getSecond());
1483+
assertEquals(row.get("negative") == null ? 123456000 : 123000, actual.getNano());
1484+
}
14631485
}
14641486

14651487
@Test

‎sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import java.util.Map;
2626
import org.apache.beam.sdk.annotations.Internal;
2727
import org.apache.beam.sdk.io.Read;
28+
import org.apache.beam.sdk.options.StreamingOptions;
2829
import org.apache.beam.sdk.schemas.Schema;
2930
import org.apache.beam.sdk.transforms.PTransform;
3031
import org.apache.beam.sdk.values.PBegin;
@@ -699,12 +700,23 @@ public PCollection<Row> expand(PBegin input) {
699700

700701
Table table = TableCache.get(getCatalogConfig(), tableId);
701702

703+
@Nullable
704+
String updateCompatibilityVersion =
705+
input
706+
.getPipeline()
707+
.getOptions()
708+
.as(StreamingOptions.class)
709+
.getUpdateCompatibilityVersion();
710+
702711
IcebergScanConfig scanConfig =
703712
IcebergScanConfig.builder()
704713
.setCatalogConfig(getCatalogConfig())
705714
.setScanType(IcebergScanConfig.ScanType.TABLE)
706715
.setTableIdentifier(tableId)
707-
.setSchema(IcebergUtils.icebergSchemaToBeamSchema(table.schema()))
716+
.setSchema(
717+
IcebergUtils.icebergSchemaToBeamSchema(
718+
table.schema(), updateCompatibilityVersion))
719+
.setUpdateCompatibilityVersion(updateCompatibilityVersion)
708720
.setFromSnapshotInclusive(getFromSnapshot())
709721
.setToSnapshot(getToSnapshot())
710722
.setFromTimestamp(getFromTimestamp())

‎sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,8 @@ public org.apache.iceberg.Schema recordIdSchema() {
164164

165165
public Schema rowIdBeamSchema() {
166166
if (cachedRowIdBeamSchema == null) {
167-
cachedRowIdBeamSchema = icebergSchemaToBeamSchema(recordIdSchema());
167+
cachedRowIdBeamSchema =
168+
icebergSchemaToBeamSchema(recordIdSchema(), getUpdateCompatibilityVersion());
168169
}
169170
return cachedRowIdBeamSchema;
170171
}
@@ -237,6 +238,9 @@ public Expression getFilter() {
237238
@Pure
238239
public abstract boolean getUseCdc();
239240

241+
@Pure
242+
public abstract @Nullable String getUpdateCompatibilityVersion();
243+
240244
@Pure
241245
public abstract @Nullable Boolean getStreaming();
242246

@@ -335,6 +339,9 @@ public Builder setTableIdentifier(String... names) {
335339

336340
public abstract Builder setUseCdc(boolean useCdc);
337341

342+
public abstract Builder setUpdateCompatibilityVersion(
343+
@Nullable String updateCompatibilityVersion);
344+
338345
public abstract Builder setStreaming(@Nullable Boolean streaming);
339346

340347
public abstract Builder setPollInterval(@Nullable Duration pollInterval);

‎sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java‎

Lines changed: 59 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,9 @@
4040
import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
4141
import org.apache.beam.sdk.schemas.logicaltypes.PassThroughLogicalType;
4242
import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
43+
import org.apache.beam.sdk.schemas.logicaltypes.Timestamp;
4344
import org.apache.beam.sdk.util.Preconditions;
45+
import org.apache.beam.sdk.util.construction.TransformUpgrader;
4446
import org.apache.beam.sdk.values.PCollection;
4547
import org.apache.beam.sdk.values.Row;
4648
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
@@ -83,7 +85,8 @@ private IcebergUtils() {}
8385
.put(MicrosInstant.IDENTIFIER, Types.TimestampType.withZone())
8486
.build();
8587

86-
private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) {
88+
private static Schema.FieldType icebergTypeToBeamFieldType(
89+
final Type type, @Nullable String updateCompatibilityVersion) {
8790
switch (type.typeId()) {
8891
case BOOLEAN:
8992
return Schema.FieldType.BOOLEAN;
@@ -102,7 +105,14 @@ private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) {
102105
case TIMESTAMP:
103106
Types.TimestampType ts = (Types.TimestampType) type.asPrimitiveType();
104107
if (ts.shouldAdjustToUTC()) {
105-
return Schema.FieldType.DATETIME;
108+
// timestamptz. The micros-precision Timestamp logical type preserves microseconds, while
109+
// the legacy DATETIME (joda) mapping truncates to millis. Gated for update compatibility.
110+
if (updateCompatibilityVersion != null
111+
&& !updateCompatibilityVersion.isEmpty()
112+
&& TransformUpgrader.compareVersions(updateCompatibilityVersion, "2.76.0") < 0) {
113+
return Schema.FieldType.DATETIME;
114+
}
115+
return Schema.FieldType.logicalType(Timestamp.MICROS);
106116
}
107117
return Schema.FieldType.logicalType(SqlTypes.DATETIME);
108118
case STRING:
@@ -114,36 +124,51 @@ private static Schema.FieldType icebergTypeToBeamFieldType(final Type type) {
114124
case DECIMAL:
115125
return Schema.FieldType.DECIMAL;
116126
case STRUCT:
117-
return Schema.FieldType.row(icebergStructTypeToBeamSchema(type.asStructType()));
127+
return Schema.FieldType.row(
128+
icebergStructTypeToBeamSchema(type.asStructType(), updateCompatibilityVersion));
118129
case LIST:
119-
return Schema.FieldType.array(icebergTypeToBeamFieldType(type.asListType().elementType()));
130+
return Schema.FieldType.array(
131+
icebergTypeToBeamFieldType(
132+
type.asListType().elementType(), updateCompatibilityVersion));
120133
case MAP:
121134
return Schema.FieldType.map(
122-
icebergTypeToBeamFieldType(type.asMapType().keyType()),
123-
icebergTypeToBeamFieldType(type.asMapType().valueType()));
135+
icebergTypeToBeamFieldType(type.asMapType().keyType(), updateCompatibilityVersion),
136+
icebergTypeToBeamFieldType(type.asMapType().valueType(), updateCompatibilityVersion));
124137
default:
125138
throw new RuntimeException("Unrecognized Iceberg Type: " + type.typeId());
126139
}
127140
}
128141

129-
private static Schema.Field icebergFieldToBeamField(final Types.NestedField field) {
130-
return Schema.Field.of(field.name(), icebergTypeToBeamFieldType(field.type()))
142+
private static Schema.Field icebergFieldToBeamField(
143+
final Types.NestedField field, @Nullable String updateCompatibilityVersion) {
144+
return Schema.Field.of(
145+
field.name(), icebergTypeToBeamFieldType(field.type(), updateCompatibilityVersion))
131146
.withNullable(field.isOptional());
132147
}
133148

134149
/** Converts an Iceberg {@link org.apache.iceberg.Schema} to a Beam {@link Schema}. */
135150
public static Schema icebergSchemaToBeamSchema(final org.apache.iceberg.Schema schema) {
151+
return icebergSchemaToBeamSchema(schema, null);
152+
}
153+
154+
/**
155+
* Converts an Iceberg {@link org.apache.iceberg.Schema} to a Beam {@link Schema}, accounting for
156+
* update compatibility.
157+
*/
158+
public static Schema icebergSchemaToBeamSchema(
159+
final org.apache.iceberg.Schema schema, @Nullable String updateCompatibilityVersion) {
136160
Schema.Builder builder = Schema.builder();
137161
for (Types.NestedField f : schema.columns()) {
138-
builder.addField(icebergFieldToBeamField(f));
162+
builder.addField(icebergFieldToBeamField(f, updateCompatibilityVersion));
139163
}
140164
return builder.build();
141165
}
142166

143-
private static Schema icebergStructTypeToBeamSchema(final Types.StructType struct) {
167+
private static Schema icebergStructTypeToBeamSchema(
168+
final Types.StructType struct, @Nullable String updateCompatibilityVersion) {
144169
Schema.Builder builder = Schema.builder();
145170
for (Types.NestedField f : struct.fields()) {
146-
builder.addField(icebergFieldToBeamField(f));
171+
builder.addField(icebergFieldToBeamField(f, updateCompatibilityVersion));
147172
}
148173
return builder.build();
149174
}
@@ -198,7 +223,17 @@ static TypeAndMaxId beamFieldTypeToIcebergFieldType(
198223
String logicalTypeIdentifier = logicalType.getIdentifier();
199224
@Nullable Type type = BEAM_LOGICAL_TYPES_TO_ICEBERG_TYPES.get(logicalTypeIdentifier);
200225
if (type == null) {
201-
throw new RuntimeException("Unsupported Beam logical type " + logicalTypeIdentifier);
226+
if (beamType.isLogicalType(Timestamp.IDENTIFIER)) {
227+
int precision = checkStateNotNull(logicalType.getArgument());
228+
if (precision == Timestamp.MICROS.getArgument()) {
229+
type = Types.TimestampType.withZone();
230+
} else {
231+
throw new UnsupportedOperationException(
232+
"Unsupported Timestamp precision: " + precision);
233+
}
234+
} else {
235+
throw new RuntimeException("Unsupported Beam logical type " + logicalTypeIdentifier);
236+
}
202237
}
203238
return new TypeAndMaxId(--nestedFieldId, type);
204239
} else if (beamType.getTypeName().isCollectionType()) { // ARRAY or ITERABLE
@@ -613,21 +648,28 @@ private static Object getLogicalTypeValue(Object icebergValue, Schema.FieldType
613648
return LocalTime.parse(strValue);
614649
} else if (type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
615650
return LocalDateTime.parse(strValue);
651+
} else if (type.isLogicalType(Timestamp.IDENTIFIER)) {
652+
return OffsetDateTime.parse(strValue).toInstant();
616653
}
617654
} else if (icebergValue instanceof Long) {
618655
if (type.isLogicalType(SqlTypes.TIME.getIdentifier())) {
619656
return DateTimeUtil.timeFromMicros((Long) icebergValue);
620657
} else if (type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
621658
return DateTimeUtil.timestampFromMicros((Long) icebergValue);
659+
} else if (type.isLogicalType(Timestamp.IDENTIFIER)) {
660+
// timestamptz stored as micros since epoch -> java.time.Instant (micros preserved).
661+
return DateTimeUtil.timestamptzFromMicros((Long) icebergValue).toInstant();
622662
}
623663
} else if (icebergValue instanceof Integer
624664
&& type.isLogicalType(SqlTypes.DATE.getIdentifier())) {
625665
return DateTimeUtil.dateFromDays((Integer) icebergValue);
626-
} else if (icebergValue instanceof OffsetDateTime
627-
&& type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
628-
return ((OffsetDateTime) icebergValue)
629-
.withOffsetSameInstant(ZoneOffset.UTC)
630-
.toLocalDateTime();
666+
} else if (icebergValue instanceof OffsetDateTime) {
667+
OffsetDateTime odt = (OffsetDateTime) icebergValue;
668+
if (type.isLogicalType(SqlTypes.DATETIME.getIdentifier())) {
669+
return odt.withOffsetSameInstant(ZoneOffset.UTC).toLocalDateTime();
670+
} else if (type.isLogicalType(Timestamp.IDENTIFIER)) {
671+
return odt.toInstant();
672+
}
631673
}
632674
// LocalDateTime, LocalDate, LocalTime
633675
return icebergValue;

‎sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IncrementalScanSource.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,9 @@ public PCollection<Row> expand(PBegin input) {
6868
.setCoder(KvCoder.of(ReadTaskDescriptor.getCoder(), ReadTask.getCoder()))
6969
.apply(Redistribute.arbitrarily())
7070
.apply("Read Rows From Tasks", ParDo.of(new ReadFromTasks(scanConfig)))
71-
.setRowSchema(IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema()));
71+
.setRowSchema(
72+
IcebergUtils.icebergSchemaToBeamSchema(
73+
scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion()));
7274
}
7375

7476
/** Continuously watches for new snapshots. */

‎sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,9 @@ public void process(
6969
return;
7070
}
7171
FileScanTask task = fileScanTasks.get((int) l);
72-
Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema());
72+
Schema beamSchema =
73+
IcebergUtils.icebergSchemaToBeamSchema(
74+
scanConfig.getProjectedSchema(), scanConfig.getUpdateCompatibilityVersion());
7375
try (CloseableIterable<Record> reader = ReadUtils.createReader(task, table, scanConfig)) {
7476

7577
for (Record record : reader) {

0 commit comments

Comments
 (0)