This is an automated email from the ASF dual-hosted git repository. ibzib pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push: new 27739f9 [BEAM-10925] Refactor ZetaSqlJavaUdfTypeTest. new 0f955b4 Merge pull request #14462 from ibzib/java-udf-types 27739f9 is described below commit 27739f959c86018f496841bfc047b8d9e721108e Author: Kyle Weaver <kcwea...@google.com> AuthorDate: Wed Apr 7 10:19:09 2021 -0700 [BEAM-10925] Refactor ZetaSqlJavaUdfTypeTest. Remove unnecessary throws clause and register all UDFs up front. --- .../sql/zetasql/ZetaSqlJavaUdfTypeTest.java | 277 ++++++++++----------- 1 file changed, 126 insertions(+), 151 deletions(-) diff --git a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlJavaUdfTypeTest.java b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlJavaUdfTypeTest.java index e45b832..d9a9ae4 100644 --- a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlJavaUdfTypeTest.java +++ b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlJavaUdfTypeTest.java @@ -17,7 +17,6 @@ */ package org.apache.beam.sdk.extensions.sql.zetasql; -import java.lang.reflect.Method; import org.apache.beam.sdk.extensions.sql.BeamSqlUdf; import org.apache.beam.sdk.extensions.sql.impl.JdbcConnection; import org.apache.beam.sdk.extensions.sql.impl.JdbcDriver; @@ -32,6 +31,7 @@ import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.SchemaPlus; import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.Frameworks; import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.RuleSet; import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; @@ -51,12 +51,6 @@ public class ZetaSqlJavaUdfTypeTest extends ZetaSqlTestBase { @Rule public transient TestPipeline pipeline = TestPipeline.create(); @Rule public ExpectedException thrown = ExpectedException.none(); - private Method boolUdf; - private Method longUdf; - private Method stringUdf; - private Method bytesUdf; - private Method doubleUdf; - private static final TestBoundedTable table = TestBoundedTable.of( Schema.builder() @@ -111,21 +105,39 @@ public class ZetaSqlJavaUdfTypeTest extends ZetaSqlTestBase { public void setUp() throws NoSuchMethodException { initialize(); - // Add BeamJavaUdfCalcRule to planner to enable UDFs. + // Register test table. + JdbcConnection jdbcConnection = + JdbcDriver.connect( + new ReadOnlyTableProvider("table_provider", ImmutableMap.of("table", table)), + PipelineOptionsFactory.create()); + + // Register UDFs. + SchemaPlus schema = jdbcConnection.getCurrentSchemaPlus(); + schema.add( + "test_boolean", + ScalarFunctionImpl.create(BooleanIdentityFn.class.getMethod("eval", Boolean.class))); + schema.add( + "test_int64", + ScalarFunctionImpl.create(Int64IdentityFn.class.getMethod("eval", Long.class))); + schema.add( + "test_string", + ScalarFunctionImpl.create(StringIdentityFn.class.getMethod("eval", String.class))); + schema.add( + "test_bytes", + ScalarFunctionImpl.create(BytesIdentityFn.class.getMethod("eval", byte[].class))); + schema.add( + "test_float64", + ScalarFunctionImpl.create(DoubleIdentityFn.class.getMethod("eval", Double.class))); + this.config = Frameworks.newConfigBuilder(config) + .defaultSchema(schema) + // Add BeamJavaUdfCalcRule to planner to enable UDFs. .ruleSets( ZetaSQLQueryPlanner.getZetaSqlRuleSets( ImmutableList.of(BeamJavaUdfCalcRule.INSTANCE)) .toArray(new RuleSet[0])) .build(); - - // Look up UDF methods. - this.boolUdf = BooleanIdentityFn.class.getMethod("eval", Boolean.class); - this.longUdf = Int64IdentityFn.class.getMethod("eval", Long.class); - this.stringUdf = StringIdentityFn.class.getMethod("eval", String.class); - this.bytesUdf = BytesIdentityFn.class.getMethod("eval", byte[].class); - this.doubleUdf = DoubleIdentityFn.class.getMethod("eval", Double.class); } public static class BooleanIdentityFn implements BeamSqlUdf { @@ -158,19 +170,7 @@ public class ZetaSqlJavaUdfTypeTest extends ZetaSqlTestBase { } } - private void runUdfTypeTest(String query, Object result, Schema.TypeName typeName, Method udf) - throws NoSuchMethodException { - // Add UDF to Calcite schema. - JdbcConnection jdbcConnection = - JdbcDriver.connect( - new ReadOnlyTableProvider("table_provider", ImmutableMap.of("table", table)), - PipelineOptionsFactory.create()); - jdbcConnection.getCurrentSchemaPlus().add("test", ScalarFunctionImpl.create(udf)); - config = - Frameworks.newConfigBuilder(config) - .defaultSchema(jdbcConnection.getCurrentSchemaPlus()) - .build(); - + private void runUdfTypeTest(String query, Object result, Schema.TypeName typeName) { ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config); BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(query); PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode); @@ -181,293 +181,268 @@ public class ZetaSqlJavaUdfTypeTest extends ZetaSqlTestBase { } @Test - public void testTrueLiteral() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(true);", true, Schema.TypeName.BOOLEAN, boolUdf); + public void testTrueLiteral() { + runUdfTypeTest("SELECT test_boolean(true);", true, Schema.TypeName.BOOLEAN); } @Test - public void testTrueInput() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(boolean_true) FROM table;", true, Schema.TypeName.BOOLEAN, boolUdf); + public void testTrueInput() { + runUdfTypeTest("SELECT test_boolean(boolean_true) FROM table;", true, Schema.TypeName.BOOLEAN); } @Test - public void testFalseLiteral() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(false);", false, Schema.TypeName.BOOLEAN, boolUdf); + public void testFalseLiteral() { + runUdfTypeTest("SELECT test_boolean(false);", false, Schema.TypeName.BOOLEAN); } @Test - public void testFalseInput() throws NoSuchMethodException { + public void testFalseInput() { runUdfTypeTest( - "SELECT test(boolean_false) FROM table;", false, Schema.TypeName.BOOLEAN, boolUdf); + "SELECT test_boolean(boolean_false) FROM table;", false, Schema.TypeName.BOOLEAN); } @Test - public void testZeroInt64Literal() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(0);", 0L, Schema.TypeName.INT64, longUdf); + public void testZeroInt64Literal() { + runUdfTypeTest("SELECT test_int64(0);", 0L, Schema.TypeName.INT64); } @Test - public void testZeroInt64Input() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(int64_0) FROM table;", 0L, Schema.TypeName.INT64, longUdf); + public void testZeroInt64Input() { + runUdfTypeTest("SELECT test_int64(int64_0) FROM table;", 0L, Schema.TypeName.INT64); } @Test - public void testPosInt64Literal() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(123);", 123L, Schema.TypeName.INT64, longUdf); + public void testPosInt64Literal() { + runUdfTypeTest("SELECT test_int64(123);", 123L, Schema.TypeName.INT64); } @Test - public void testPosInt64Input() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(int64_pos) FROM table;", 123L, Schema.TypeName.INT64, longUdf); + public void testPosInt64Input() { + runUdfTypeTest("SELECT test_int64(int64_pos) FROM table;", 123L, Schema.TypeName.INT64); } @Test - public void testNegInt64Literal() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(-123);", -123L, Schema.TypeName.INT64, longUdf); + public void testNegInt64Literal() { + runUdfTypeTest("SELECT test_int64(-123);", -123L, Schema.TypeName.INT64); } @Test - public void testNegInt64Input() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(int64_neg) FROM table;", -123L, Schema.TypeName.INT64, longUdf); + public void testNegInt64Input() { + runUdfTypeTest("SELECT test_int64(int64_neg) FROM table;", -123L, Schema.TypeName.INT64); } @Test - public void testMaxInt64Literal() throws NoSuchMethodException { + public void testMaxInt64Literal() { runUdfTypeTest( - "SELECT test(9223372036854775807);", 9223372036854775807L, Schema.TypeName.INT64, longUdf); + "SELECT test_int64(9223372036854775807);", 9223372036854775807L, Schema.TypeName.INT64); } @Test - public void testMaxInt64Input() throws NoSuchMethodException { + public void testMaxInt64Input() { runUdfTypeTest( - "SELECT test(int64_max) FROM table;", 9223372036854775807L, Schema.TypeName.INT64, longUdf); + "SELECT test_int64(int64_max) FROM table;", 9223372036854775807L, Schema.TypeName.INT64); } @Test - public void testMinInt64Literal() throws NoSuchMethodException { + public void testMinInt64Literal() { runUdfTypeTest( - "SELECT test(-9223372036854775808);", - -9223372036854775808L, - Schema.TypeName.INT64, - longUdf); + "SELECT test_int64(-9223372036854775808);", -9223372036854775808L, Schema.TypeName.INT64); } @Test - public void testMinInt64Input() throws NoSuchMethodException { + public void testMinInt64Input() { runUdfTypeTest( - "SELECT test(int64_min) FROM table;", - -9223372036854775808L, - Schema.TypeName.INT64, - longUdf); + "SELECT test_int64(int64_min) FROM table;", -9223372036854775808L, Schema.TypeName.INT64); } @Test - public void testEmptyStringLiteral() throws NoSuchMethodException { - runUdfTypeTest("SELECT test('');", "", Schema.TypeName.STRING, stringUdf); + public void testEmptyStringLiteral() { + runUdfTypeTest("SELECT test_string('');", "", Schema.TypeName.STRING); } @Test - public void testEmptyStringInput() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(string_empty) FROM table;", "", Schema.TypeName.STRING, stringUdf); + public void testEmptyStringInput() { + runUdfTypeTest("SELECT test_string(string_empty) FROM table;", "", Schema.TypeName.STRING); } @Test - public void testAsciiStringLiteral() throws NoSuchMethodException { - runUdfTypeTest("SELECT test('abc');", "abc", Schema.TypeName.STRING, stringUdf); + public void testAsciiStringLiteral() { + runUdfTypeTest("SELECT test_string('abc');", "abc", Schema.TypeName.STRING); } @Test - public void testAsciiStringInput() throws NoSuchMethodException { - runUdfTypeTest( - "SELECT test(string_ascii) FROM table;", "abc", Schema.TypeName.STRING, stringUdf); + public void testAsciiStringInput() { + runUdfTypeTest("SELECT test_string(string_ascii) FROM table;", "abc", Schema.TypeName.STRING); } @Test - public void testUnicodeStringLiteral() throws NoSuchMethodException { - runUdfTypeTest("SELECT test('スタリング');", "スタリング", Schema.TypeName.STRING, stringUdf); + public void testUnicodeStringLiteral() { + runUdfTypeTest("SELECT test_string('スタリング');", "スタリング", Schema.TypeName.STRING); } @Test - public void testUnicodeStringInput() throws NoSuchMethodException { + public void testUnicodeStringInput() { runUdfTypeTest( - "SELECT test(string_unicode) FROM table;", "スタリング", Schema.TypeName.STRING, stringUdf); + "SELECT test_string(string_unicode) FROM table;", "スタリング", Schema.TypeName.STRING); } @Test - public void testEmptyBytesLiteral() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(b'');", new byte[] {}, Schema.TypeName.BYTES, bytesUdf); + public void testEmptyBytesLiteral() { + runUdfTypeTest("SELECT test_bytes(b'');", new byte[] {}, Schema.TypeName.BYTES); } @Test - public void testEmptyBytesInput() throws NoSuchMethodException { + public void testEmptyBytesInput() { runUdfTypeTest( - "SELECT test(bytes_empty) FROM table;", new byte[] {}, Schema.TypeName.BYTES, bytesUdf); + "SELECT test_bytes(bytes_empty) FROM table;", new byte[] {}, Schema.TypeName.BYTES); } @Test - public void testAsciiBytesLiteral() throws NoSuchMethodException { - runUdfTypeTest( - "SELECT test(b'abc');", new byte[] {'a', 'b', 'c'}, Schema.TypeName.BYTES, bytesUdf); + public void testAsciiBytesLiteral() { + runUdfTypeTest("SELECT test_bytes(b'abc');", new byte[] {'a', 'b', 'c'}, Schema.TypeName.BYTES); } @Test - public void testAsciiBytesInput() throws NoSuchMethodException { + public void testAsciiBytesInput() { runUdfTypeTest( - "SELECT test(bytes_ascii) FROM table;", + "SELECT test_bytes(bytes_ascii) FROM table;", new byte[] {'a', 'b', 'c'}, - Schema.TypeName.BYTES, - bytesUdf); + Schema.TypeName.BYTES); } @Test - public void testUnicodeBytesLiteral() throws NoSuchMethodException { - runUdfTypeTest( - "SELECT test(b'ス');", new byte[] {-29, -126, -71}, Schema.TypeName.BYTES, bytesUdf); + public void testUnicodeBytesLiteral() { + runUdfTypeTest("SELECT test_bytes(b'ス');", new byte[] {-29, -126, -71}, Schema.TypeName.BYTES); } @Test - public void testUnicodeBytesInput() throws NoSuchMethodException { + public void testUnicodeBytesInput() { runUdfTypeTest( - "SELECT test(bytes_unicode) FROM table;", + "SELECT test_bytes(bytes_unicode) FROM table;", new byte[] {-29, -126, -71}, - Schema.TypeName.BYTES, - bytesUdf); + Schema.TypeName.BYTES); } @Test - public void testZeroFloat64Literal() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(0.0);", 0.0, Schema.TypeName.DOUBLE, doubleUdf); + public void testZeroFloat64Literal() { + runUdfTypeTest("SELECT test_float64(0.0);", 0.0, Schema.TypeName.DOUBLE); } @Test - public void testZeroFloat64Input() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(float64_0) FROM table;", 0.0, Schema.TypeName.DOUBLE, doubleUdf); + public void testZeroFloat64Input() { + runUdfTypeTest("SELECT test_float64(float64_0) FROM table;", 0.0, Schema.TypeName.DOUBLE); } @Test - public void testNonIntegerFloat64Literal() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(0.123);", 0.123, Schema.TypeName.DOUBLE, doubleUdf); + public void testNonIntegerFloat64Literal() { + runUdfTypeTest("SELECT test_float64(0.123);", 0.123, Schema.TypeName.DOUBLE); } @Test - public void testNonIntegerFloat64Input() throws NoSuchMethodException { + public void testNonIntegerFloat64Input() { runUdfTypeTest( - "SELECT test(float64_noninteger) FROM table;", 0.123, Schema.TypeName.DOUBLE, doubleUdf); + "SELECT test_float64(float64_noninteger) FROM table;", 0.123, Schema.TypeName.DOUBLE); } @Test - public void testPosFloat64Literal() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(123.0);", 123.0, Schema.TypeName.DOUBLE, doubleUdf); + public void testPosFloat64Literal() { + runUdfTypeTest("SELECT test_float64(123.0);", 123.0, Schema.TypeName.DOUBLE); } @Test - public void testPosFloat64Input() throws NoSuchMethodException { - runUdfTypeTest( - "SELECT test(float64_pos) FROM table;", 123.0, Schema.TypeName.DOUBLE, doubleUdf); + public void testPosFloat64Input() { + runUdfTypeTest("SELECT test_float64(float64_pos) FROM table;", 123.0, Schema.TypeName.DOUBLE); } @Test - public void testNegFloat64Literal() throws NoSuchMethodException { - runUdfTypeTest("SELECT test(-123.0);", -123.0, Schema.TypeName.DOUBLE, doubleUdf); + public void testNegFloat64Literal() { + runUdfTypeTest("SELECT test_float64(-123.0);", -123.0, Schema.TypeName.DOUBLE); } @Test - public void testNegFloat64Input() throws NoSuchMethodException { - runUdfTypeTest( - "SELECT test(float64_neg) FROM table;", -123.0, Schema.TypeName.DOUBLE, doubleUdf); + public void testNegFloat64Input() { + runUdfTypeTest("SELECT test_float64(float64_neg) FROM table;", -123.0, Schema.TypeName.DOUBLE); } @Test - public void testMaxFloat64Literal() throws NoSuchMethodException { + public void testMaxFloat64Literal() { runUdfTypeTest( - "SELECT test(1.7976931348623157e+308);", + "SELECT test_float64(1.7976931348623157e+308);", 1.7976931348623157e+308, - Schema.TypeName.DOUBLE, - doubleUdf); + Schema.TypeName.DOUBLE); } @Test - public void testMaxFloat64Input() throws NoSuchMethodException { + public void testMaxFloat64Input() { runUdfTypeTest( - "SELECT test(float64_max) FROM table;", + "SELECT test_float64(float64_max) FROM table;", 1.7976931348623157e+308, - Schema.TypeName.DOUBLE, - doubleUdf); + Schema.TypeName.DOUBLE); } @Test - public void testMinPosFloat64Literal() throws NoSuchMethodException { + public void testMinPosFloat64Literal() { runUdfTypeTest( - "SELECT test(2.2250738585072014e-308);", + "SELECT test_float64(2.2250738585072014e-308);", 2.2250738585072014e-308, - Schema.TypeName.DOUBLE, - doubleUdf); + Schema.TypeName.DOUBLE); } @Test - public void testMinPosFloat64Input() throws NoSuchMethodException { + public void testMinPosFloat64Input() { runUdfTypeTest( - "SELECT test(float64_min_pos) FROM table;", + "SELECT test_float64(float64_min_pos) FROM table;", 2.2250738585072014e-308, - Schema.TypeName.DOUBLE, - doubleUdf); + Schema.TypeName.DOUBLE); } @Test @Ignore( "+inf is implemented as a ZetaSQL builtin function, so combining it with a UDF requires Calc splitting (BEAM-12009).") - public void testPosInfFloat64Literal() throws NoSuchMethodException { + public void testPosInfFloat64Literal() { runUdfTypeTest( - "SELECT test(CAST('+inf' AS FLOAT64));", + "SELECT test_float64(CAST('+inf' AS FLOAT64));", Double.POSITIVE_INFINITY, - Schema.TypeName.DOUBLE, - doubleUdf); + Schema.TypeName.DOUBLE); } @Test - public void testPosInfFloat64Input() throws NoSuchMethodException { + public void testPosInfFloat64Input() { runUdfTypeTest( - "SELECT test(float64_inf) FROM table;", + "SELECT test_float64(float64_inf) FROM table;", Double.POSITIVE_INFINITY, - Schema.TypeName.DOUBLE, - doubleUdf); + Schema.TypeName.DOUBLE); } @Test @Ignore( "-inf is implemented as a ZetaSQL builtin function, so combining it with a UDF requires Calc splitting (BEAM-12009).") - public void testNegInfFloat64Literal() throws NoSuchMethodException { + public void testNegInfFloat64Literal() { runUdfTypeTest( - "SELECT test(CAST('-inf' AS FLOAT64));", + "SELECT test_float64(CAST('-inf' AS FLOAT64));", Double.NEGATIVE_INFINITY, - Schema.TypeName.DOUBLE, - doubleUdf); + Schema.TypeName.DOUBLE); } @Test - public void testNegInfFloat64Input() throws NoSuchMethodException { + public void testNegInfFloat64Input() { runUdfTypeTest( - "SELECT test(float64_neg_inf) FROM table;", + "SELECT test_float64(float64_neg_inf) FROM table;", Double.NEGATIVE_INFINITY, - Schema.TypeName.DOUBLE, - doubleUdf); + Schema.TypeName.DOUBLE); } @Test @Ignore( "NaN is implemented as a ZetaSQL builtin function, so combining it with a UDF requires Calc splitting (BEAM-12009).") - public void testNaNFloat64Literal() throws NoSuchMethodException { + public void testNaNFloat64Literal() { runUdfTypeTest( - "SELECT test(CAST('NaN' AS FLOAT64));", Double.NaN, Schema.TypeName.DOUBLE, doubleUdf); + "SELECT test_float64(CAST('NaN' AS FLOAT64));", Double.NaN, Schema.TypeName.DOUBLE); } @Test - public void testNaNFloat64Input() throws NoSuchMethodException { + public void testNaNFloat64Input() { runUdfTypeTest( - "SELECT test(float64_nan) FROM table;", Double.NaN, Schema.TypeName.DOUBLE, doubleUdf); + "SELECT test_float64(float64_nan) FROM table;", Double.NaN, Schema.TypeName.DOUBLE); } - - // TODO(ibzib) Test that dates and times are rejected. - // TODO(ibzib) Test arrays. - // TODO(ibzib) Test structs. }