[FLINK-22378][table] Support TIMESTAMP as StreamRecord timestamp in toDataStream

twalthr · twalthr · commit d3978edc39bc · 2021-04-23T10:46:41.000+02:00
diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/connector/sink/DynamicTableSink.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/connector/sink/DynamicTableSink.java
@@ -20,7 +20,7 @@
 
 import org.apache.flink.annotation.PublicEvolving;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
-import org.apache.flink.table.api.TableSchema;
+import org.apache.flink.table.catalog.ResolvedSchema;
 import org.apache.flink.table.connector.ChangelogMode;
 import org.apache.flink.table.connector.ParallelismProvider;
 import org.apache.flink.table.connector.RuntimeConverter;
@@ -96,8 +96,7 @@ public interface DynamicTableSink {
      * interfaces might be located in other Flink modules.
      *
      * <p>Independent of the provider interface, the table runtime expects that a sink
-     * implementation accepts internal data structures (see {@link
-     * org.apache.flink.table.data.RowData} for more information).
+     * implementation accepts internal data structures (see {@link RowData} for more information).
      *
      * <p>The given {@link Context} offers utilities by the planner for creating runtime
      * implementation with minimal dependencies to internal data structures.
@@ -146,7 +145,7 @@ interface Context {
          * Creates type information describing the internal data structures of the given {@link
          * DataType}.
          *
-         * @see TableSchema#toPhysicalRowDataType()
+         * @see ResolvedSchema#toPhysicalRowDataType()
          */
         <T> TypeInformation<T> createTypeInformation(DataType consumedDataType);
 
diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/data/TimestampData.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/data/TimestampData.java
@@ -23,11 +23,15 @@
 import org.apache.flink.table.types.logical.TimestampType;
 import org.apache.flink.util.Preconditions;
 
+import javax.annotation.Nullable;
+
 import java.sql.Timestamp;
 import java.time.Instant;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
 
 /**
  * An internal data structure representing data of {@link TimestampType} and {@link
@@ -88,7 +92,7 @@ public LocalDateTime toLocalDateTime() {
         return LocalDateTime.of(localDate, localTime);
     }
 
-    /** Converts this {@link TimestampData} object to a {@link Instant}. */
+    /** Converts this {@link TimestampData} object to an {@link Instant}. */
     public Instant toInstant() {
         long epochSecond = millisecond / 1000;
         int milliOfSecond = (int) (millisecond % 1000);
@@ -100,6 +104,29 @@ public Instant toInstant() {
         return Instant.ofEpochSecond(epochSecond, nanoAdjustment);
     }
 
+    /**
+     * Converts this {@link TimestampData} object to milliseconds adjusted by the given time zone to
+     * epoch millis.
+     *
+     * @param zoneId zone for offset calculation or null if {@link TimestampData} is already in UTC
+     *     (i.e. for {@link LocalZonedTimestampType})
+     */
+    public long toZonedMillisecond(@Nullable ZoneId zoneId) {
+        if (zoneId == null) {
+            return millisecond;
+        }
+        int date = (int) (millisecond / MILLIS_PER_DAY);
+        int time = (int) (millisecond % MILLIS_PER_DAY);
+        if (time < 0) {
+            --date;
+            time += MILLIS_PER_DAY;
+        }
+        long nanoOfDay = time * 1_000_000L + nanoOfMillisecond;
+        LocalDate localDate = LocalDate.ofEpochDay(date);
+        LocalTime localTime = LocalTime.ofNanoOfDay(nanoOfDay);
+        return ZonedDateTime.of(localDate, localTime, zoneId).toInstant().toEpochMilli();
+    }
+
     @Override
     public int compareTo(TimestampData that) {
         int cmp = Long.compare(this.millisecond, that.millisecond);
diff --git a/flink-table/flink-table-planner-blink/src/main/java/org/apache/flink/table/planner/connectors/DynamicSinkUtils.java b/flink-table/flink-table-planner-blink/src/main/java/org/apache/flink/table/planner/connectors/DynamicSinkUtils.java
@@ -61,6 +61,7 @@
 
 import javax.annotation.Nullable;
 
+import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -133,13 +134,15 @@ public static RelNode convertExternalToRel(
             FlinkRelBuilder relBuilder,
             RelNode input,
             ExternalModifyOperation externalModifyOperation) {
+        final ZoneId localTimeZone = unwrapContext(relBuilder).getTableConfig().getLocalTimeZone();
         final ResolvedSchema schema = externalModifyOperation.getResolvedSchema();
         final CatalogTable unresolvedTable = new InlineCatalogTable(schema);
         final ResolvedCatalogTable catalogTable = new ResolvedCatalogTable(unresolvedTable, schema);
         final DynamicTableSink tableSink =
                 new ExternalDynamicSink(
                         externalModifyOperation.getChangelogMode().orElse(null),
-                        externalModifyOperation.getPhysicalDataType());
+                        externalModifyOperation.getPhysicalDataType(),
+                        localTimeZone);
         return convertSinkToRel(
                 relBuilder,
                 input,
diff --git a/flink-table/flink-table-planner-blink/src/main/java/org/apache/flink/table/planner/connectors/ExternalDynamicSink.java b/flink-table/flink-table-planner-blink/src/main/java/org/apache/flink/table/planner/connectors/ExternalDynamicSink.java
@@ -31,14 +31,18 @@
 import org.apache.flink.table.runtime.typeutils.ExternalTypeInfo;
 import org.apache.flink.table.types.DataType;
 import org.apache.flink.table.types.logical.LogicalType;
+import org.apache.flink.table.types.logical.LogicalTypeRoot;
 import org.apache.flink.table.types.logical.utils.LogicalTypeChecks;
 
 import javax.annotation.Nullable;
 
+import java.time.ZoneId;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
+import static org.apache.flink.table.types.logical.utils.LogicalTypeChecks.hasRoot;
+
 /** Table sink for connecting to the external {@link DataStream} API. */
 @Internal
 final class ExternalDynamicSink implements DynamicTableSink, SupportsWritingMetadata {
@@ -51,13 +55,17 @@ final class ExternalDynamicSink implements DynamicTableSink, SupportsWritingMeta
 
     private final DataType physicalDataType;
 
+    private final ZoneId rowtimeShift;
+
     // mutable attributes
 
     private boolean consumeRowtimeMetadata;
 
-    ExternalDynamicSink(@Nullable ChangelogMode changelogMode, DataType physicalDataType) {
+    ExternalDynamicSink(
+            @Nullable ChangelogMode changelogMode, DataType physicalDataType, ZoneId rowtimeShift) {
         this.changelogMode = changelogMode;
         this.physicalDataType = physicalDataType;
+        this.rowtimeShift = rowtimeShift;
     }
 
     @Override
@@ -86,19 +94,33 @@ public SinkRuntimeProvider getSinkRuntimeProvider(Context context) {
                         atomicFieldGetter = RowData.createFieldGetter(physicalType, 0);
                     }
 
+                    final int rowtimeIndex = transformationContext.getRowtimeIndex();
+
                     return new OneInputTransformation<>(
                             input,
                             generateOperatorName(),
                             new OutputConversionOperator(
                                     atomicFieldGetter,
                                     physicalConverter,
-                                    transformationContext.getRowtimeIndex(),
+                                    rowtimeIndex,
+                                    generateRowtimeShift(rowtimeIndex),
                                     consumeRowtimeMetadata),
                             ExternalTypeInfo.of(physicalDataType),
                             input.getParallelism());
                 };
     }
 
+    private @Nullable ZoneId generateRowtimeShift(int rowtimeIndex) {
+        if (rowtimeIndex != -1) {
+            final LogicalType rowtimeType =
+                    physicalDataType.getChildren().get(rowtimeIndex).getLogicalType();
+            if (hasRoot(rowtimeType, LogicalTypeRoot.TIMESTAMP_WITHOUT_TIME_ZONE)) {
+                return rowtimeShift;
+            }
+        }
+        return null;
+    }
+
     private String generateOperatorName() {
         return String.format(
                 "TableToDataSteam(type=%s, rowtime=%s)",
@@ -107,7 +129,7 @@ private String generateOperatorName() {
 
     @Override
     public DynamicTableSink copy() {
-        return new ExternalDynamicSink(changelogMode, physicalDataType);
+        return new ExternalDynamicSink(changelogMode, physicalDataType, rowtimeShift);
     }
 
     @Override
diff --git a/flink-table/flink-table-planner-blink/src/test/java/org/apache/flink/table/planner/runtime/stream/sql/DataStreamJavaITCase.java b/flink-table/flink-table-planner-blink/src/test/java/org/apache/flink/table/planner/runtime/stream/sql/DataStreamJavaITCase.java
@@ -28,6 +28,7 @@
 import org.apache.flink.api.java.typeutils.TupleTypeInfo;
 import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.ProcessFunction;
 import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
 import org.apache.flink.streaming.api.windowing.time.Time;
 import org.apache.flink.table.api.DataTypes;
@@ -48,6 +49,7 @@
 import org.apache.flink.types.RowKind;
 import org.apache.flink.util.CloseableIterator;
 import org.apache.flink.util.CollectionUtil;
+import org.apache.flink.util.Collector;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -57,6 +59,8 @@
 import org.junit.runners.Parameterized.Parameters;
 
 import java.time.DayOfWeek;
+import java.time.LocalDateTime;
+import java.time.ZoneId;
 import java.time.ZoneOffset;
 import java.util.Arrays;
 import java.util.Collections;
@@ -433,6 +437,51 @@ public void testFromAndToChangelogStreamUpsert() throws Exception {
                 getOutput(inputOrOutput));
     }
 
+    @Test
+    public void testToDataStreamCustomEventTime() throws Exception {
+        final ZoneId originalZone = tableEnv.getConfig().getLocalTimeZone();
+
+        final ZoneId customZone = ZoneId.of("Europe/Berlin");
+
+        tableEnv.getConfig().setLocalTimeZone(customZone);
+
+        final LocalDateTime localDateTime1 = LocalDateTime.parse("1970-01-01T00:00:00.000");
+        final LocalDateTime localDateTime2 = LocalDateTime.parse("1970-01-01T01:00:00.000");
+
+        final DataStream<Tuple2<LocalDateTime, String>> dataStream =
+                env.fromElements(
+                        new Tuple2<>(localDateTime1, "alice"), new Tuple2<>(localDateTime2, "bob"));
+
+        final Table table =
+                tableEnv.fromDataStream(
+                        dataStream,
+                        Schema.newBuilder()
+                                .column("f0", DataTypes.TIMESTAMP(3))
+                                .column("f1", DataTypes.STRING())
+                                .watermark("f0", "f0 - INTERVAL '0.01' SECOND")
+                                .build());
+
+        final DataStream<Long> rowtimeStream =
+                tableEnv.toDataStream(table)
+                        .process(
+                                new ProcessFunction<Row, Long>() {
+                                    @Override
+                                    public void processElement(
+                                            Row value, Context ctx, Collector<Long> out) {
+                                        out.collect(ctx.timestamp());
+                                    }
+                                });
+
+        // TODO timestamp and watermark are off by the time zone at the moment see FLINK-22356
+
+        testResult(
+                rowtimeStream,
+                localDateTime1.atZone(customZone).toInstant().toEpochMilli(),
+                localDateTime2.atZone(customZone).toInstant().toEpochMilli());
+
+        tableEnv.getConfig().setLocalTimeZone(originalZone);
+    }
+
     // --------------------------------------------------------------------------------------------
     // Helper methods
     // --------------------------------------------------------------------------------------------
diff --git a/flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/sink/OutputConversionOperator.java b/flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/sink/OutputConversionOperator.java
@@ -29,6 +29,8 @@
 
 import javax.annotation.Nullable;
 
+import java.time.ZoneId;
+
 /** Operator that converts to external data structures and unwraps atomic records if necessary. */
 @Internal
 public class OutputConversionOperator extends TableStreamOperator<Object>
@@ -40,6 +42,8 @@ public class OutputConversionOperator extends TableStreamOperator<Object>
 
     private final int rowtimeIndex;
 
+    private final @Nullable ZoneId rowtimeShift;
+
     private final boolean consumeRowtimeMetadata;
 
     private transient StreamRecord<Object> outRecord;
@@ -48,10 +52,12 @@ public OutputConversionOperator(
             @Nullable RowData.FieldGetter atomicFieldGetter,
             DataStructureConverter converter,
             int rowtimeIndex,
+            @Nullable ZoneId rowtimeShift,
             boolean consumeRowtimeMetadata) {
         this.atomicFieldGetter = atomicFieldGetter;
         this.converter = converter;
         this.rowtimeIndex = rowtimeIndex;
+        this.rowtimeShift = rowtimeShift;
         this.consumeRowtimeMetadata = consumeRowtimeMetadata;
     }
 
@@ -70,10 +76,13 @@ public void processElement(StreamRecord<RowData> element) throws Exception {
         final RowData rowData = element.getValue();
 
         if (consumeRowtimeMetadata) {
+            // timestamp is TIMESTAMP_LTZ
             final long rowtime = rowData.getTimestamp(rowData.getArity() - 1, 3).getMillisecond();
             outRecord.setTimestamp(rowtime);
         } else if (rowtimeIndex != -1) {
-            final long rowtime = rowData.getTimestamp(rowtimeIndex, 3).getMillisecond();
+            // timestamp might be TIMESTAMP or TIMESTAMP_LTZ
+            final long rowtime =
+                    rowData.getTimestamp(rowtimeIndex, 3).toZonedMillisecond(rowtimeShift);
             outRecord.setTimestamp(rowtime);
         }
 
diff --git a/flink-table/flink-table-runtime-blink/src/test/java/org/apache/flink/table/data/TimestampDataTest.java b/flink-table/flink-table-runtime-blink/src/test/java/org/apache/flink/table/data/TimestampDataTest.java
