Add support for reading/writing timestamps without timezone. (#2757)

Previously Spark could not handle Iceberg tables which contained Timestamp.withoutTimeZone. New parameters are introduced to allow Timestamp without TimeZone to be treated as Timestamp with Timezone.  

Co-authored-by: bkahloon <kahlonbakht@gmail.com>
Co-authored-by: shardulm94

Author: sshkvar

spark/src/main/java/org/apache/iceberg/spark/PruneColumnsWithoutReordering.java

@@ -202,11 +202,6 @@ public Type primitive(Type.PrimitiveType primitive) {
             "Cannot project decimal with incompatible precision: %s < %s",
             requestedDecimal.precision(), decimal.precision());
         break;
-      case TIMESTAMP:
-        Types.TimestampType timestamp = (Types.TimestampType) primitive;
-        Preconditions.checkArgument(timestamp.shouldAdjustToUTC(),
-            "Cannot project timestamp (without time zone) as timestamptz (with time zone)");
-        break;
       default:
     }
 

spark/src/main/java/org/apache/iceberg/spark/SparkFixupTimestampType.java

@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iceberg.spark;
+
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.types.FixupTypes;
+import org.apache.iceberg.types.Type;
+import org.apache.iceberg.types.TypeUtil;
+import org.apache.iceberg.types.Types;
+
+/**
+ * By default Spark type {@link org.apache.iceberg.types.Types.TimestampType} should be converted to
+ * {@link Types.TimestampType#withZone()} iceberg type. But we also can convert
+ * {@link org.apache.iceberg.types.Types.TimestampType} to {@link Types.TimestampType#withoutZone()} iceberg type
+ * by setting {@link SparkUtil#USE_TIMESTAMP_WITHOUT_TIME_ZONE_IN_NEW_TABLES} to 'true'
+ */
+class SparkFixupTimestampType extends FixupTypes {
+
+  private SparkFixupTimestampType(Schema referenceSchema) {
+    super(referenceSchema);
+  }
+
+  static Schema fixup(Schema schema) {
+    return new Schema(TypeUtil.visit(schema,
+        new SparkFixupTimestampType(schema)).asStructType().fields());
+  }
+
+  @Override
+  public Type primitive(Type.PrimitiveType primitive) {
+    if (primitive.typeId() == Type.TypeID.TIMESTAMP) {
+      return Types.TimestampType.withoutZone();
+    }
+    return primitive;
+  }
+
+  @Override
+  protected boolean fixupPrimitive(Type.PrimitiveType type, Type source) {
+    return Type.TypeID.TIMESTAMP.equals(type.typeId());
+  }
+}

spark/src/main/java/org/apache/iceberg/spark/SparkFixupTypes.java

@@ -52,6 +52,11 @@ protected boolean fixupPrimitive(Type.PrimitiveType type, Type source) {
           return true;
         }
         break;
+      case TIMESTAMP:
+        if (source.typeId() == Type.TypeID.TIMESTAMP) {
+          return true;
+        }
+        break;
       default:
     }
     return false;

spark/src/main/java/org/apache/iceberg/spark/SparkSchemaUtil.java

@@ -122,8 +122,31 @@ public static DataType convert(Type type) {
    * @throws IllegalArgumentException if the type cannot be converted
    */
   public static Schema convert(StructType sparkType) {
+    return convert(sparkType, false);
+  }
+
+  /**
+   * Convert a Spark {@link StructType struct} to a {@link Schema} with new field ids.
+   * <p>
+   * This conversion assigns fresh ids.
+   * <p>
+   * Some data types are represented as the same Spark type. These are converted to a default type.
+   * <p>
+   * To convert using a reference schema for field ids and ambiguous types, use
+   * {@link #convert(Schema, StructType)}.
+   *
+   * @param sparkType a Spark StructType
+   * @param useTimestampWithoutZone boolean flag indicates that timestamp should be stored without timezone
+   * @return the equivalent Schema
+   * @throws IllegalArgumentException if the type cannot be converted
+   */
+  public static Schema convert(StructType sparkType, boolean useTimestampWithoutZone) {
     Type converted = SparkTypeVisitor.visit(sparkType, new SparkTypeToType(sparkType));
-    return new Schema(converted.asNestedType().asStructType().fields());
+    Schema schema = new Schema(converted.asNestedType().asStructType().fields());
+    if (useTimestampWithoutZone) {
+      schema = SparkFixupTimestampType.fixup(schema);
+    }
+    return schema;
   }
 
   /**

spark/src/main/java/org/apache/iceberg/spark/SparkUtil.java

@@ -20,21 +20,38 @@
 package org.apache.iceberg.spark;
 
 import java.util.List;
+import java.util.Map;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.apache.iceberg.PartitionField;
 import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.hadoop.HadoopConfigurable;
 import org.apache.iceberg.io.FileIO;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.transforms.Transform;
 import org.apache.iceberg.transforms.UnknownTransform;
+import org.apache.iceberg.types.TypeUtil;
+import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.Pair;
+import org.apache.spark.sql.RuntimeConfig;
 import org.apache.spark.util.SerializableConfiguration;
 
 public class SparkUtil {
+
+  public static final String HANDLE_TIMESTAMP_WITHOUT_TIMEZONE =
+          "spark.sql.iceberg.handle-timestamp-without-timezone";
+  public static final String TIMESTAMP_WITHOUT_TIMEZONE_ERROR = String.format("Cannot handle timestamp without" +
+          " timezone fields in Spark. Spark does not natively support this type but if you would like to handle all" +
+          " timestamps as timestamp with timezone set '%s' to true. This will not change the underlying values stored" +
+          " but will change their displayed values in Spark. For more information please see" +
+          " https://docs.databricks.com/spark/latest/dataframes-datasets/dates-timestamps.html#ansi-sql-and" +
+          "-spark-sql-timestamps", HANDLE_TIMESTAMP_WITHOUT_TIMEZONE);
+  public static final String USE_TIMESTAMP_WITHOUT_TIME_ZONE_IN_NEW_TABLES =
+          "spark.sql.iceberg.use-timestamp-without-timezone-in-new-tables";
+
   private SparkUtil() {
   }
 
@@ -99,4 +116,58 @@ public static <C, T> Pair<C, T> catalogAndIdentifier(List<String> nameParts,
       }
     }
   }
+
+  /**
+   * Responsible for checking if the table schema has a timestamp without timezone column
+   * @param schema table schema to check if it contains a timestamp without timezone column
+   * @return boolean indicating if the schema passed in has a timestamp field without a timezone
+   */
+  public static boolean hasTimestampWithoutZone(Schema schema) {
+    return TypeUtil.find(schema, t -> Types.TimestampType.withoutZone().equals(t)) != null;
+  }
+
+  /**
+   * Allow reading/writing timestamp without time zone as timestamp with time zone. Generally,
+   * this is not safe as timestamp without time zone is supposed to represent wall clock time semantics,
+   * i.e. no matter the reader/writer timezone 3PM should always be read as 3PM,
+   * but timestamp with time zone represents instant semantics, i.e the timestamp
+   * is adjusted so that the corresponding time in the reader timezone is displayed.
+   * When set to false (default), we throw an exception at runtime
+   * "Spark does not support timestamp without time zone fields" if reading timestamp without time zone fields
+   *
+   * @param readerConfig table read options
+   * @param sessionConf spark session configurations
+   * @return boolean indicating if reading timestamps without timezone is allowed
+   */
+  public static boolean canHandleTimestampWithoutZone(Map<String, String> readerConfig, RuntimeConfig sessionConf) {
+    String readerOption = readerConfig.get(HANDLE_TIMESTAMP_WITHOUT_TIMEZONE);
+    if (readerOption != null) {
+      return Boolean.parseBoolean(readerOption);
+    }
+    String sessionConfValue = sessionConf.get(HANDLE_TIMESTAMP_WITHOUT_TIMEZONE, null);
+    if (sessionConfValue != null) {
+      return Boolean.parseBoolean(sessionConfValue);
+    }
+    return false;
+  }
+
+  /**
+   * Check whether the spark session config contains a {@link SparkUtil#USE_TIMESTAMP_WITHOUT_TIME_ZONE_IN_NEW_TABLES}
+   * property.
+   * Default value - false
+   * If true in new table all timestamp fields will be stored as {@link Types.TimestampType#withoutZone()},
+   * otherwise {@link Types.TimestampType#withZone()} will be used
+   *
+   * @param sessionConf a spark runtime config
+   * @return true if the session config has {@link SparkUtil#USE_TIMESTAMP_WITHOUT_TIME_ZONE_IN_NEW_TABLES} property
+   * and this property is set to true
+   */
+  public static boolean useTimestampWithoutZoneInNewTables(RuntimeConfig sessionConf) {
+    String sessionConfValue = sessionConf.get(USE_TIMESTAMP_WITHOUT_TIME_ZONE_IN_NEW_TABLES, null);
+    if (sessionConfValue != null) {
+      return Boolean.parseBoolean(sessionConfValue);
+    }
+    return false;
+  }
+
 }

spark/src/main/java/org/apache/iceberg/spark/TypeToSparkType.java

@@ -104,12 +104,7 @@ public DataType primitive(Type.PrimitiveType primitive) {
         throw new UnsupportedOperationException(
             "Spark does not support time fields");
       case TIMESTAMP:
-        Types.TimestampType timestamp = (Types.TimestampType) primitive;
-        if (timestamp.shouldAdjustToUTC()) {
-          return TimestampType$.MODULE$;
-        }
-        throw new UnsupportedOperationException(
-            "Spark does not support timestamp without time zone fields");
+        return TimestampType$.MODULE$;
       case STRING:
         return StringType$.MODULE$;
       case UUID:

spark/src/main/java/org/apache/iceberg/spark/data/SparkOrcReader.java

@@ -105,6 +105,7 @@ public OrcValueReader<?> primitive(Type.PrimitiveType iPrimitive, TypeDescriptio
         case DOUBLE:
           return OrcValueReaders.doubles();
         case TIMESTAMP_INSTANT:
+        case TIMESTAMP:
           return SparkOrcValueReaders.timestampTzs();
         case DECIMAL:
           return SparkOrcValueReaders.decimals(primitive.getPrecision(), primitive.getScale());

spark/src/main/java/org/apache/iceberg/spark/data/SparkOrcWriter.java

@@ -118,6 +118,7 @@ public SparkOrcValueWriter primitive(Type.PrimitiveType iPrimitive, TypeDescript
         case DECIMAL:
           return SparkOrcValueWriters.decimal(primitive.getPrecision(), primitive.getScale());
         case TIMESTAMP_INSTANT:
+        case TIMESTAMP:
           return SparkOrcValueWriters.timestampTz();
         default:
           throw new IllegalArgumentException("Unhandled type " + primitive);

spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkOrcReaders.java

@@ -128,6 +128,7 @@ public Converter primitive(Type.PrimitiveType iPrimitive, TypeDescription primit
           primitiveValueReader = OrcValueReaders.doubles();
           break;
         case TIMESTAMP_INSTANT:
+        case TIMESTAMP:
           primitiveValueReader = SparkOrcValueReaders.timestampTzs();
           break;
         case DECIMAL:

spark/src/test/java/org/apache/iceberg/spark/data/AvroDataTest.java

@@ -20,15 +20,20 @@
 package org.apache.iceberg.spark.data;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.iceberg.Schema;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.spark.SparkUtil;
 import org.apache.iceberg.types.TypeUtil;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.types.Types.ListType;
 import org.apache.iceberg.types.Types.LongType;
 import org.apache.iceberg.types.Types.MapType;
 import org.apache.iceberg.types.Types.StructType;
+import org.apache.spark.sql.internal.SQLConf;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
@@ -185,4 +190,51 @@ public void testMixedTypes() throws IOException {
 
     writeAndValidate(schema);
   }
+
+  @Test
+  public void testTimestampWithoutZone() throws IOException {
+    withSQLConf(ImmutableMap.of(SparkUtil.HANDLE_TIMESTAMP_WITHOUT_TIMEZONE, "true"), () -> {
+      Schema schema = TypeUtil.assignIncreasingFreshIds(new Schema(
+              required(0, "id", LongType.get()),
+              optional(1, "ts_without_zone", Types.TimestampType.withoutZone())));
+
+      writeAndValidate(schema);
+    });
+  }
+
+  protected void withSQLConf(Map<String, String> conf, Action action) throws IOException {
+    SQLConf sqlConf = SQLConf.get();
+
+    Map<String, String> currentConfValues = Maps.newHashMap();
+    conf.keySet().forEach(confKey -> {
+      if (sqlConf.contains(confKey)) {
+        String currentConfValue = sqlConf.getConfString(confKey);
+        currentConfValues.put(confKey, currentConfValue);
+      }
+    });
+
+    conf.forEach((confKey, confValue) -> {
+      if (SQLConf.staticConfKeys().contains(confKey)) {
+        throw new RuntimeException("Cannot modify the value of a static config: " + confKey);
+      }
+      sqlConf.setConfString(confKey, confValue);
+    });
+
+    try {
+      action.invoke();
+    } finally {
+      conf.forEach((confKey, confValue) -> {
+        if (currentConfValues.containsKey(confKey)) {
+          sqlConf.setConfString(confKey, currentConfValues.get(confKey));
+        } else {
+          sqlConf.unsetConf(confKey);
+        }
+      });
+    }
+  }
+
+  @FunctionalInterface
+  protected interface Action {
+    void invoke() throws IOException;
+  }
 }