Hive: Fix predicate pushdown for Timestamp.withZone()

hive3/src/main/java/org/apache/iceberg/mr/hive/serde/objectinspector/IcebergTimestampObjectInspectorHive3.java

@@ -55,9 +55,7 @@ public Timestamp getPrimitiveJavaObject(Object o) {
       return null;
     }
     LocalDateTime time = (LocalDateTime) o;
-    Timestamp timestamp = Timestamp.ofEpochMilli(time.toInstant(ZoneOffset.UTC).toEpochMilli());
-    timestamp.setNanos(time.getNano());
-    return timestamp;
+    return Timestamp.ofEpochMilli(time.toInstant(ZoneOffset.UTC).toEpochMilli(), time.getNano());
   }
 
   @Override

mr/src/main/java/org/apache/iceberg/mr/hive/serde/objectinspector/IcebergTimestampWithZoneObjectInspector.java

@@ -42,12 +42,13 @@ private IcebergTimestampWithZoneObjectInspector() {
 
   @Override
   public OffsetDateTime convert(Object o) {
-    return o == null ? null : OffsetDateTime.ofInstant(((Timestamp) o).toInstant(), ZoneOffset.UTC);
+    return o == null ? null : OffsetDateTime.of(((Timestamp) o).toLocalDateTime(), ZoneOffset.UTC);
   }
 
   @Override
   public Timestamp getPrimitiveJavaObject(Object o) {
-    return o == null ? null : Timestamp.from(((OffsetDateTime) o).toInstant());
+    return o == null ? null :
+        Timestamp.valueOf(((OffsetDateTime) o).withOffsetSameInstant(ZoneOffset.UTC).toLocalDateTime());
   }
 
   @Override

mr/src/test/java/org/apache/iceberg/mr/hive/HiveIcebergTestUtils.java

@@ -32,6 +32,10 @@
 import java.time.LocalTime;
 import java.time.OffsetDateTime;
 import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
+import java.time.format.DateTimeFormatter;
+import java.time.format.DateTimeFormatterBuilder;
+import java.time.temporal.ChronoField;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Comparator;
@@ -59,7 +63,9 @@
 import org.apache.iceberg.data.GenericRecord;
 import org.apache.iceberg.data.IcebergGenerics;
 import org.apache.iceberg.data.Record;
+import org.apache.iceberg.hive.MetastoreUtil;
 import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.ByteBuffers;
 import org.junit.Assert;
@@ -107,6 +113,13 @@ public class HiveIcebergTestUtils {
               PrimitiveObjectInspectorFactory.writableStringObjectInspector
           ));
 
+  public static final DateTimeFormatter TIMESTAMP_WITH_TZ_FORMATTER = new DateTimeFormatterBuilder()
+      .appendPattern("yyyy-MM-dd HH:mm:ss")
+      .appendFraction(ChronoField.MICRO_OF_SECOND, 0, 9, true)
+      .appendLiteral(' ')
+      .appendZoneOrOffsetId()
+      .toFormatter();
+
   private HiveIcebergTestUtils() {
     // Empty constructor for the utility class
   }
@@ -210,14 +223,14 @@ public static void assertEquals(Record expected, Record actual) {
   }
 
   /**
-   * Validates whether the table contains the expected records. The results should be sorted by a unique key so we do
-   * not end up with flaky tests.
+   * Validates whether the table contains the expected records reading the data through the Iceberg API.
+   * The results should be sorted by a unique key so we do not end up with flaky tests.
    * @param table The table we should read the records from
    * @param expected The expected list of Records
    * @param sortBy The column position by which we will sort
    * @throws IOException Exceptions when reading the table data
    */
-  public static void validateData(Table table, List<Record> expected, int sortBy) throws IOException {
+  public static void validateDataWithIceberg(Table table, List<Record> expected, int sortBy) throws IOException {
     // Refresh the table, so we get the new data as well
     table.refresh();
     List<Record> records = new ArrayList<>(expected.size());
@@ -265,4 +278,58 @@ public static void validateFiles(Table table, Configuration conf, JobID jobId, i
     Assert.assertEquals(dataFileNum, dataFiles.size());
     Assert.assertFalse(new File(HiveIcebergOutputCommitter.generateJobLocation(conf, jobId)).exists());
   }
+
+  /**
+   * Simplified implementation for checking that the returned results from a SELECT statement are the same than the
+   * inserted values. We expect that the values are inserted using {@link #getStringValueForInsert(Object, Type)}.
+   * <p>
+   * For the full implementation we might want to add sorting so the check could work for every table/query.
+   * @param shell The shell used for executing the query
+   * @param tableName The name of the table to query
+   * @param expected The records we inserted
+   */
+  public static void validateDataWithSql(TestHiveShell shell, String tableName, List<Record> expected) {
+    List<Object[]> actual = shell.executeStatement("SELECT * from " + tableName);
+
+    for (int rowId = 0; rowId < expected.size(); ++rowId) {
+      Record record = expected.get(rowId);
+      Object[] row = actual.get(rowId);
+      Assert.assertEquals(record.size(), row.length);
+      for (int fieldId = 0; fieldId < record.size(); ++fieldId) {
+        Types.NestedField field = record.struct().fields().get(fieldId);
+        String inserted = getStringValueForInsert(record.getField(field.name()), field.type())
+            // If there are enclosing quotes then remove them
+            .replaceAll("'(.*)'", "$1");
+        String returned = row[fieldId].toString();
+        if (field.type().equals(Types.TimestampType.withZone()) && MetastoreUtil.hive3PresentOnClasspath()) {
+          Timestamp timestamp = Timestamp.from(ZonedDateTime.parse(returned, TIMESTAMP_WITH_TZ_FORMATTER).toInstant());
+          returned = timestamp.toString();
+        }
+        Assert.assertEquals(inserted, returned);
+      }
+    }
+  }
+
+  public static String getStringValueForInsert(Object value, Type type) {
+    String template = "\'%s\'";
+    if (type.equals(Types.TimestampType.withoutZone())) {
+      return String.format(template, Timestamp.valueOf((LocalDateTime) value).toString());
+    } else if (type.equals(Types.TimestampType.withZone())) {
+      Timestamp timestamp;
+      // Hive2 stores Timestamps with local TZ, Hive3 stores Timestamps in UTC so we have to insert different timestamp
+      // to get the same expected values in the Iceberg rows. The Hive query should return the same values as inserted
+      // in both cases
+      if (MetastoreUtil.hive3PresentOnClasspath()) {
+        timestamp = Timestamp.from(((OffsetDateTime) value).toInstant());
+      } else {
+        timestamp = Timestamp.valueOf(((OffsetDateTime) value).withOffsetSameInstant(ZoneOffset.UTC).toLocalDateTime());
+      }
+      return String.format(template, timestamp.toString());
+    } else if (type.equals(Types.BooleanType.get())) {
+      // in hive2 boolean type values must not be surrounded in apostrophes. Otherwise the value is translated to true.
+      return value.toString();
+    } else {
+      return String.format(template, value.toString());
+    }
+  }
 }

mr/src/test/java/org/apache/iceberg/mr/hive/TestHiveIcebergOutputCommitter.java

@@ -111,7 +111,7 @@ public void testSuccessfulUnpartitionedWrite() throws IOException {
     committer.commitJob(new JobContextImpl(conf, JOB_ID));
 
     HiveIcebergTestUtils.validateFiles(table, conf, JOB_ID, 1);
-    HiveIcebergTestUtils.validateData(table, expected, 0);
+    HiveIcebergTestUtils.validateDataWithIceberg(table, expected, 0);
   }
 
   @Test
@@ -123,7 +123,7 @@ public void testSuccessfulPartitionedWrite() throws IOException {
     committer.commitJob(new JobContextImpl(conf, JOB_ID));
 
     HiveIcebergTestUtils.validateFiles(table, conf, JOB_ID, 3);
-    HiveIcebergTestUtils.validateData(table, expected, 0);
+    HiveIcebergTestUtils.validateDataWithIceberg(table, expected, 0);
   }
 
   @Test
@@ -135,7 +135,7 @@ public void testSuccessfulMultipleTasksUnpartitionedWrite() throws IOException {
     committer.commitJob(new JobContextImpl(conf, JOB_ID));
 
     HiveIcebergTestUtils.validateFiles(table, conf, JOB_ID, 2);
-    HiveIcebergTestUtils.validateData(table, expected, 0);
+    HiveIcebergTestUtils.validateDataWithIceberg(table, expected, 0);
   }
 
   @Test
@@ -147,7 +147,7 @@ public void testSuccessfulMultipleTasksPartitionedWrite() throws IOException {
     committer.commitJob(new JobContextImpl(conf, JOB_ID));
 
     HiveIcebergTestUtils.validateFiles(table, conf, JOB_ID, 6);
-    HiveIcebergTestUtils.validateData(table, expected, 0);
+    HiveIcebergTestUtils.validateDataWithIceberg(table, expected, 0);
   }
 
   @Test
@@ -159,19 +159,19 @@ public void testRetryTask() throws IOException {
     // Write records and abort the tasks
     writeRecords(2, 0, false, true, conf);
     HiveIcebergTestUtils.validateFiles(table, conf, JOB_ID, 0);
-    HiveIcebergTestUtils.validateData(table, Collections.emptyList(), 0);
+    HiveIcebergTestUtils.validateDataWithIceberg(table, Collections.emptyList(), 0);
 
     // Write records but do not abort the tasks
     // The data files remain since we can not identify them but should not be read
     writeRecords(2, 1, false, false, conf);
     HiveIcebergTestUtils.validateFiles(table, conf, JOB_ID, 2);
-    HiveIcebergTestUtils.validateData(table, Collections.emptyList(), 0);
+    HiveIcebergTestUtils.validateDataWithIceberg(table, Collections.emptyList(), 0);
 
     // Write and commit the records
     List<Record> expected = writeRecords(2, 2, true, false, conf);
     committer.commitJob(new JobContextImpl(conf, JOB_ID));
     HiveIcebergTestUtils.validateFiles(table, conf, JOB_ID, 4);
-    HiveIcebergTestUtils.validateData(table, expected, 0);
+    HiveIcebergTestUtils.validateDataWithIceberg(table, expected, 0);
   }
 
   @Test
@@ -183,7 +183,7 @@ public void testAbortJob() throws IOException {
     committer.abortJob(new JobContextImpl(conf, JOB_ID), JobStatus.State.FAILED);
 
     HiveIcebergTestUtils.validateFiles(table, conf, JOB_ID, 0);
-    HiveIcebergTestUtils.validateData(table, Collections.emptyList(), 0);
+    HiveIcebergTestUtils.validateDataWithIceberg(table, Collections.emptyList(), 0);
   }
 
   @Test

mr/src/test/java/org/apache/iceberg/mr/hive/TestHiveIcebergStorageHandlerTimezone.java

@@ -23,16 +23,23 @@
 import java.text.DateFormat;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
+import java.time.OffsetDateTime;
+import java.time.ZoneId;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
 import java.util.Collection;
 import java.util.List;
 import java.util.Optional;
 import java.util.TimeZone;
 import org.apache.hadoop.hive.serde2.io.DateWritable;
 import org.apache.hadoop.hive.serde2.io.TimestampWritable;
+import org.apache.hadoop.hive.serde2.typeinfo.PrimitiveTypeInfo;
+import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory;
 import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.common.DynFields;
 import org.apache.iceberg.data.Record;
+import org.apache.iceberg.hive.MetastoreUtil;
 import org.apache.iceberg.mr.TestHelper;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
 import org.apache.iceberg.types.Types;
@@ -47,26 +54,38 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 
+import static org.apache.iceberg.mr.hive.HiveIcebergTestUtils.TIMESTAMP_WITH_TZ_FORMATTER;
 import static org.apache.iceberg.types.Types.NestedField.optional;
 import static org.junit.runners.Parameterized.Parameter;
 import static org.junit.runners.Parameterized.Parameters;
 
 @RunWith(Parameterized.class)
 public class TestHiveIcebergStorageHandlerTimezone {
-  private static final Optional<ThreadLocal<DateFormat>> dateFormat =
+  private static final Optional<ThreadLocal<DateFormat>> DATE_FORMAT =
       Optional.ofNullable((ThreadLocal<DateFormat>) DynFields.builder()
           .hiddenImpl(TimestampWritable.class, "threadLocalDateFormat")
           .defaultAlwaysNull()
           .buildStatic()
           .get());
 
-  private static final Optional<ThreadLocal<TimeZone>> localTimeZone =
+  private static final Optional<ThreadLocal<TimeZone>> LOCAL_TIMEZONE =
       Optional.ofNullable((ThreadLocal<TimeZone>) DynFields.builder()
           .hiddenImpl(DateWritable.class, "LOCAL_TIMEZONE")
           .defaultAlwaysNull()
           .buildStatic()
           .get());
 
+  private static final PrimitiveTypeInfo TIMESTAMP_LOCAL_TZ_TYPE_INFO =
+      (PrimitiveTypeInfo) DynFields.builder().hiddenImpl(TypeInfoFactory.class, "timestampLocalTZTypeInfo")
+          .defaultAlwaysNull()
+          .buildStatic().get();
+
+  private static final Optional<DynFields.BoundField<ZoneId>> ZONE_ID =
+      Optional.ofNullable(TIMESTAMP_LOCAL_TZ_TYPE_INFO == null ? null : DynFields.builder()
+          .hiddenImpl("org.apache.hadoop.hive.serde2.typeinfo.TimestampLocalTZTypeInfo", "timeZone")
+          .defaultAlwaysNull()
+          .build(TIMESTAMP_LOCAL_TZ_TYPE_INFO));
+
   @Parameters(name = "timezone={0}")
   public static Collection<Object[]> parameters() {
     return ImmutableList.of(
@@ -102,8 +121,9 @@ public void before() throws IOException {
 
     // Magic to clean cached date format and local timezone for Hive where the default timezone is used/stored in the
     // cached object
-    dateFormat.ifPresent(ThreadLocal::remove);
-    localTimeZone.ifPresent(ThreadLocal::remove);
+    DATE_FORMAT.ifPresent(ThreadLocal::remove);
+    LOCAL_TIMEZONE.ifPresent(ThreadLocal::remove);
+    ZONE_ID.ifPresent(field -> field.set(TimeZone.getDefault().toZoneId()));
 
     this.testTables = HiveIcebergStorageHandlerTestUtils.testTables(shell, TestTables.TestTableType.HIVE_CATALOG, temp);
     // Uses spark as an engine so we can detect if we unintentionally try to use any execution engines
@@ -169,4 +189,70 @@ public void testTimestampQuery() throws IOException {
     result = shell.executeStatement("SELECT * FROM ts_test WHERE d_ts='2017-01-01 22:30:57.3'");
     Assert.assertEquals(0, result.size());
   }
+
+  @Test
+  public void testTimestampTzQuery() throws IOException {
+    Schema timestampSchema = new Schema(optional(1, "d_ts", Types.TimestampType.withZone()));
+
+    if (MetastoreUtil.hive3PresentOnClasspath()) {
+      LocalDateTime localDateTime = LocalDateTime.of(2019, 1, 22, 9, 44, 54, 100000000);
+      ZoneOffset zoneOffSet = TimeZone.getDefault().toZoneId().getRules().getOffset(localDateTime);
+
+      OffsetDateTime offsetDateTime1 = localDateTime.atOffset(zoneOffSet);
+
+      localDateTime = LocalDateTime.of(2019, 2, 22, 9, 44, 54, 200000000);
+      zoneOffSet = TimeZone.getDefault().toZoneId().getRules().getOffset(localDateTime);
+
+      OffsetDateTime offsetDateTime2 = localDateTime.atOffset(zoneOffSet);
+
+      List<Record> records = TestHelper.RecordsBuilder.newInstance(timestampSchema)
+          .add(offsetDateTime1)
+          .add(offsetDateTime2)
+          .build();
+
+      testTables.createTable(shell, "ts_test", timestampSchema, FileFormat.PARQUET, records);
+
+      List<Object[]> result = shell.executeStatement("SELECT d_ts FROM ts_test WHERE d_ts='2019-02-22 09:44:54.2'");
+      Assert.assertEquals(1, result.size());
+      Assert.assertEquals(offsetDateTime2.toInstant(),
+          ZonedDateTime.parse(result.get(0)[0].toString(), TIMESTAMP_WITH_TZ_FORMATTER).toInstant());
+
+      // Skip testing `in` as I was not able come up with the constants that would work (maybe some Hive bug?)
+
+      result = shell.executeStatement("SELECT d_ts FROM ts_test WHERE d_ts < '2019-02-22 09:44:54.2'");
+      Assert.assertEquals(1, result.size());
+      Assert.assertEquals(offsetDateTime1.toInstant(),
+          ZonedDateTime.parse(result.get(0)[0].toString(), TIMESTAMP_WITH_TZ_FORMATTER).toInstant());
+
+      result = shell.executeStatement("SELECT d_ts FROM ts_test WHERE d_ts='2017-01-01 22:30:57.3'");
+      Assert.assertEquals(0, result.size());
+    } else {
+      List<Record> records = TestHelper.RecordsBuilder.newInstance(timestampSchema)
+          .add(OffsetDateTime.of(
+              LocalDateTime.of(2019, 1, 22, 9, 44, 54, 100000000),
+              ZoneOffset.UTC))
+          .add(OffsetDateTime.of(
+              LocalDateTime.of(2019, 2, 22, 9, 44, 54, 200000000),
+              ZoneOffset.UTC))
+          .build();
+
+      testTables.createTable(shell, "ts_test", timestampSchema, FileFormat.PARQUET, records);
+
+      List<Object[]> result = shell.executeStatement("SELECT d_ts FROM ts_test WHERE d_ts='2019-02-22 09:44:54.2'");
+      Assert.assertEquals(1, result.size());
+      Assert.assertEquals("2019-02-22 09:44:54.2", result.get(0)[0]);
+
+      result = shell.executeStatement(
+          "SELECT d_ts FROM ts_test WHERE d_ts in ('2017-01-01 22:30:57.1', '2019-02-22 09:44:54.2')");
+      Assert.assertEquals(1, result.size());
+      Assert.assertEquals("2019-02-22 09:44:54.2", result.get(0)[0]);
+
+      result = shell.executeStatement("SELECT d_ts FROM ts_test WHERE d_ts < '2019-02-22 09:44:54.2'");
+      Assert.assertEquals(1, result.size());
+      Assert.assertEquals("2019-01-22 09:44:54.1", result.get(0)[0]);
+
+      result = shell.executeStatement("SELECT d_ts FROM ts_test WHERE d_ts='2017-01-01 22:30:57.3'");
+      Assert.assertEquals(0, result.size());
+    }
+  }
 }