From 6919b32e527351a564bab6a62c5c7fe8e882f31e Mon Sep 17 00:00:00 2001 From: zhangchen Date: Mon, 25 Oct 2021 18:10:34 +0800 Subject: [PATCH 1/7] v3.0 --- .../apache/iceberg/spark/IcebergSpark.java | 3 +- .../iceberg/spark/SparkValueConverter.java | 43 ++++++ .../spark/source/TestIcebergSpark.java | 132 +++++++++++++++++- 3 files changed, 170 insertions(+), 8 deletions(-) diff --git a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java index ac659f6c7b13..6924851cefd0 100644 --- a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java +++ b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java @@ -34,6 +34,7 @@ public static void registerBucketUDF(SparkSession session, String funcName, Data SparkTypeToType typeConverter = new SparkTypeToType(); Type sourceIcebergType = typeConverter.atomic(sourceType); Transform bucket = Transforms.bucket(sourceIcebergType, numBuckets); - session.udf().register(funcName, bucket::apply, DataTypes.IntegerType); + session.udf().register(funcName, + value -> bucket.apply(SparkValueConverter.convertAtomicValue(sourceType, value)), DataTypes.IntegerType); } } diff --git a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java index 92c812a9b979..35513c12f8d5 100644 --- a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java +++ b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java @@ -33,6 +33,21 @@ import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.catalyst.util.DateTimeUtils; +import org.apache.spark.sql.types.BinaryType; +import org.apache.spark.sql.types.BooleanType; +import org.apache.spark.sql.types.ByteType; +import org.apache.spark.sql.types.CharType; +import org.apache.spark.sql.types.DataType; +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; +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.TimestampType; +import org.apache.spark.sql.types.VarcharType; /** * A utility class that converts Spark values to Iceberg's internal representation. @@ -117,4 +132,32 @@ private static Record convert(Types.StructType struct, Row row) { } return record; } + + @SuppressWarnings("checkstyle:CyclomaticComplexity") + public static Object convertAtomicValue(DataType atomic, Object object) { + if (atomic instanceof BooleanType || + atomic instanceof IntegerType || + atomic instanceof LongType || + atomic instanceof StringType || + atomic instanceof FloatType || + atomic instanceof DoubleType || + atomic instanceof DecimalType || + atomic instanceof CharType || + atomic instanceof VarcharType) { + return object; + } else if (atomic instanceof ShortType) { + return ((Short) object).intValue(); + } else if (atomic instanceof ByteType) { + return ((Byte) object).intValue(); + } else if (atomic instanceof DateType) { + return DateTimeUtils.fromJavaDate((Date) object); + } else if (atomic instanceof TimestampType) { + return DateTimeUtils.fromJavaTimestamp((Timestamp) object); + } else if (atomic instanceof BinaryType) { + return ByteBuffer.wrap((byte[]) object); + } + + throw new UnsupportedOperationException( + "Not a supported type: " + atomic.catalogString()); + } } diff --git a/spark/v3.0/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java b/spark/v3.0/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java index e83709a3b2a9..39ff681ac74e 100644 --- a/spark/v3.0/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java +++ b/spark/v3.0/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java @@ -19,13 +19,22 @@ package org.apache.iceberg.spark.source; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.sql.Date; +import java.sql.Timestamp; import java.util.List; import org.apache.iceberg.spark.IcebergSpark; import org.apache.iceberg.transforms.Transforms; import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.util.DateTimeUtils; +import org.apache.spark.sql.types.CharType; import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.DecimalType; +import org.apache.spark.sql.types.VarcharType; +import org.assertj.core.api.Assertions; import org.junit.AfterClass; import org.junit.Assert; import org.junit.BeforeClass; @@ -48,23 +57,132 @@ public static void stopSpark() { } @Test - public void testRegisterBucketUDF() { + public void testRegisterIntegerBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_int_16", DataTypes.IntegerType, 16); List results = spark.sql("SELECT iceberg_bucket_int_16(1)").collectAsList(); Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), results.get(0).getInt(0)); + } + + @Test + public void testRegisterShortBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_short_16", DataTypes.ShortType, 16); + List results = spark.sql("SELECT iceberg_bucket_short_16(1S)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterByteBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_byte_16", DataTypes.ByteType, 16); + List results = spark.sql("SELECT iceberg_bucket_byte_16(1Y)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), + results.get(0).getInt(0)); + } + @Test + public void testRegisterLongBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_long_16", DataTypes.LongType, 16); - List results2 = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); - Assert.assertEquals(1, results2.size()); + List results = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); + Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.LongType.get(), 16).apply(1L), - results2.get(0).getInt(0)); + results.get(0).getInt(0)); + } + @Test + public void testRegisterStringBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_string_16", DataTypes.StringType, 16); - List results3 = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); - Assert.assertEquals(1, results3.size()); + List results = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterCharBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_char_16", new CharType(5), 16); + List results = spark.sql("SELECT iceberg_bucket_char_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterVarCharBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_varchar_16", new VarcharType(5), 16); + List results = spark.sql("SELECT iceberg_bucket_varchar_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results3.get(0).getInt(0)); + results.get(0).getInt(0)); + } + + @Test + public void testRegisterDateBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_date_16", DataTypes.DateType, 16); + List results = spark.sql("SELECT iceberg_bucket_date_16(DATE '2021-06-30')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.DateType.get(), 16) + .apply(DateTimeUtils.fromJavaDate(Date.valueOf("2021-06-30"))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterTimestampBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_timestamp_16", DataTypes.TimestampType, 16); + List results = + spark.sql("SELECT iceberg_bucket_timestamp_16(TIMESTAMP '2021-06-30 00:00:00.000')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.TimestampType.withZone(), 16) + .apply(DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2021-06-30 00:00:00.000"))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterBinaryBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_binary_16", DataTypes.BinaryType, 16); + List results = + spark.sql("SELECT iceberg_bucket_binary_16(X'0020001F')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.BinaryType.get(), 16) + .apply(ByteBuffer.wrap((new byte[]{0x00, 0x20, 0x00, 0x1F}))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterDecimalBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_decimal_binary_16", new DecimalType(4, 2), 16); + List results = + spark.sql("SELECT iceberg_decimal_binary_16(11.11)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.DecimalType.of(4, 2), 16) + .apply(new BigDecimal("11.11")), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterBooleanBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_boolean_16", DataTypes.BooleanType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: boolean"); + } + + @Test + public void testRegisterDoubleBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_double_16", DataTypes.DoubleType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: double"); + } + + @Test + public void testRegisterFloatBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_float_16", DataTypes.FloatType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: float"); } } From 5b7201ff67e2656675baaf9502c8dcf9dff77f00 Mon Sep 17 00:00:00 2001 From: zhangchen Date: Mon, 25 Oct 2021 18:21:03 +0800 Subject: [PATCH 2/7] v2.4 --- .../apache/iceberg/spark/IcebergSpark.java | 3 +- .../iceberg/spark/SparkValueConverter.java | 43 ++++++ .../spark/source/TestIcebergSpark.java | 132 +++++++++++++++++- 3 files changed, 170 insertions(+), 8 deletions(-) diff --git a/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java b/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java index ac659f6c7b13..6924851cefd0 100644 --- a/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java +++ b/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java @@ -34,6 +34,7 @@ public static void registerBucketUDF(SparkSession session, String funcName, Data SparkTypeToType typeConverter = new SparkTypeToType(); Type sourceIcebergType = typeConverter.atomic(sourceType); Transform bucket = Transforms.bucket(sourceIcebergType, numBuckets); - session.udf().register(funcName, bucket::apply, DataTypes.IntegerType); + session.udf().register(funcName, + value -> bucket.apply(SparkValueConverter.convertAtomicValue(sourceType, value)), DataTypes.IntegerType); } } diff --git a/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java b/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java index 92c812a9b979..a0487b3b3963 100644 --- a/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java +++ b/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java @@ -33,6 +33,21 @@ import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.catalyst.util.DateTimeUtils; +import org.apache.spark.sql.types.BinaryType; +import org.apache.spark.sql.types.BooleanType; +import org.apache.spark.sql.types.ByteType; +import org.apache.spark.sql.types.CharType; +import org.apache.spark.sql.types.DataType; +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; +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.TimestampType; +import org.apache.spark.sql.types.VarcharType; /** * A utility class that converts Spark values to Iceberg's internal representation. @@ -117,4 +132,32 @@ private static Record convert(Types.StructType struct, Row row) { } return record; } + + @SuppressWarnings("checkstyle:CyclomaticComplexity") + public static Object convertAtomicValue(DataType atomic, Object object) { + if (atomic instanceof BooleanType || + atomic instanceof IntegerType || + atomic instanceof LongType || + atomic instanceof StringType || + atomic instanceof FloatType || + atomic instanceof DoubleType || + atomic instanceof DecimalType || + atomic instanceof CharType || + atomic instanceof VarcharType) { + return object; + } else if (atomic instanceof ShortType) { + return ((Short) object).intValue(); + } else if (atomic instanceof ByteType) { + return ((Byte) object).intValue(); + } else if (atomic instanceof DateType) { + return DateTimeUtils.fromJavaDate((Date) object); + } else if (atomic instanceof TimestampType) { + return DateTimeUtils.fromJavaTimestamp((Timestamp) object); + } else if (atomic instanceof BinaryType) { + return ByteBuffer.wrap((byte[]) object); + } + + throw new UnsupportedOperationException( + "Not a supported type: " + atomic.catalogString()); + } } diff --git a/spark/v2.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java b/spark/v2.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java index e83709a3b2a9..39ff681ac74e 100644 --- a/spark/v2.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java +++ b/spark/v2.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java @@ -19,13 +19,22 @@ package org.apache.iceberg.spark.source; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.sql.Date; +import java.sql.Timestamp; import java.util.List; import org.apache.iceberg.spark.IcebergSpark; import org.apache.iceberg.transforms.Transforms; import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.util.DateTimeUtils; +import org.apache.spark.sql.types.CharType; import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.DecimalType; +import org.apache.spark.sql.types.VarcharType; +import org.assertj.core.api.Assertions; import org.junit.AfterClass; import org.junit.Assert; import org.junit.BeforeClass; @@ -48,23 +57,132 @@ public static void stopSpark() { } @Test - public void testRegisterBucketUDF() { + public void testRegisterIntegerBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_int_16", DataTypes.IntegerType, 16); List results = spark.sql("SELECT iceberg_bucket_int_16(1)").collectAsList(); Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), results.get(0).getInt(0)); + } + + @Test + public void testRegisterShortBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_short_16", DataTypes.ShortType, 16); + List results = spark.sql("SELECT iceberg_bucket_short_16(1S)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterByteBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_byte_16", DataTypes.ByteType, 16); + List results = spark.sql("SELECT iceberg_bucket_byte_16(1Y)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), + results.get(0).getInt(0)); + } + @Test + public void testRegisterLongBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_long_16", DataTypes.LongType, 16); - List results2 = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); - Assert.assertEquals(1, results2.size()); + List results = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); + Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.LongType.get(), 16).apply(1L), - results2.get(0).getInt(0)); + results.get(0).getInt(0)); + } + @Test + public void testRegisterStringBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_string_16", DataTypes.StringType, 16); - List results3 = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); - Assert.assertEquals(1, results3.size()); + List results = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterCharBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_char_16", new CharType(5), 16); + List results = spark.sql("SELECT iceberg_bucket_char_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterVarCharBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_varchar_16", new VarcharType(5), 16); + List results = spark.sql("SELECT iceberg_bucket_varchar_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results3.get(0).getInt(0)); + results.get(0).getInt(0)); + } + + @Test + public void testRegisterDateBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_date_16", DataTypes.DateType, 16); + List results = spark.sql("SELECT iceberg_bucket_date_16(DATE '2021-06-30')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.DateType.get(), 16) + .apply(DateTimeUtils.fromJavaDate(Date.valueOf("2021-06-30"))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterTimestampBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_timestamp_16", DataTypes.TimestampType, 16); + List results = + spark.sql("SELECT iceberg_bucket_timestamp_16(TIMESTAMP '2021-06-30 00:00:00.000')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.TimestampType.withZone(), 16) + .apply(DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2021-06-30 00:00:00.000"))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterBinaryBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_binary_16", DataTypes.BinaryType, 16); + List results = + spark.sql("SELECT iceberg_bucket_binary_16(X'0020001F')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.BinaryType.get(), 16) + .apply(ByteBuffer.wrap((new byte[]{0x00, 0x20, 0x00, 0x1F}))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterDecimalBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_decimal_binary_16", new DecimalType(4, 2), 16); + List results = + spark.sql("SELECT iceberg_decimal_binary_16(11.11)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.DecimalType.of(4, 2), 16) + .apply(new BigDecimal("11.11")), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterBooleanBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_boolean_16", DataTypes.BooleanType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: boolean"); + } + + @Test + public void testRegisterDoubleBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_double_16", DataTypes.DoubleType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: double"); + } + + @Test + public void testRegisterFloatBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_float_16", DataTypes.FloatType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: float"); } } From 0bc7172c44ca61e98783f698306bc2140b385887 Mon Sep 17 00:00:00 2001 From: zhangchen Date: Mon, 25 Oct 2021 18:21:16 +0800 Subject: [PATCH 3/7] v3.2 --- .../apache/iceberg/spark/IcebergSpark.java | 3 +- .../iceberg/spark/SparkValueConverter.java | 43 ++++++ .../spark/source/TestIcebergSpark.java | 132 +++++++++++++++++- 3 files changed, 170 insertions(+), 8 deletions(-) diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java index ac659f6c7b13..6924851cefd0 100644 --- a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java +++ b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java @@ -34,6 +34,7 @@ public static void registerBucketUDF(SparkSession session, String funcName, Data SparkTypeToType typeConverter = new SparkTypeToType(); Type sourceIcebergType = typeConverter.atomic(sourceType); Transform bucket = Transforms.bucket(sourceIcebergType, numBuckets); - session.udf().register(funcName, bucket::apply, DataTypes.IntegerType); + session.udf().register(funcName, + value -> bucket.apply(SparkValueConverter.convertAtomicValue(sourceType, value)), DataTypes.IntegerType); } } diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java index 92c812a9b979..a0487b3b3963 100644 --- a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java +++ b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java @@ -33,6 +33,21 @@ import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.catalyst.util.DateTimeUtils; +import org.apache.spark.sql.types.BinaryType; +import org.apache.spark.sql.types.BooleanType; +import org.apache.spark.sql.types.ByteType; +import org.apache.spark.sql.types.CharType; +import org.apache.spark.sql.types.DataType; +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; +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.TimestampType; +import org.apache.spark.sql.types.VarcharType; /** * A utility class that converts Spark values to Iceberg's internal representation. @@ -117,4 +132,32 @@ private static Record convert(Types.StructType struct, Row row) { } return record; } + + @SuppressWarnings("checkstyle:CyclomaticComplexity") + public static Object convertAtomicValue(DataType atomic, Object object) { + if (atomic instanceof BooleanType || + atomic instanceof IntegerType || + atomic instanceof LongType || + atomic instanceof StringType || + atomic instanceof FloatType || + atomic instanceof DoubleType || + atomic instanceof DecimalType || + atomic instanceof CharType || + atomic instanceof VarcharType) { + return object; + } else if (atomic instanceof ShortType) { + return ((Short) object).intValue(); + } else if (atomic instanceof ByteType) { + return ((Byte) object).intValue(); + } else if (atomic instanceof DateType) { + return DateTimeUtils.fromJavaDate((Date) object); + } else if (atomic instanceof TimestampType) { + return DateTimeUtils.fromJavaTimestamp((Timestamp) object); + } else if (atomic instanceof BinaryType) { + return ByteBuffer.wrap((byte[]) object); + } + + throw new UnsupportedOperationException( + "Not a supported type: " + atomic.catalogString()); + } } diff --git a/spark/v3.2/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java b/spark/v3.2/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java index e83709a3b2a9..39ff681ac74e 100644 --- a/spark/v3.2/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java +++ b/spark/v3.2/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java @@ -19,13 +19,22 @@ package org.apache.iceberg.spark.source; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.sql.Date; +import java.sql.Timestamp; import java.util.List; import org.apache.iceberg.spark.IcebergSpark; import org.apache.iceberg.transforms.Transforms; import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.util.DateTimeUtils; +import org.apache.spark.sql.types.CharType; import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.DecimalType; +import org.apache.spark.sql.types.VarcharType; +import org.assertj.core.api.Assertions; import org.junit.AfterClass; import org.junit.Assert; import org.junit.BeforeClass; @@ -48,23 +57,132 @@ public static void stopSpark() { } @Test - public void testRegisterBucketUDF() { + public void testRegisterIntegerBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_int_16", DataTypes.IntegerType, 16); List results = spark.sql("SELECT iceberg_bucket_int_16(1)").collectAsList(); Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), results.get(0).getInt(0)); + } + + @Test + public void testRegisterShortBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_short_16", DataTypes.ShortType, 16); + List results = spark.sql("SELECT iceberg_bucket_short_16(1S)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterByteBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_byte_16", DataTypes.ByteType, 16); + List results = spark.sql("SELECT iceberg_bucket_byte_16(1Y)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), + results.get(0).getInt(0)); + } + @Test + public void testRegisterLongBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_long_16", DataTypes.LongType, 16); - List results2 = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); - Assert.assertEquals(1, results2.size()); + List results = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); + Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.LongType.get(), 16).apply(1L), - results2.get(0).getInt(0)); + results.get(0).getInt(0)); + } + @Test + public void testRegisterStringBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_string_16", DataTypes.StringType, 16); - List results3 = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); - Assert.assertEquals(1, results3.size()); + List results = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterCharBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_char_16", new CharType(5), 16); + List results = spark.sql("SELECT iceberg_bucket_char_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterVarCharBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_varchar_16", new VarcharType(5), 16); + List results = spark.sql("SELECT iceberg_bucket_varchar_16('hello')").collectAsList(); + Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results3.get(0).getInt(0)); + results.get(0).getInt(0)); + } + + @Test + public void testRegisterDateBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_date_16", DataTypes.DateType, 16); + List results = spark.sql("SELECT iceberg_bucket_date_16(DATE '2021-06-30')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.DateType.get(), 16) + .apply(DateTimeUtils.fromJavaDate(Date.valueOf("2021-06-30"))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterTimestampBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_timestamp_16", DataTypes.TimestampType, 16); + List results = + spark.sql("SELECT iceberg_bucket_timestamp_16(TIMESTAMP '2021-06-30 00:00:00.000')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.TimestampType.withZone(), 16) + .apply(DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2021-06-30 00:00:00.000"))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterBinaryBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_binary_16", DataTypes.BinaryType, 16); + List results = + spark.sql("SELECT iceberg_bucket_binary_16(X'0020001F')").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.BinaryType.get(), 16) + .apply(ByteBuffer.wrap((new byte[]{0x00, 0x20, 0x00, 0x1F}))), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterDecimalBucketUDF() { + IcebergSpark.registerBucketUDF(spark, "iceberg_decimal_binary_16", new DecimalType(4, 2), 16); + List results = + spark.sql("SELECT iceberg_decimal_binary_16(11.11)").collectAsList(); + Assert.assertEquals(1, results.size()); + Assert.assertEquals((int) Transforms.bucket(Types.DecimalType.of(4, 2), 16) + .apply(new BigDecimal("11.11")), + results.get(0).getInt(0)); + } + + @Test + public void testRegisterBooleanBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_boolean_16", DataTypes.BooleanType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: boolean"); + } + + @Test + public void testRegisterDoubleBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_double_16", DataTypes.DoubleType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: double"); + } + + @Test + public void testRegisterFloatBucketUDF() { + Assertions.assertThatThrownBy(() -> + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_float_16", DataTypes.FloatType, 16)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Cannot bucket by type: float"); } } From 4c306b020785c397cae2770761cd2f1aa10e785c Mon Sep 17 00:00:00 2001 From: zhangchen Date: Tue, 26 Oct 2021 00:41:56 +0800 Subject: [PATCH 4/7] Revert "v3.0" This reverts commit 6919b32e527351a564bab6a62c5c7fe8e882f31e. --- .../apache/iceberg/spark/IcebergSpark.java | 3 +- .../iceberg/spark/SparkValueConverter.java | 43 ------ .../spark/source/TestIcebergSpark.java | 132 +----------------- 3 files changed, 8 insertions(+), 170 deletions(-) diff --git a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java index 6924851cefd0..ac659f6c7b13 100644 --- a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java +++ b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java @@ -34,7 +34,6 @@ public static void registerBucketUDF(SparkSession session, String funcName, Data SparkTypeToType typeConverter = new SparkTypeToType(); Type sourceIcebergType = typeConverter.atomic(sourceType); Transform bucket = Transforms.bucket(sourceIcebergType, numBuckets); - session.udf().register(funcName, - value -> bucket.apply(SparkValueConverter.convertAtomicValue(sourceType, value)), DataTypes.IntegerType); + session.udf().register(funcName, bucket::apply, DataTypes.IntegerType); } } diff --git a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java index 35513c12f8d5..92c812a9b979 100644 --- a/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java +++ b/spark/v3.0/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java @@ -33,21 +33,6 @@ import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.catalyst.util.DateTimeUtils; -import org.apache.spark.sql.types.BinaryType; -import org.apache.spark.sql.types.BooleanType; -import org.apache.spark.sql.types.ByteType; -import org.apache.spark.sql.types.CharType; -import org.apache.spark.sql.types.DataType; -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; -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.TimestampType; -import org.apache.spark.sql.types.VarcharType; /** * A utility class that converts Spark values to Iceberg's internal representation. @@ -132,32 +117,4 @@ private static Record convert(Types.StructType struct, Row row) { } return record; } - - @SuppressWarnings("checkstyle:CyclomaticComplexity") - public static Object convertAtomicValue(DataType atomic, Object object) { - if (atomic instanceof BooleanType || - atomic instanceof IntegerType || - atomic instanceof LongType || - atomic instanceof StringType || - atomic instanceof FloatType || - atomic instanceof DoubleType || - atomic instanceof DecimalType || - atomic instanceof CharType || - atomic instanceof VarcharType) { - return object; - } else if (atomic instanceof ShortType) { - return ((Short) object).intValue(); - } else if (atomic instanceof ByteType) { - return ((Byte) object).intValue(); - } else if (atomic instanceof DateType) { - return DateTimeUtils.fromJavaDate((Date) object); - } else if (atomic instanceof TimestampType) { - return DateTimeUtils.fromJavaTimestamp((Timestamp) object); - } else if (atomic instanceof BinaryType) { - return ByteBuffer.wrap((byte[]) object); - } - - throw new UnsupportedOperationException( - "Not a supported type: " + atomic.catalogString()); - } } diff --git a/spark/v3.0/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java b/spark/v3.0/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java index 39ff681ac74e..e83709a3b2a9 100644 --- a/spark/v3.0/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java +++ b/spark/v3.0/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java @@ -19,22 +19,13 @@ package org.apache.iceberg.spark.source; -import java.math.BigDecimal; -import java.nio.ByteBuffer; -import java.sql.Date; -import java.sql.Timestamp; import java.util.List; import org.apache.iceberg.spark.IcebergSpark; import org.apache.iceberg.transforms.Transforms; import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; -import org.apache.spark.sql.catalyst.util.DateTimeUtils; -import org.apache.spark.sql.types.CharType; import org.apache.spark.sql.types.DataTypes; -import org.apache.spark.sql.types.DecimalType; -import org.apache.spark.sql.types.VarcharType; -import org.assertj.core.api.Assertions; import org.junit.AfterClass; import org.junit.Assert; import org.junit.BeforeClass; @@ -57,132 +48,23 @@ public static void stopSpark() { } @Test - public void testRegisterIntegerBucketUDF() { + public void testRegisterBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_int_16", DataTypes.IntegerType, 16); List results = spark.sql("SELECT iceberg_bucket_int_16(1)").collectAsList(); Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), results.get(0).getInt(0)); - } - - @Test - public void testRegisterShortBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_short_16", DataTypes.ShortType, 16); - List results = spark.sql("SELECT iceberg_bucket_short_16(1S)").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterByteBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_byte_16", DataTypes.ByteType, 16); - List results = spark.sql("SELECT iceberg_bucket_byte_16(1Y)").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), - results.get(0).getInt(0)); - } - @Test - public void testRegisterLongBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_long_16", DataTypes.LongType, 16); - List results = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); - Assert.assertEquals(1, results.size()); + List results2 = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); + Assert.assertEquals(1, results2.size()); Assert.assertEquals((int) Transforms.bucket(Types.LongType.get(), 16).apply(1L), - results.get(0).getInt(0)); - } + results2.get(0).getInt(0)); - @Test - public void testRegisterStringBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_string_16", DataTypes.StringType, 16); - List results = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterCharBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_char_16", new CharType(5), 16); - List results = spark.sql("SELECT iceberg_bucket_char_16('hello')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterVarCharBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_varchar_16", new VarcharType(5), 16); - List results = spark.sql("SELECT iceberg_bucket_varchar_16('hello')").collectAsList(); - Assert.assertEquals(1, results.size()); + List results3 = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); + Assert.assertEquals(1, results3.size()); Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterDateBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_date_16", DataTypes.DateType, 16); - List results = spark.sql("SELECT iceberg_bucket_date_16(DATE '2021-06-30')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.DateType.get(), 16) - .apply(DateTimeUtils.fromJavaDate(Date.valueOf("2021-06-30"))), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterTimestampBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_timestamp_16", DataTypes.TimestampType, 16); - List results = - spark.sql("SELECT iceberg_bucket_timestamp_16(TIMESTAMP '2021-06-30 00:00:00.000')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.TimestampType.withZone(), 16) - .apply(DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2021-06-30 00:00:00.000"))), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterBinaryBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_binary_16", DataTypes.BinaryType, 16); - List results = - spark.sql("SELECT iceberg_bucket_binary_16(X'0020001F')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.BinaryType.get(), 16) - .apply(ByteBuffer.wrap((new byte[]{0x00, 0x20, 0x00, 0x1F}))), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterDecimalBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_decimal_binary_16", new DecimalType(4, 2), 16); - List results = - spark.sql("SELECT iceberg_decimal_binary_16(11.11)").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.DecimalType.of(4, 2), 16) - .apply(new BigDecimal("11.11")), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterBooleanBucketUDF() { - Assertions.assertThatThrownBy(() -> - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_boolean_16", DataTypes.BooleanType, 16)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Cannot bucket by type: boolean"); - } - - @Test - public void testRegisterDoubleBucketUDF() { - Assertions.assertThatThrownBy(() -> - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_double_16", DataTypes.DoubleType, 16)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Cannot bucket by type: double"); - } - - @Test - public void testRegisterFloatBucketUDF() { - Assertions.assertThatThrownBy(() -> - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_float_16", DataTypes.FloatType, 16)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Cannot bucket by type: float"); + results3.get(0).getInt(0)); } } From 078dfe45e9bec99eb533d418accae7326263161a Mon Sep 17 00:00:00 2001 From: zhangchen Date: Tue, 26 Oct 2021 00:42:03 +0800 Subject: [PATCH 5/7] Revert "v2.4" This reverts commit 5b7201ff67e2656675baaf9502c8dcf9dff77f00. --- .../apache/iceberg/spark/IcebergSpark.java | 3 +- .../iceberg/spark/SparkValueConverter.java | 43 ------ .../spark/source/TestIcebergSpark.java | 132 +----------------- 3 files changed, 8 insertions(+), 170 deletions(-) diff --git a/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java b/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java index 6924851cefd0..ac659f6c7b13 100644 --- a/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java +++ b/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java @@ -34,7 +34,6 @@ public static void registerBucketUDF(SparkSession session, String funcName, Data SparkTypeToType typeConverter = new SparkTypeToType(); Type sourceIcebergType = typeConverter.atomic(sourceType); Transform bucket = Transforms.bucket(sourceIcebergType, numBuckets); - session.udf().register(funcName, - value -> bucket.apply(SparkValueConverter.convertAtomicValue(sourceType, value)), DataTypes.IntegerType); + session.udf().register(funcName, bucket::apply, DataTypes.IntegerType); } } diff --git a/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java b/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java index a0487b3b3963..92c812a9b979 100644 --- a/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java +++ b/spark/v2.4/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java @@ -33,21 +33,6 @@ import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.catalyst.util.DateTimeUtils; -import org.apache.spark.sql.types.BinaryType; -import org.apache.spark.sql.types.BooleanType; -import org.apache.spark.sql.types.ByteType; -import org.apache.spark.sql.types.CharType; -import org.apache.spark.sql.types.DataType; -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; -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.TimestampType; -import org.apache.spark.sql.types.VarcharType; /** * A utility class that converts Spark values to Iceberg's internal representation. @@ -132,32 +117,4 @@ private static Record convert(Types.StructType struct, Row row) { } return record; } - - @SuppressWarnings("checkstyle:CyclomaticComplexity") - public static Object convertAtomicValue(DataType atomic, Object object) { - if (atomic instanceof BooleanType || - atomic instanceof IntegerType || - atomic instanceof LongType || - atomic instanceof StringType || - atomic instanceof FloatType || - atomic instanceof DoubleType || - atomic instanceof DecimalType || - atomic instanceof CharType || - atomic instanceof VarcharType) { - return object; - } else if (atomic instanceof ShortType) { - return ((Short) object).intValue(); - } else if (atomic instanceof ByteType) { - return ((Byte) object).intValue(); - } else if (atomic instanceof DateType) { - return DateTimeUtils.fromJavaDate((Date) object); - } else if (atomic instanceof TimestampType) { - return DateTimeUtils.fromJavaTimestamp((Timestamp) object); - } else if (atomic instanceof BinaryType) { - return ByteBuffer.wrap((byte[]) object); - } - - throw new UnsupportedOperationException( - "Not a supported type: " + atomic.catalogString()); - } } diff --git a/spark/v2.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java b/spark/v2.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java index 39ff681ac74e..e83709a3b2a9 100644 --- a/spark/v2.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java +++ b/spark/v2.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java @@ -19,22 +19,13 @@ package org.apache.iceberg.spark.source; -import java.math.BigDecimal; -import java.nio.ByteBuffer; -import java.sql.Date; -import java.sql.Timestamp; import java.util.List; import org.apache.iceberg.spark.IcebergSpark; import org.apache.iceberg.transforms.Transforms; import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; -import org.apache.spark.sql.catalyst.util.DateTimeUtils; -import org.apache.spark.sql.types.CharType; import org.apache.spark.sql.types.DataTypes; -import org.apache.spark.sql.types.DecimalType; -import org.apache.spark.sql.types.VarcharType; -import org.assertj.core.api.Assertions; import org.junit.AfterClass; import org.junit.Assert; import org.junit.BeforeClass; @@ -57,132 +48,23 @@ public static void stopSpark() { } @Test - public void testRegisterIntegerBucketUDF() { + public void testRegisterBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_int_16", DataTypes.IntegerType, 16); List results = spark.sql("SELECT iceberg_bucket_int_16(1)").collectAsList(); Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), results.get(0).getInt(0)); - } - - @Test - public void testRegisterShortBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_short_16", DataTypes.ShortType, 16); - List results = spark.sql("SELECT iceberg_bucket_short_16(1S)").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterByteBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_byte_16", DataTypes.ByteType, 16); - List results = spark.sql("SELECT iceberg_bucket_byte_16(1Y)").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.IntegerType.get(), 16).apply(1), - results.get(0).getInt(0)); - } - @Test - public void testRegisterLongBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_long_16", DataTypes.LongType, 16); - List results = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); - Assert.assertEquals(1, results.size()); + List results2 = spark.sql("SELECT iceberg_bucket_long_16(1L)").collectAsList(); + Assert.assertEquals(1, results2.size()); Assert.assertEquals((int) Transforms.bucket(Types.LongType.get(), 16).apply(1L), - results.get(0).getInt(0)); - } + results2.get(0).getInt(0)); - @Test - public void testRegisterStringBucketUDF() { IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_string_16", DataTypes.StringType, 16); - List results = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterCharBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_char_16", new CharType(5), 16); - List results = spark.sql("SELECT iceberg_bucket_char_16('hello')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterVarCharBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_varchar_16", new VarcharType(5), 16); - List results = spark.sql("SELECT iceberg_bucket_varchar_16('hello')").collectAsList(); - Assert.assertEquals(1, results.size()); + List results3 = spark.sql("SELECT iceberg_bucket_string_16('hello')").collectAsList(); + Assert.assertEquals(1, results3.size()); Assert.assertEquals((int) Transforms.bucket(Types.StringType.get(), 16).apply("hello"), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterDateBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_date_16", DataTypes.DateType, 16); - List results = spark.sql("SELECT iceberg_bucket_date_16(DATE '2021-06-30')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.DateType.get(), 16) - .apply(DateTimeUtils.fromJavaDate(Date.valueOf("2021-06-30"))), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterTimestampBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_timestamp_16", DataTypes.TimestampType, 16); - List results = - spark.sql("SELECT iceberg_bucket_timestamp_16(TIMESTAMP '2021-06-30 00:00:00.000')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.TimestampType.withZone(), 16) - .apply(DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2021-06-30 00:00:00.000"))), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterBinaryBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_binary_16", DataTypes.BinaryType, 16); - List results = - spark.sql("SELECT iceberg_bucket_binary_16(X'0020001F')").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.BinaryType.get(), 16) - .apply(ByteBuffer.wrap((new byte[]{0x00, 0x20, 0x00, 0x1F}))), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterDecimalBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_decimal_binary_16", new DecimalType(4, 2), 16); - List results = - spark.sql("SELECT iceberg_decimal_binary_16(11.11)").collectAsList(); - Assert.assertEquals(1, results.size()); - Assert.assertEquals((int) Transforms.bucket(Types.DecimalType.of(4, 2), 16) - .apply(new BigDecimal("11.11")), - results.get(0).getInt(0)); - } - - @Test - public void testRegisterBooleanBucketUDF() { - Assertions.assertThatThrownBy(() -> - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_boolean_16", DataTypes.BooleanType, 16)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Cannot bucket by type: boolean"); - } - - @Test - public void testRegisterDoubleBucketUDF() { - Assertions.assertThatThrownBy(() -> - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_double_16", DataTypes.DoubleType, 16)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Cannot bucket by type: double"); - } - - @Test - public void testRegisterFloatBucketUDF() { - Assertions.assertThatThrownBy(() -> - IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_float_16", DataTypes.FloatType, 16)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("Cannot bucket by type: float"); + results3.get(0).getInt(0)); } } From b87e1265f68fa6ebb82f838e2a4c8909e4a76c71 Mon Sep 17 00:00:00 2001 From: zhangchen Date: Tue, 26 Oct 2021 00:54:57 +0800 Subject: [PATCH 6/7] updating convert method --- .../apache/iceberg/spark/IcebergSpark.java | 2 +- .../iceberg/spark/SparkValueConverter.java | 46 +------------------ 2 files changed, 3 insertions(+), 45 deletions(-) diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java index 6924851cefd0..862626d0cd6d 100644 --- a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java +++ b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/IcebergSpark.java @@ -35,6 +35,6 @@ public static void registerBucketUDF(SparkSession session, String funcName, Data Type sourceIcebergType = typeConverter.atomic(sourceType); Transform bucket = Transforms.bucket(sourceIcebergType, numBuckets); session.udf().register(funcName, - value -> bucket.apply(SparkValueConverter.convertAtomicValue(sourceType, value)), DataTypes.IntegerType); + value -> bucket.apply(SparkValueConverter.convert(sourceIcebergType, value)), DataTypes.IntegerType); } } diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java index a0487b3b3963..ef453c0cef2b 100644 --- a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java +++ b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkValueConverter.java @@ -33,21 +33,6 @@ import org.apache.iceberg.types.Types; import org.apache.spark.sql.Row; import org.apache.spark.sql.catalyst.util.DateTimeUtils; -import org.apache.spark.sql.types.BinaryType; -import org.apache.spark.sql.types.BooleanType; -import org.apache.spark.sql.types.ByteType; -import org.apache.spark.sql.types.CharType; -import org.apache.spark.sql.types.DataType; -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; -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.TimestampType; -import org.apache.spark.sql.types.VarcharType; /** * A utility class that converts Spark values to Iceberg's internal representation. @@ -94,8 +79,9 @@ public static Object convert(Type type, Object object) { return DateTimeUtils.fromJavaTimestamp((Timestamp) object); case BINARY: return ByteBuffer.wrap((byte[]) object); - case BOOLEAN: case INTEGER: + return ((Number) object).intValue(); + case BOOLEAN: case LONG: case FLOAT: case DOUBLE: @@ -132,32 +118,4 @@ private static Record convert(Types.StructType struct, Row row) { } return record; } - - @SuppressWarnings("checkstyle:CyclomaticComplexity") - public static Object convertAtomicValue(DataType atomic, Object object) { - if (atomic instanceof BooleanType || - atomic instanceof IntegerType || - atomic instanceof LongType || - atomic instanceof StringType || - atomic instanceof FloatType || - atomic instanceof DoubleType || - atomic instanceof DecimalType || - atomic instanceof CharType || - atomic instanceof VarcharType) { - return object; - } else if (atomic instanceof ShortType) { - return ((Short) object).intValue(); - } else if (atomic instanceof ByteType) { - return ((Byte) object).intValue(); - } else if (atomic instanceof DateType) { - return DateTimeUtils.fromJavaDate((Date) object); - } else if (atomic instanceof TimestampType) { - return DateTimeUtils.fromJavaTimestamp((Timestamp) object); - } else if (atomic instanceof BinaryType) { - return ByteBuffer.wrap((byte[]) object); - } - - throw new UnsupportedOperationException( - "Not a supported type: " + atomic.catalogString()); - } } From cf0db4be39c55e9f7cc7fb4506cba5f05917ad3a Mon Sep 17 00:00:00 2001 From: zhangchen Date: Tue, 26 Oct 2021 11:32:40 +0800 Subject: [PATCH 7/7] fix spelling mistakes --- .../org/apache/iceberg/spark/source/TestIcebergSpark.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/spark/v3.2/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java b/spark/v3.2/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java index 39ff681ac74e..5803390a66f4 100644 --- a/spark/v3.2/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java +++ b/spark/v3.2/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSpark.java @@ -153,9 +153,9 @@ public void testRegisterBinaryBucketUDF() { @Test public void testRegisterDecimalBucketUDF() { - IcebergSpark.registerBucketUDF(spark, "iceberg_decimal_binary_16", new DecimalType(4, 2), 16); + IcebergSpark.registerBucketUDF(spark, "iceberg_bucket_decimal_16", new DecimalType(4, 2), 16); List results = - spark.sql("SELECT iceberg_decimal_binary_16(11.11)").collectAsList(); + spark.sql("SELECT iceberg_bucket_decimal_16(11.11)").collectAsList(); Assert.assertEquals(1, results.size()); Assert.assertEquals((int) Transforms.bucket(Types.DecimalType.of(4, 2), 16) .apply(new BigDecimal("11.11")),