You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by ao...@apache.org on 2022/11/24 00:10:00 UTC
[iceberg] branch master updated: Spark 3.3: Add years transform function (#6207)
This is an automated email from the ASF dual-hosted git repository.
aokolnychyi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/master by this push:
new 3a95e6afe0 Spark 3.3: Add years transform function (#6207)
3a95e6afe0 is described below
commit 3a95e6afe032fc2a6e6fc5055356e5916f3515f8
Author: Anton Okolnychyi <ao...@apple.com>
AuthorDate: Wed Nov 23 16:09:55 2022 -0800
Spark 3.3: Add years transform function (#6207)
---
.../iceberg/spark/functions/SparkFunctions.java | 1 +
.../spark/functions/UnaryUnboundFunction.java | 44 ++++++++
.../iceberg/spark/functions/YearsFunction.java | 118 +++++++++++++++++++++
.../iceberg/spark/sql/TestSparkYearsFunction.java | 117 ++++++++++++++++++++
4 files changed, 280 insertions(+)
diff --git a/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/SparkFunctions.java b/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/SparkFunctions.java
index 7ca157ad40..aa71a02a22 100644
--- a/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/SparkFunctions.java
+++ b/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/SparkFunctions.java
@@ -32,6 +32,7 @@ public class SparkFunctions {
private static final Map<String, UnboundFunction> FUNCTIONS =
ImmutableMap.of(
"iceberg_version", new IcebergVersionFunction(),
+ "years", new YearsFunction(),
"bucket", new BucketFunction(),
"truncate", new TruncateFunction());
diff --git a/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/UnaryUnboundFunction.java b/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/UnaryUnboundFunction.java
new file mode 100644
index 0000000000..9003c68919
--- /dev/null
+++ b/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/UnaryUnboundFunction.java
@@ -0,0 +1,44 @@
+/*
+ * 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.iceberg.spark.functions;
+
+import org.apache.spark.sql.connector.catalog.functions.BoundFunction;
+import org.apache.spark.sql.connector.catalog.functions.UnboundFunction;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.StructType;
+
+/** An unbound function that accepts only one argument */
+abstract class UnaryUnboundFunction implements UnboundFunction {
+
+ @Override
+ public BoundFunction bind(StructType inputType) {
+ DataType valueType = valueType(inputType);
+ return doBind(valueType);
+ }
+
+ protected abstract BoundFunction doBind(DataType valueType);
+
+ private DataType valueType(StructType inputType) {
+ if (inputType.size() != 1) {
+ throw new UnsupportedOperationException("Wrong number of inputs (expected value)");
+ }
+
+ return inputType.fields()[0].dataType();
+ }
+}
diff --git a/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/YearsFunction.java b/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/YearsFunction.java
new file mode 100644
index 0000000000..779e7a28ca
--- /dev/null
+++ b/spark/v3.3/spark/src/main/java/org/apache/iceberg/spark/functions/YearsFunction.java
@@ -0,0 +1,118 @@
+/*
+ * 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.iceberg.spark.functions;
+
+import org.apache.iceberg.util.DateTimeUtil;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.catalog.functions.BoundFunction;
+import org.apache.spark.sql.connector.catalog.functions.ScalarFunction;
+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.TimestampType;
+
+/**
+ * A Spark function implementation for the Iceberg year transform.
+ *
+ * <p>Example usage: {@code SELECT system.years('source_col')}.
+ */
+public class YearsFunction extends UnaryUnboundFunction {
+
+ @Override
+ protected BoundFunction doBind(DataType valueType) {
+ if (valueType instanceof DateType) {
+ return new DateToYearsFunction();
+ } else if (valueType instanceof TimestampType) {
+ return new TimestampToYearsFunction();
+ } else {
+ throw new UnsupportedOperationException(
+ "Expected value to be date or timestamp: " + valueType.catalogString());
+ }
+ }
+
+ @Override
+ public String description() {
+ return name()
+ + "(col) - Call Iceberg's year transform\n"
+ + " col :: source column (must be date or timestamp)";
+ }
+
+ @Override
+ public String name() {
+ return "years";
+ }
+
+ private abstract static class BaseToYearsFunction implements ScalarFunction<Integer> {
+ @Override
+ public String name() {
+ return "years";
+ }
+
+ @Override
+ public DataType resultType() {
+ return DataTypes.IntegerType;
+ }
+ }
+
+ public static class DateToYearsFunction extends BaseToYearsFunction {
+ // magic method used in codegen
+ public static int invoke(int days) {
+ return DateTimeUtil.daysToYears(days);
+ }
+
+ @Override
+ public DataType[] inputTypes() {
+ return new DataType[] {DataTypes.DateType};
+ }
+
+ @Override
+ public String canonicalName() {
+ return "iceberg.years(date)";
+ }
+
+ @Override
+ public Integer produceResult(InternalRow input) {
+ // return null for null input to match what Spark does in codegen
+ return input.isNullAt(0) ? null : invoke(input.getInt(0));
+ }
+ }
+
+ public static class TimestampToYearsFunction extends BaseToYearsFunction {
+ // magic method used in codegen
+ public static int invoke(long micros) {
+ return DateTimeUtil.microsToYears(micros);
+ }
+
+ @Override
+ public DataType[] inputTypes() {
+ return new DataType[] {DataTypes.TimestampType};
+ }
+
+ @Override
+ public String canonicalName() {
+ return "iceberg.years(timestamp)";
+ }
+
+ @Override
+ public Integer produceResult(InternalRow input) {
+ // return null for null input to match what Spark does in codegen
+ return input.isNullAt(0) ? null : invoke(input.getLong(0));
+ }
+ }
+}
diff --git a/spark/v3.3/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkYearsFunction.java b/spark/v3.3/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkYearsFunction.java
new file mode 100644
index 0000000000..d4676716a6
--- /dev/null
+++ b/spark/v3.3/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkYearsFunction.java
@@ -0,0 +1,117 @@
+/*
+ * 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.iceberg.spark.sql;
+
+import org.apache.iceberg.AssertHelpers;
+import org.apache.iceberg.spark.SparkTestBaseWithCatalog;
+import org.apache.iceberg.spark.functions.YearsFunction;
+import org.apache.spark.sql.AnalysisException;
+import org.assertj.core.api.Assertions;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestSparkYearsFunction extends SparkTestBaseWithCatalog {
+
+ @Before
+ public void useCatalog() {
+ sql("USE %s", catalogName);
+ }
+
+ @Test
+ public void testDates() {
+ Assert.assertEquals(
+ "Expected to produce 2017 - 1970 = 47",
+ 47,
+ scalarSql("SELECT system.years(date('2017-12-01'))"));
+ Assert.assertEquals(
+ "Expected to produce 1970 - 1970 = 0",
+ 0,
+ scalarSql("SELECT system.years(date('1970-01-01'))"));
+ Assert.assertEquals(
+ "Expected to produce 1969 - 1970 = -1",
+ -1,
+ scalarSql("SELECT system.years(date('1969-12-31'))"));
+ Assert.assertNull(scalarSql("SELECT system.years(CAST(null AS DATE))"));
+ }
+
+ @Test
+ public void testTimestamps() {
+ Assert.assertEquals(
+ "Expected to produce 2017 - 1970 = 47",
+ 47,
+ scalarSql("SELECT system.years(TIMESTAMP '2017-12-01 10:12:55.038194 UTC+00:00')"));
+ Assert.assertEquals(
+ "Expected to produce 1970 - 1970 = 0",
+ 0,
+ scalarSql("SELECT system.years(TIMESTAMP '1970-01-01 00:00:01.000001 UTC+00:00')"));
+ Assert.assertEquals(
+ "Expected to produce 1969 - 1970 = -1",
+ -1,
+ scalarSql("SELECT system.years(TIMESTAMP '1969-12-31 23:59:58.999999 UTC+00:00')"));
+ Assert.assertNull(scalarSql("SELECT system.years(CAST(null AS TIMESTAMP))"));
+ }
+
+ @Test
+ public void testWrongNumberOfArguments() {
+ AssertHelpers.assertThrows(
+ "Function resolution should not work with zero arguments",
+ AnalysisException.class,
+ "Function 'years' cannot process input: (): Wrong number of inputs",
+ () -> scalarSql("SELECT system.years()"));
+
+ AssertHelpers.assertThrows(
+ "Function resolution should not work with more than one argument",
+ AnalysisException.class,
+ "Function 'years' cannot process input: (date, date): Wrong number of inputs",
+ () -> scalarSql("SELECT system.years(date('1969-12-31'), date('1969-12-31'))"));
+ }
+
+ @Test
+ public void testInvalidInputTypes() {
+ AssertHelpers.assertThrows(
+ "Int type should not be coercible to date/timestamp",
+ AnalysisException.class,
+ "Function 'years' cannot process input: (int): Expected value to be date or timestamp",
+ () -> scalarSql("SELECT system.years(1)"));
+
+ AssertHelpers.assertThrows(
+ "Long type should not be coercible to date/timestamp",
+ AnalysisException.class,
+ "Function 'years' cannot process input: (bigint): Expected value to be date or timestamp",
+ () -> scalarSql("SELECT system.years(1L)"));
+ }
+
+ @Test
+ public void testThatMagicFunctionsAreInvoked() {
+ String dateValue = "date('2017-12-01')";
+ String dateTransformClass = YearsFunction.DateToYearsFunction.class.getName();
+ Assertions.assertThat(scalarSql("EXPLAIN EXTENDED SELECT system.years(%s)", dateValue))
+ .asString()
+ .isNotNull()
+ .contains("staticinvoke(class " + dateTransformClass);
+
+ String timestampValue = "TIMESTAMP '2017-12-01 10:12:55.038194 UTC+00:00'";
+ String timestampTransformClass = YearsFunction.TimestampToYearsFunction.class.getName();
+ Assertions.assertThat(scalarSql("EXPLAIN EXTENDED SELECT system.years(%s)", timestampValue))
+ .asString()
+ .isNotNull()
+ .contains("staticinvoke(class " + timestampTransformClass);
+ }
+}