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 2ec825906d [spark] Support decimal columns in zorder clustering
(#10196)
2ec825906d is described below
commit 2ec825906d8970ac23be8ef40fff5907c06c19f4
Author: jackylee <[email protected]>
AuthorDate: Sat Oct 3 17:08:30 2026 +0800
[spark] Support decimal columns in zorder clustering (#10196)
---
.../apache/paimon/spark/sort/SparkZOrderUDF.java | 50 ++++++++
.../paimon/spark/sort/SparkZOrderUDFTest.java | 127 +++++++++++++++++++++
2 files changed, 177 insertions(+)
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 7bb80daaf3..9abdc4b83c 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
@@ -29,6 +29,7 @@ import org.apache.spark.sql.types.ByteType;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.DateType;
+import org.apache.spark.sql.types.DecimalType;
import org.apache.spark.sql.types.DoubleType;
import org.apache.spark.sql.types.FloatType;
import org.apache.spark.sql.types.IntegerType;
@@ -39,9 +40,12 @@ import org.apache.spark.sql.types.TimestampNTZType;
import org.apache.spark.sql.types.TimestampType;
import java.io.Serializable;
+import java.math.BigDecimal;
+import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
+import java.util.Arrays;
import scala.collection.JavaConverters;
import scala.collection.Seq;
@@ -256,6 +260,25 @@ public class SparkZOrderUDF implements Serializable {
return udf;
}
+ private UserDefinedFunction decimalToOrderedBytesUDF(int width) {
+ UserDefinedFunction udf =
+ functions
+ .udf(
+ (BigDecimal value) -> {
+ if (value == null) {
+ return new byte[width];
+ }
+ return
decimalToOrderedBytes(value.unscaledValue(), width);
+ },
+ DataTypes.BinaryType)
+ .withName("DECIMAL_ORDERED_BYTES");
+
+ this.inputCol++;
+ increaseOutputSize(width);
+
+ return udf;
+ }
+
private UserDefinedFunction booleanToOrderedBytesUDF() {
int position = inputCol;
UserDefinedFunction udf =
@@ -364,6 +387,13 @@ public class SparkZOrderUDF implements Serializable {
return
longToOrderedBytesUDF().apply(column.cast(DataTypes.LongType));
} else if (type instanceof TimestampNTZType) {
return
longToOrderedBytesUDF().apply(timestampNtzToLongUDF().apply(column));
+ } else if (type instanceof DecimalType) {
+ // Encode the unscaled value as fixed-width, sign-flipped two's
complement (mirroring
+ // the core ZIndexer). Casting to LongType instead would raise
CAST_OVERFLOW under ANSI
+ // (Spark 4's default) for values outside the BIGINT range,
aborting the z-order
+ // compact.
+ int width = decimalBytes(((DecimalType) type).precision());
+ return decimalToOrderedBytesUDF(width).apply(column);
} else if (type instanceof DateType) {
return
longToOrderedBytesUDF().apply(column.cast(DataTypes.LongType));
} else {
@@ -377,4 +407,24 @@ public class SparkZOrderUDF implements Serializable {
private void increaseOutputSize(int bytes) {
totalOutputBytes = Math.min(totalOutputBytes + bytes, maxOutputSize);
}
+
+ /** Two's-complement byte width that holds any unscaled value with {@code
precision} digits. */
+ static int decimalBytes(int precision) {
+ return BigInteger.TEN.pow(precision).bitLength() / 8 + 1;
+ }
+
+ /**
+ * Fixed-width, sign-flipped big-endian two's complement of {@code
unscaled}, so signed order
+ * becomes the unsigned lexicographic order a z-value compares with.
{@code width} is wide
+ * enough for any value of the column's precision, so values past the long
range keep distinct
+ * keys.
+ */
+ static byte[] decimalToOrderedBytes(BigInteger unscaled, int width) {
+ byte[] buffer = new byte[width];
+ Arrays.fill(buffer, 0, width, unscaled.signum() < 0 ? (byte) 0xFF :
(byte) 0x00);
+ byte[] magnitude = unscaled.toByteArray();
+ System.arraycopy(magnitude, 0, buffer, width - magnitude.length,
magnitude.length);
+ buffer[0] ^= (byte) 0x80;
+ return buffer;
+ }
}
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 7a091f8b63..4c57893f93 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,7 @@ import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.junit.jupiter.api.Test;
+import java.math.BigDecimal;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Arrays;
@@ -152,4 +153,130 @@ public class SparkZOrderUDFTest {
SparkSession.clearDefaultSession();
}
}
+
+ @Test
+ void testDecimalColumnIsSupported() {
+ SparkSession spark =
+ SparkSession.builder()
+ .master("local[1]")
+ .appName("spark-zorder-udf-decimal-test")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ try {
+ StructType schema =
+ new StructType(
+ new StructField[] {
+ new StructField(
+ "a",
+ DataTypes.createDecimalType(10, 2),
+ true,
+ Metadata.empty())
+ });
+ Dataset<Row> df =
+ spark.createDataFrame(
+ Arrays.asList(
+ RowFactory.create(new BigDecimal("1.50")),
+ RowFactory.create(new
BigDecimal("100.25")),
+ RowFactory.create(new
BigDecimal("12345.67")),
+ RowFactory.create((BigDecimal) null)),
+ schema);
+
+ SparkZOrderUDF udf = new SparkZOrderUDF(1, 8, Integer.MAX_VALUE);
+ // A decimal column used to fall through to the default branch and
throw "the type is
+ // unsupported". It now encodes the unscaled value as fixed-width
ordered bytes.
+ List<Row> rows =
+ df.select(
+ df.col("a"),
+ functions
+ .hex(
+
udf.sortedLexicographically(
+ df.col("a"),
+
DataTypes.createDecimalType(10, 2)))
+ .as("zvalue"))
+ .orderBy(df.col("a").asc_nulls_first())
+ .collectAsList();
+
+ assertThat(rows).hasSize(4);
+ assertThat(rows.get(0).isNullAt(0)).isTrue();
+ // NULL is the all-zero sentinel, width bytes wide (two hex chars
each).
+ int width = SparkZOrderUDF.decimalBytes(10);
+ assertThat(rows.get(0).getString(1))
+ .isEqualTo(String.join("",
java.util.Collections.nCopies(width, "00")));
+
+ // the z-value rises with the decimal value
+ 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();
+ }
+ }
+
+ @Test
+ void testDecimalOutsideBigIntRangeIsOrderedWithoutOverflow() {
+ SparkSession spark =
+ SparkSession.builder()
+ .master("local[1]")
+ .appName("spark-zorder-udf-decimal-wide-test")
+ .config("spark.ui.enabled", "false")
+ // ANSI (Spark 4's default) would make a decimal->long
cast raise
+ // CAST_OVERFLOW on these values; the fixed-width
encoding must not.
+ .config("spark.sql.ansi.enabled", "true")
+ .getOrCreate();
+ try {
+ StructType schema =
+ new StructType(
+ new StructField[] {
+ new StructField(
+ "a",
+ DataTypes.createDecimalType(38, 2),
+ true,
+ Metadata.empty())
+ });
+ // Values straddling the BIGINT range, including both just past
Long.MAX and Long.MIN.
+ List<BigDecimal> values =
+ Arrays.asList(
+ new BigDecimal("-9223372036854775809.00"),
+ new BigDecimal("-1.00"),
+ new BigDecimal("0.00"),
+ new BigDecimal("1.00"),
+ new BigDecimal("9223372036854775808.00"),
+ new BigDecimal("9223372036854775809.00"));
+ List<Row> input = new ArrayList<>();
+ for (BigDecimal v : values) {
+ input.add(RowFactory.create(v));
+ }
+ Dataset<Row> df = spark.createDataFrame(input, schema);
+
+ SparkZOrderUDF udf = new SparkZOrderUDF(1, 8, Integer.MAX_VALUE);
+ List<Row> rows =
+ df.select(
+ df.col("a"),
+ functions
+ .hex(
+
udf.sortedLexicographically(
+ df.col("a"),
+
DataTypes.createDecimalType(38, 2)))
+ .as("zvalue"))
+ .orderBy(df.col("a").asc())
+ .collectAsList();
+
+ // No CAST_OVERFLOW, and the z-value is strictly increasing with
the decimal across the
+ // whole range (negatives, zero, and both out-of-long positives
stay distinct/ordered).
+ assertThat(rows).hasSize(values.size());
+ for (int i = 1; i < rows.size(); i++) {
+ assertThat(rows.get(i).getString(1)).isGreaterThan(rows.get(i
- 1).getString(1));
+ }
+ } finally {
+ spark.stop();
+ SparkSession.clearActiveSession();
+ SparkSession.clearDefaultSession();
+ }
+ }
}
