This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new bb5498ff32 [spark] Support timestamp_ntz columns in zorder and hilbert
clustering (#10011)
bb5498ff32 is described below
commit bb5498ff324f58c53a8292e5c8339222959c61c7
Author: cxzl25 <[email protected]>
AuthorDate: Mon Sep 21 10:42:10 2026 +0800
[spark] Support timestamp_ntz columns in zorder and hilbert clustering
(#10011)
---
.../apache/paimon/spark/sort/SparkHilbertUDF.java | 17 ++++++
.../apache/paimon/spark/sort/SparkZOrderUDF.java | 17 ++++++
.../paimon/spark/sort/SparkHilbertUDFTest.java | 57 ++++++++++++++++++++
.../paimon/spark/sort/SparkZOrderUDFTest.java | 61 ++++++++++++++++++++++
4 files changed, 152 insertions(+)
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/sort/SparkHilbertUDF.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/sort/SparkHilbertUDF.java
index 971e8342e5..0e60000389 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/sort/SparkHilbertUDF.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/sort/SparkHilbertUDF.java
@@ -37,11 +37,14 @@ import org.apache.spark.sql.types.IntegerType;
import org.apache.spark.sql.types.LongType;
import org.apache.spark.sql.types.ShortType;
import org.apache.spark.sql.types.StringType;
+import org.apache.spark.sql.types.TimestampNTZType;
import org.apache.spark.sql.types.TimestampType;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.Serializable;
+import java.time.LocalDateTime;
+import java.time.ZoneOffset;
import scala.collection.JavaConverters;
import scala.collection.Seq;
@@ -110,6 +113,18 @@ public class SparkHilbertUDF implements Serializable {
return udf;
}
+ private UserDefinedFunction timestampNtzToLongUDF() {
+ UserDefinedFunction udf =
+ functions
+ .udf(
+ (LocalDateTime value) ->
+ value == null ? null :
value.toEpochSecond(ZoneOffset.UTC),
+ DataTypes.LongType)
+ .withName("TIMESTAMP_NTZ_LONG");
+
+ return udf;
+ }
+
private UserDefinedFunction longToOrderedLongUDF() {
UserDefinedFunction udf =
functions
@@ -227,6 +242,8 @@ public class SparkHilbertUDF implements Serializable {
return booleanToOrderedLongUDF().apply(column);
} else if (type instanceof TimestampType) {
return
longToOrderedLongUDF().apply(column.cast(DataTypes.LongType));
+ } else if (type instanceof TimestampNTZType) {
+ return
longToOrderedLongUDF().apply(timestampNtzToLongUDF().apply(column));
} else if (type instanceof DecimalType) {
return
longToOrderedLongUDF().apply(column.cast(DataTypes.LongType));
} else if (type instanceof DateType) {
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/sort/SparkZOrderUDF.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/sort/SparkZOrderUDF.java
index fbd944e9a1..7bb80daaf3 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/sort/SparkZOrderUDF.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/sort/SparkZOrderUDF.java
@@ -35,10 +35,13 @@ import org.apache.spark.sql.types.IntegerType;
import org.apache.spark.sql.types.LongType;
import org.apache.spark.sql.types.ShortType;
import org.apache.spark.sql.types.StringType;
+import org.apache.spark.sql.types.TimestampNTZType;
import org.apache.spark.sql.types.TimestampType;
import java.io.Serializable;
import java.nio.ByteBuffer;
+import java.time.LocalDateTime;
+import java.time.ZoneOffset;
import scala.collection.JavaConverters;
import scala.collection.Seq;
@@ -166,6 +169,18 @@ public class SparkZOrderUDF implements Serializable {
return udf;
}
+ private UserDefinedFunction timestampNtzToLongUDF() {
+ UserDefinedFunction udf =
+ functions
+ .udf(
+ (LocalDateTime value) ->
+ value == null ? null :
value.toEpochSecond(ZoneOffset.UTC),
+ DataTypes.LongType)
+ .withName("TIMESTAMP_NTZ_LONG");
+
+ return udf;
+ }
+
private UserDefinedFunction longToOrderedBytesUDF() {
int position = inputCol;
UserDefinedFunction udf =
@@ -347,6 +362,8 @@ public class SparkZOrderUDF implements Serializable {
return booleanToOrderedBytesUDF().apply(column);
} else if (type instanceof TimestampType) {
return
longToOrderedBytesUDF().apply(column.cast(DataTypes.LongType));
+ } else if (type instanceof TimestampNTZType) {
+ return
longToOrderedBytesUDF().apply(timestampNtzToLongUDF().apply(column));
} else if (type instanceof DateType) {
return
longToOrderedBytesUDF().apply(column.cast(DataTypes.LongType));
} else {
diff --git
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/sort/SparkHilbertUDFTest.java
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/sort/SparkHilbertUDFTest.java
index d96b68388d..abb24eb731 100644
---
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/sort/SparkHilbertUDFTest.java
+++
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/sort/SparkHilbertUDFTest.java
@@ -28,6 +28,8 @@ import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.junit.jupiter.api.Test;
+import java.time.LocalDateTime;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
@@ -85,4 +87,59 @@ public class SparkHilbertUDFTest {
SparkSession.clearDefaultSession();
}
}
+
+ @Test
+ void testTimestampNtzColumnIsSupported() {
+ SparkSession spark =
+ SparkSession.builder()
+ .master("local[1]")
+ .appName("spark-hilbert-udf-ntz-test")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ try {
+ StructType schema =
+ new StructType(
+ new StructField[] {
+ new StructField(
+ "a", DataTypes.TimestampNTZType, true,
Metadata.empty())
+ });
+ Dataset<Row> df =
+ spark.createDataFrame(
+ Arrays.asList(
+ RowFactory.create(LocalDateTime.of(2024,
1, 1, 0, 0, 0)),
+ RowFactory.create(LocalDateTime.of(2024,
6, 1, 12, 0, 0)),
+ RowFactory.create(LocalDateTime.of(2025,
1, 1, 0, 0, 0)),
+ RowFactory.create((LocalDateTime) null)),
+ schema);
+
+ SparkHilbertUDF udf = new SparkHilbertUDF();
+ // A paimon TIMESTAMP column reads back as timestamp_ntz unless
the legacy mapping is
+ // on, and the dispatch used to reject it, so ordering by it threw
"the type is
+ // unsupported" instead of clustering.
+ List<Row> rows =
+ df.select(
+ df.col("a"),
+ udf.sortedLexicographically(
+ df.col("a"),
DataTypes.TimestampNTZType)
+ .as("hilbert"))
+ .orderBy(df.col("a").asc_nulls_first())
+ .collectAsList();
+
+ assertThat(rows).hasSize(4);
+ assertThat(rows.get(0).isNullAt(0)).isTrue();
+ assertThat(rows.get(0).getLong(1)).isEqualTo(Long.MAX_VALUE);
+
+ List<Long> nonNull = new ArrayList<>();
+ for (Row row : rows.subList(1, rows.size())) {
+ nonNull.add(row.getLong(1));
+ }
+ for (int i = 1; i < nonNull.size(); i++) {
+ assertThat(nonNull.get(i)).isGreaterThan(nonNull.get(i - 1));
+ }
+ } finally {
+ spark.stop();
+ SparkSession.clearActiveSession();
+ SparkSession.clearDefaultSession();
+ }
+ }
}
diff --git
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/sort/SparkZOrderUDFTest.java
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/sort/SparkZOrderUDFTest.java
index ef07b0f0cb..7a091f8b63 100644
---
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/sort/SparkZOrderUDFTest.java
+++
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/sort/SparkZOrderUDFTest.java
@@ -29,6 +29,8 @@ import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.junit.jupiter.api.Test;
+import java.time.LocalDateTime;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
@@ -91,4 +93,63 @@ public class SparkZOrderUDFTest {
SparkSession.clearDefaultSession();
}
}
+
+ @Test
+ void testTimestampNtzColumnIsSupported() {
+ SparkSession spark =
+ SparkSession.builder()
+ .master("local[1]")
+ .appName("spark-zorder-udf-ntz-test")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ try {
+ StructType schema =
+ new StructType(
+ new StructField[] {
+ new StructField(
+ "a", DataTypes.TimestampNTZType, true,
Metadata.empty())
+ });
+ Dataset<Row> df =
+ spark.createDataFrame(
+ Arrays.asList(
+ RowFactory.create(LocalDateTime.of(2024,
1, 1, 0, 0, 0)),
+ RowFactory.create(LocalDateTime.of(2024,
6, 1, 12, 0, 0)),
+ RowFactory.create(LocalDateTime.of(2025,
1, 1, 0, 0, 0)),
+ RowFactory.create((LocalDateTime) null)),
+ schema);
+
+ SparkZOrderUDF udf = new SparkZOrderUDF(1, 8, Integer.MAX_VALUE);
+ // A paimon TIMESTAMP column reads back as timestamp_ntz unless
the legacy mapping is
+ // on, and the dispatch used to reject it, so ordering by it threw
"the type is
+ // unsupported" instead of clustering.
+ List<Row> rows =
+ df.select(
+ df.col("a"),
+ functions
+ .hex(
+
udf.sortedLexicographically(
+ df.col("a"),
+
DataTypes.TimestampNTZType))
+ .as("zvalue"))
+ .orderBy(df.col("a").asc_nulls_first())
+ .collectAsList();
+
+ assertThat(rows).hasSize(4);
+ assertThat(rows.get(0).isNullAt(0)).isTrue();
+ assertThat(rows.get(0).getString(1)).isEqualTo("0000000000000000");
+
+ // the z-value of a non-null timestamp has to rise with the
timestamp itself
+ List<String> nonNull = new ArrayList<>();
+ for (Row row : rows.subList(1, rows.size())) {
+ nonNull.add(row.getString(1));
+ }
+ for (int i = 1; i < nonNull.size(); i++) {
+ assertThat(nonNull.get(i)).isGreaterThan(nonNull.get(i - 1));
+ }
+ } finally {
+ spark.stop();
+ SparkSession.clearActiveSession();
+ SparkSession.clearDefaultSession();
+ }
+ }
}
