Add support for spark.sql.datetime.java8API.enabled

Motivation: when using connector with thrift server, any queries on data containing date/time types are crashing. It's caused by thrift server enabling `spark.sql.datetime.java8API.enabled` flag for all of its sessions.

Supported JSR-310 types which are already supported by Spark SQL:
- LocalDate
- LocalDateTime
- Instant

Other types are causing Encoder to fail, so they won't work even if taken into account.

As `spark.sql.datetime.java8API.enabled` is disabled by default, it shouldn't cause compatibility issues.

Adds support for the Spark 3.4 TimeTimestampNTZType

SPARK-453

Original PR: #125
---------
Co-authored-by: Ross Lawley <[email protected]>

Author: VVBondarenko

src/integrationTest/java/com/mongodb/spark/sql/connector/RoundTripTest.java

@@ -18,10 +18,13 @@
 package com.mongodb.spark.sql.connector;
 
 import static com.mongodb.spark.sql.connector.mongodb.MongoSparkConnectorHelper.CATALOG;
+import static com.mongodb.spark.sql.connector.schema.ConverterHelper.TIMESTAMP_NTZ_TYPE;
+import static java.time.ZoneOffset.UTC;
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertIterableEquals;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
 import com.mongodb.spark.sql.connector.beans.BoxedBean;
 import com.mongodb.spark.sql.connector.beans.ComplexBean;
@@ -34,7 +37,7 @@
 import java.sql.Timestamp;
 import java.time.Instant;
 import java.time.LocalDate;
-import java.time.ZoneOffset;
+import java.time.LocalDateTime;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -48,6 +51,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.ValueSource;
 
 public class RoundTripTest extends MongoSparkConnectorTestCase {
 
@@ -105,24 +109,29 @@ void testBoxedBean(final TruncateMode mode) {
     assertIterableEquals(dataSetOriginal, dataSetMongo);
   }
 
-  @Test
-  void testDateTimeBean() {
+  @ParameterizedTest()
+  @ValueSource(strings = {"true", "false"})
+  void testDateTimeBean(final String java8DateTimeAPI) {
+    assumeTrue(TIMESTAMP_NTZ_TYPE != null);
     TimeZone original = TimeZone.getDefault();
     try {
-      TimeZone.setDefault(TimeZone.getTimeZone(ZoneOffset.UTC));
+      TimeZone.setDefault(TimeZone.getTimeZone(UTC));
 
       // Given
       long oneHour = TimeUnit.MILLISECONDS.convert(1, TimeUnit.HOURS);
       long oneDay = oneHour * 24;
 
+      Instant epoch = Instant.EPOCH;
       List<DateTimeBean> dataSetOriginal = singletonList(new DateTimeBean(
           new Date(oneDay * 365),
           new Timestamp(oneDay + oneHour),
           LocalDate.of(2000, 1, 1),
-          Instant.EPOCH));
+          epoch,
+          LocalDateTime.ofInstant(epoch, UTC)));
 
       // when
-      SparkSession spark = getOrCreateSparkSession();
+      SparkSession spark = getOrCreateSparkSession(
+          getSparkConf().set("spark.sql.datetime.java8API.enabled", java8DateTimeAPI));
       Encoder<DateTimeBean> encoder = Encoders.bean(DateTimeBean.class);
 
       Dataset<DateTimeBean> dataset = spark.createDataset(dataSetOriginal, encoder);

src/integrationTest/java/com/mongodb/spark/sql/connector/beans/DateTimeBean.java

@@ -21,24 +21,28 @@
 import java.sql.Timestamp;
 import java.time.Instant;
 import java.time.LocalDate;
+import java.time.LocalDateTime;
 import java.util.Objects;
 
 public class DateTimeBean implements Serializable {
   private java.sql.Date sqlDate;
   private java.sql.Timestamp sqlTimestamp;
   private java.time.LocalDate localDate;
   private java.time.Instant instant;
+  private java.time.LocalDateTime localDateTime;
 
   public DateTimeBean() {}
 
   public DateTimeBean(
       final Date sqlDate,
       final Timestamp sqlTimestamp,
       final LocalDate localDate,
-      final Instant instant) {
+      final Instant instant,
+      final LocalDateTime localDateTime) {
     this.sqlDate = sqlDate;
     this.sqlTimestamp = sqlTimestamp;
     this.localDate = localDate;
+    this.localDateTime = localDateTime;
     this.instant = instant;
   }
 
@@ -66,6 +70,14 @@ public void setLocalDate(final LocalDate localDate) {
     this.localDate = localDate;
   }
 
+  public LocalDateTime getLocalDateTime() {
+    return localDateTime;
+  }
+
+  public void setLocalDateTime(final LocalDateTime localDateTime) {
+    this.localDateTime = localDateTime;
+  }
+
   public Instant getInstant() {
     return instant;
   }
@@ -86,20 +98,28 @@ public boolean equals(final Object o) {
     return Objects.equals(sqlDate, that.sqlDate)
         && Objects.equals(sqlTimestamp, that.sqlTimestamp)
         && Objects.equals(localDate, that.localDate)
+        && Objects.equals(localDateTime, that.localDateTime)
         && Objects.equals(instant, that.instant);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(sqlDate, sqlTimestamp, localDate, instant);
+    return Objects.hash(sqlDate, sqlTimestamp, localDate, localDateTime, instant);
   }
 
   @Override
   public String toString() {
-    return "DateTimeBean{" + "sqlDate="
-        + sqlDate + ", sqlTimestamp="
-        + sqlTimestamp + ", localDate="
-        + localDate + ", instant="
-        + instant + '}';
+    return "DateTimeBean{"
+        + "sqlDate="
+        + sqlDate
+        + ", sqlTimestamp="
+        + sqlTimestamp
+        + ", localDate="
+        + localDate
+        + ", localDateTime="
+        + localDateTime
+        + ", instant="
+        + instant
+        + '}';
   }
 }

src/main/java/com/mongodb/spark/sql/connector/schema/BsonDocumentToRowConverter.java

@@ -19,6 +19,7 @@
 
 import static com.mongodb.spark.sql.connector.schema.ConverterHelper.BSON_VALUE_CODEC;
 import static com.mongodb.spark.sql.connector.schema.ConverterHelper.getJsonWriterSettings;
+import static com.mongodb.spark.sql.connector.schema.ConverterHelper.isTimestampNTZ;
 import static com.mongodb.spark.sql.connector.schema.ConverterHelper.toJson;
 import static java.lang.String.format;
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -44,6 +45,7 @@
 import org.apache.spark.sql.Row;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema;
+import org.apache.spark.sql.internal.SQLConf;
 import org.apache.spark.sql.types.ArrayType;
 import org.apache.spark.sql.types.BinaryType;
 import org.apache.spark.sql.types.BooleanType;
@@ -76,6 +78,8 @@
 import org.bson.types.Decimal128;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * The helper for conversion of BsonDocuments to GenericRowWithSchema instances.
@@ -89,13 +93,15 @@
 @NotNull
 public final class BsonDocumentToRowConverter implements Serializable {
   private static final long serialVersionUID = 1L;
+  private static final Logger LOGGER = LoggerFactory.getLogger(BsonDocumentToRowConverter.class);
   private final Function<Row, InternalRow> rowToInternalRowFunction;
   private final StructType schema;
   private final boolean outputExtendedJson;
   private final boolean isPermissive;
   private final boolean dropMalformed;
   private final String columnNameOfCorruptRecord;
   private final boolean schemaContainsCorruptRecordColumn;
+  private final boolean dataTimeJava8APIEnabled;
 
   private boolean corruptedRecord;
 
@@ -114,6 +120,7 @@ public BsonDocumentToRowConverter(final StructType originalSchema, final ReadCon
     this.columnNameOfCorruptRecord = readConfig.getColumnNameOfCorruptRecord();
     this.schemaContainsCorruptRecordColumn = !columnNameOfCorruptRecord.isEmpty()
         && Arrays.asList(schema.fieldNames()).contains(columnNameOfCorruptRecord);
+    this.dataTimeJava8APIEnabled = SQLConf.get().datetimeJava8ApiEnabled();
   }
 
   /** @return the schema for the converter */
@@ -165,6 +172,7 @@ GenericRowWithSchema toRow(final BsonDocument bsonDocument) {
   @VisibleForTesting
   Object convertBsonValue(
       final String fieldName, final DataType dataType, final BsonValue bsonValue) {
+    LOGGER.info("converting bson to value: {} {} {}", fieldName, dataType, bsonValue);
     try {
       if (bsonValue.isNull()) {
         return null;
@@ -179,9 +187,22 @@ Object convertBsonValue(
       } else if (dataType instanceof BooleanType) {
         return convertToBoolean(fieldName, dataType, bsonValue);
       } else if (dataType instanceof DateType) {
-        return convertToDate(fieldName, dataType, bsonValue);
+        Date date = convertToDate(fieldName, dataType, bsonValue);
+        if (dataTimeJava8APIEnabled) {
+          return date.toLocalDate();
+        } else {
+          return date;
+        }
       } else if (dataType instanceof TimestampType) {
-        return convertToTimestamp(fieldName, dataType, bsonValue);
+        Timestamp timestamp = convertToTimestamp(fieldName, dataType, bsonValue);
+        if (dataTimeJava8APIEnabled) {
+          return timestamp.toInstant();
+        } else {
+          return timestamp;
+        }
+      } else if (isTimestampNTZ(dataType)) {
+        Timestamp timestamp = convertToTimestamp(fieldName, dataType, bsonValue);
+        return timestamp.toLocalDateTime();
       } else if (dataType instanceof FloatType) {
         return convertToFloat(fieldName, dataType, bsonValue);
       } else if (dataType instanceof IntegerType) {

src/main/java/com/mongodb/spark/sql/connector/schema/ConverterHelper.java

@@ -22,6 +22,8 @@
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
 import java.util.Base64;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.DataTypes;
 import org.bson.BsonDocument;
 import org.bson.BsonValue;
 import org.bson.codecs.BsonValueCodec;
@@ -39,6 +41,29 @@ static JsonWriterSettings getJsonWriterSettings(final boolean outputExtendedJson
     return outputExtendedJson ? EXTENDED_JSON_WRITER_SETTINGS : RELAXED_JSON_WRITER_SETTINGS;
   }
 
+  /**
+   * The {{TimestampNTZType}} if available or null
+   *
+   * <p>Only available in Spark 3.4+
+   * <p>TODO: SPARK-450 remove code for Spark 4.0
+   */
+  public static final DataType TIMESTAMP_NTZ_TYPE;
+
+  static {
+    DataType timestampNTZType;
+    try {
+      timestampNTZType =
+          (DataType) DataTypes.class.getDeclaredField("TimestampNTZType").get(DataType.class);
+    } catch (IllegalAccessException | NoSuchFieldException e) {
+      timestampNTZType = null;
+    }
+    TIMESTAMP_NTZ_TYPE = timestampNTZType;
+  }
+
+  static boolean isTimestampNTZ(final DataType dataType) {
+    return TIMESTAMP_NTZ_TYPE != null && TIMESTAMP_NTZ_TYPE.acceptsType(dataType);
+  }
+
   private static final JsonWriterSettings RELAXED_JSON_WRITER_SETTINGS =
       JsonWriterSettings.builder()
           .outputMode(JsonMode.RELAXED)

src/main/java/com/mongodb/spark/sql/connector/schema/RowToBsonDocumentConverter.java

@@ -17,6 +17,7 @@
 
 package com.mongodb.spark.sql.connector.schema;
 
+import static com.mongodb.spark.sql.connector.schema.ConverterHelper.isTimestampNTZ;
 import static java.lang.String.format;
 import static java.util.Arrays.asList;
 
@@ -26,6 +27,10 @@
 import com.mongodb.spark.sql.connector.interop.JavaScala;
 import java.io.Serializable;
 import java.math.BigDecimal;
+import java.sql.Timestamp;
+import java.time.Instant;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
@@ -137,9 +142,11 @@ public static ObjectToBsonValue createObjectToBsonValue(
       try {
         return cachedObjectToBsonValue.apply(data);
       } catch (Exception e) {
-        throw new DataException(format(
-            "Cannot cast %s into a BsonValue. %s has no matching BsonValue. Error: %s",
-            data, dataType, e.getMessage()));
+        throw new DataException(
+            format(
+                "Cannot cast %s into a BsonValue. %s has no matching BsonValue. Error: %s",
+                data, dataType, e.getMessage()),
+            e);
       }
     };
   }
@@ -177,12 +184,30 @@ private static ObjectToBsonValue objectToBsonValue(
     } else if (DataTypes.StringType.acceptsType(dataType)) {
       return (data) -> processString((String) data, convertJson);
     } else if (DataTypes.DateType.acceptsType(dataType)
-        || DataTypes.TimestampType.acceptsType(dataType)) {
+        || DataTypes.TimestampType.acceptsType(dataType)
+        || isTimestampNTZ(dataType)) {
       return (data) -> {
         if (data instanceof Date) {
           // Covers java.util.Date, java.sql.Date, java.sql.Timestamp
           return new BsonDateTime(((Date) data).getTime());
         }
+        if (data instanceof Instant) {
+          return new BsonDateTime(((Instant) data).toEpochMilli());
+        }
+        if (data instanceof LocalDateTime) {
+          LocalDateTime dateTime = (LocalDateTime) data;
+          return new BsonDateTime(Timestamp.valueOf(dateTime).getTime());
+        }
+        if (data instanceof LocalDate) {
+          long epochSeconds = ((LocalDate) data).toEpochDay() * 24L * 3600L;
+          return new BsonDateTime(epochSeconds * 1000L);
+        }
+
+        /*
+        NOTE 1: ZonedDateTime, OffsetDateTime, OffsetTime are not explicitly supported by Spark and cause the Encoder resolver to fail
+        due to cyclic dependency in the ZoneOffset. Subject for review after it changes (if ever).
+        NOTE 2: LocalTime type is not represented neither in Bson nor in Spark
+         */
         throw new MongoSparkException(
             "Unsupported date type: " + data.getClass().getSimpleName());
       };