[spark] Support timestamp_ntz for Spark 3.4+ (apache#3569)

docs/content/spark/quick-start.md

@@ -310,7 +310,12 @@ All Spark's data types are available in package `org.apache.spark.sql.types`.
     </tr>
     <tr>
       <td><code>TimestampType</code></td>
-      <td><code>TimestampType</code>, <code>LocalZonedTimestamp</code></td>
+      <td><code>LocalZonedTimestamp</code></td>
+      <td>true</td>
+    </tr>
+    <tr>
+      <td><code>TimestampNTZType(Spark3.4+)</code></td>
+      <td><code>TimestampType</code></td>
       <td>true</td>
     </tr>
     <tr>
@@ -325,3 +330,11 @@ All Spark's data types are available in package `org.apache.spark.sql.types`.
     </tr>
     </tbody>
 </table>
+
+{{< hint warning >}}
+Due to the previous design, in Spark3.3 and below, Paimon will map both Paimon's TimestampType and LocalZonedTimestamp to Spark's TimestampType, and only correctly handle with TimestampType.
+
+Therefore, when using Spark3.3 and below, reads Paimon table with LocalZonedTimestamp type written by other engines, such as Flink, the query result of LocalZonedTimestamp type will have time zone offset, which needs to be adjusted manually.
+
+When using Spark3.4 and above, all timestamp types can be parsed correctly.
+{{< /hint >}}

paimon-spark/paimon-spark-3.1/src/main/scala/org/apache/paimon/spark/util/shim/TypeUtils.scala

@@ -0,0 +1,25 @@
+/*
+ * 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.paimon.spark.util.shim
+
+object TypeUtils {
+
+  // Since Spark 3.3 and below do not support timestamp ntz, treat Paimon TimestampType as Spark TimestampType
+  def treatPaimonTimestampTypeAsSparkTimestampType(): Boolean = true
+}

paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/util/shim/TypeUtils.scala

@@ -0,0 +1,25 @@
+/*
+ * 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.paimon.spark.util.shim
+
+object TypeUtils {
+
+  // Since Spark 3.3 and below do not support timestamp ntz, treat Paimon TimestampType as Spark TimestampType
+  def treatPaimonTimestampTypeAsSparkTimestampType(): Boolean = true
+}

paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/util/shim/TypeUtils.scala

@@ -0,0 +1,25 @@
+/*
+ * 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.paimon.spark.util.shim
+
+object TypeUtils {
+
+  // Since Spark 3.3 and below do not support timestamp ntz, treat Paimon TimestampType as Spark TimestampType
+  def treatPaimonTimestampTypeAsSparkTimestampType(): Boolean = true
+}

paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkInternalRow.java

@@ -23,6 +23,7 @@
 import org.apache.paimon.data.InternalMap;
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.data.Timestamp;
+import org.apache.paimon.spark.util.shim.TypeUtils;
 import org.apache.paimon.types.ArrayType;
 import org.apache.paimon.types.BigIntType;
 import org.apache.paimon.types.DataType;
@@ -283,7 +284,11 @@ public static org.apache.spark.sql.catalyst.InternalRow fromPaimon(
     }
 
     public static long fromPaimon(Timestamp timestamp) {
-        return DateTimeUtils.fromJavaTimestamp(timestamp.toSQLTimestamp());
+        if (TypeUtils.treatPaimonTimestampTypeAsSparkTimestampType()) {
+            return DateTimeUtils.fromJavaTimestamp(timestamp.toSQLTimestamp());
+        } else {
+            return timestamp.toMicros();
+        }
     }
 
     public static ArrayData fromPaimon(InternalArray array, ArrayType arrayType) {

paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkRow.java

@@ -24,6 +24,7 @@
 import org.apache.paimon.data.InternalMap;
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.data.Timestamp;
+import org.apache.paimon.spark.util.shim.TypeUtils;
 import org.apache.paimon.types.ArrayType;
 import org.apache.paimon.types.DataType;
 import org.apache.paimon.types.DateType;
@@ -170,11 +171,21 @@ private static int toPaimonDate(Object object) {
 
     private static Timestamp toPaimonTimestamp(Object object) {
         if (object instanceof java.sql.Timestamp) {
-            return Timestamp.fromSQLTimestamp((java.sql.Timestamp) object);
+            java.sql.Timestamp ts = (java.sql.Timestamp) object;
+            if (TypeUtils.treatPaimonTimestampTypeAsSparkTimestampType()) {
+                return Timestamp.fromSQLTimestamp(ts);
+            } else {
+                return Timestamp.fromInstant(ts.toInstant());
+            }
         } else if (object instanceof java.time.Instant) {
-            LocalDateTime localDateTime =
-                    LocalDateTime.ofInstant((Instant) object, ZoneId.systemDefault());
-            return Timestamp.fromLocalDateTime(localDateTime);
+            Instant instant = (Instant) object;
+            if (TypeUtils.treatPaimonTimestampTypeAsSparkTimestampType()) {
+                LocalDateTime localDateTime =
+                        LocalDateTime.ofInstant((Instant) object, ZoneId.systemDefault());
+                return Timestamp.fromLocalDateTime(localDateTime);
+            } else {
+                return Timestamp.fromInstant(instant);
+            }
         } else {
             return Timestamp.fromLocalDateTime((LocalDateTime) object);
         }

paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkTypeUtils.java

@@ -18,6 +18,7 @@
 
 package org.apache.paimon.spark;
 
+import org.apache.paimon.spark.util.shim.TypeUtils;
 import org.apache.paimon.types.ArrayType;
 import org.apache.paimon.types.BigIntType;
 import org.apache.paimon.types.BinaryType;
@@ -150,7 +151,11 @@ public DataType visit(TimeType timeType) {
 
         @Override
         public DataType visit(TimestampType timestampType) {
-            return DataTypes.TimestampType;
+            if (TypeUtils.treatPaimonTimestampTypeAsSparkTimestampType()) {
+                return DataTypes.TimestampType;
+            } else {
+                return DataTypes.TimestampNTZType;
+            }
         }
 
         @Override
@@ -308,13 +313,20 @@ public org.apache.paimon.types.DataType atomic(DataType atomic) {
             } else if (atomic instanceof org.apache.spark.sql.types.DateType) {
                 return new DateType();
             } else if (atomic instanceof org.apache.spark.sql.types.TimestampType) {
-                return new TimestampType();
+                if (TypeUtils.treatPaimonTimestampTypeAsSparkTimestampType()) {
+                    return new TimestampType();
+                } else {
+                    return new LocalZonedTimestampType();
+                }
             } else if (atomic instanceof org.apache.spark.sql.types.DecimalType) {
                 return new DecimalType(
                         ((org.apache.spark.sql.types.DecimalType) atomic).precision(),
                         ((org.apache.spark.sql.types.DecimalType) atomic).scale());
             } else if (atomic instanceof org.apache.spark.sql.types.BinaryType) {
                 return new VarBinaryType(VarBinaryType.MAX_LENGTH);
+            } else if (atomic instanceof org.apache.spark.sql.types.TimestampNTZType) {
+                // Move TimestampNTZType to the end for compatibility with spark3.3 and below
+                return new TimestampType();
             }
 
             throw new UnsupportedOperationException(

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/util/shim/TypeUtils.scala

@@ -0,0 +1,24 @@
+/*
+ * 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.paimon.spark.util.shim
+
+object TypeUtils {
+
+  def treatPaimonTimestampTypeAsSparkTimestampType(): Boolean = false
+}

paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/SparkInternalRowTest.java

@@ -37,10 +37,12 @@
 import java.time.LocalDateTime;
 import java.util.AbstractMap;
 import java.util.Map;
+import java.util.TimeZone;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import scala.Function1;
+import scala.collection.JavaConverters;
 
 import static org.apache.paimon.data.BinaryString.fromString;
 import static org.apache.paimon.spark.SparkTypeTest.ALL_TYPES;
@@ -51,6 +53,8 @@ public class SparkInternalRowTest {
 
     @Test
     public void test() {
+        TimeZone tz = TimeZone.getDefault();
+        TimeZone.setDefault(TimeZone.getTimeZone("Asia/Shanghai"));
         InternalRow rowData =
                 GenericRow.of(
                         1,
@@ -77,6 +81,7 @@ public void test() {
                         23567222L,
                         "varbinary_v".getBytes(StandardCharsets.UTF_8),
                         Timestamp.fromLocalDateTime(LocalDateTime.parse("2007-12-03T10:15:30")),
+                        Timestamp.fromLocalDateTime(LocalDateTime.parse("2007-12-03T10:15:30")),
                         DateTimeUtils.toInternal(LocalDate.parse("2022-05-02")),
                         Decimal.fromBigDecimal(BigDecimal.valueOf(0.21), 2, 2),
                         Decimal.fromBigDecimal(BigDecimal.valueOf(65782123123.01), 38, 2),
@@ -93,32 +98,40 @@ public void test() {
                         sparkConverter.apply(new SparkInternalRow(ALL_TYPES).replace(rowData));
 
         String expected =
-                "{"
-                        + "\"id\":1,"
-                        + "\"name\":\"jingsong\","
-                        + "\"char\":\"apache\","
-                        + "\"varchar\":\"paimon\","
-                        + "\"salary\":22.2,"
-                        + "\"locations\":{\"key1\":{\"posX\":1.2,\"posY\":2.3},\"key2\":{\"posX\":2.4,\"posY\":3.5}},"
-                        + "\"strArray\":[\"v1\",\"v5\"],"
-                        + "\"intArray\":[10,30],"
-                        + "\"boolean\":true,"
-                        + "\"tinyint\":22,"
-                        + "\"smallint\":356,"
-                        + "\"bigint\":23567222,"
-                        + "\"bytes\":\"dmFyYmluYXJ5X3Y=\","
-                        + "\"timestamp\":\"2007-12-03 10:15:30\","
-                        + "\"date\":\"2022-05-02\","
-                        + "\"decimal\":0.21,"
-                        + "\"decimal2\":65782123123.01,"
-                        + "\"decimal3\":62123123.5"
-                        + "}";
-        assertThat(sparkRow.json()).isEqualTo(expected);
+                "1,"
+                        + "jingsong,"
+                        + "apache,"
+                        + "paimon,"
+                        + "22.2,"
+                        + "Map(key2 -> [2.4,3.5], key1 -> [1.2,2.3]),"
+                        + "WrappedArray(v1, v5),"
+                        + "WrappedArray(10, 30),"
+                        + "true,"
+                        + "22,"
+                        + "356,"
+                        + "23567222,"
+                        + "[B@,"
+                        + "2007-12-03 18:15:30.0,"
+                        + "2007-12-03T10:15:30,"
+                        + "2022-05-02,"
+                        + "0.21,"
+                        + "65782123123.01,"
+                        + "62123123.5";
+        assertThat(sparkRowToString(sparkRow)).isEqualTo(expected);
 
         SparkRow sparkRowData = new SparkRow(ALL_TYPES, sparkRow);
         sparkRow =
                 (org.apache.spark.sql.Row)
                         sparkConverter.apply(new SparkInternalRow(ALL_TYPES).replace(sparkRowData));
-        assertThat(sparkRow.json()).isEqualTo(expected);
+        assertThat(sparkRowToString(sparkRow)).isEqualTo(expected);
+        TimeZone.setDefault(tz);
+    }
+
+    private String sparkRowToString(org.apache.spark.sql.Row row) {
+        return JavaConverters.seqAsJavaList(row.toSeq()).stream()
+                .map(Object::toString)
+                // Since the toString result of Spark's binary col is unstable, replace it
+                .map(x -> x.startsWith("[B@") ? "[B@" : x)
+                .collect(Collectors.joining(","));
     }
 }

paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/SparkTypeTest.java

@@ -65,7 +65,8 @@ public class SparkTypeTest {
                     .field("smallint", DataTypes.SMALLINT())
                     .field("bigint", DataTypes.BIGINT())
                     .field("bytes", DataTypes.BYTES())
-                    .field("timestamp", DataTypes.TIMESTAMP())
+                    .field("timestamp", DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE())
+                    .field("timestamp_ntz", DataTypes.TIMESTAMP())
                     .field("date", DataTypes.DATE())
                     .field("decimal", DataTypes.DECIMAL(2, 2))
                     .field("decimal2", DataTypes.DECIMAL(38, 2))
@@ -95,6 +96,7 @@ public void testAllTypes() {
                         + "StructField(bigint,LongType,true),"
                         + "StructField(bytes,BinaryType,true),"
                         + "StructField(timestamp,TimestampType,true),"
+                        + "StructField(timestamp_ntz,TimestampNTZType,true),"
                         + "StructField(date,DateType,true),"
                         + "StructField(decimal,DecimalType(2,2),true),"
                         + "StructField(decimal2,DecimalType(38,2),true),"

paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/PaimonSparkTestBase.scala

@@ -36,6 +36,7 @@ import org.scalactic.source.Position
 import org.scalatest.Tag
 
 import java.io.File
+import java.util.TimeZone
 
 import scala.util.Random
 
@@ -101,6 +102,18 @@ class PaimonSparkTestBase
     withTempDir(file1 => withTempDir(file2 => f(file1, file2)))
   }
 
+  protected def withTimeZone(timeZone: String)(f: => Unit): Unit = {
+    withSQLConf("spark.sql.session.timeZone" -> timeZone) {
+      val originTimeZone = TimeZone.getDefault
+      try {
+        TimeZone.setDefault(TimeZone.getTimeZone(timeZone))
+        f
+      } finally {
+        TimeZone.setDefault(originTimeZone)
+      }
+    }
+  }
+
   override def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
       pos: Position): Unit = {
     println(testName)