apilloud commented on a change in pull request #14418: URL: https://github.com/apache/beam/pull/14418#discussion_r607170581
########## File path: sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlJavaUdfTypeTest.java ########## @@ -0,0 +1,276 @@ +/* + * 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.beam.sdk.extensions.sql.zetasql; + +import static org.junit.Assert.fail; + +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; +import org.apache.beam.sdk.extensions.sql.impl.ScalarFunctionImpl; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils; +import org.apache.beam.sdk.extensions.sql.meta.provider.ReadOnlyTableProvider; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.schemas.Schema; +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.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; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; +import org.joda.time.Duration; +import org.junit.Before; +import org.junit.Ignore; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests verifying that various data types can be passed through Java UDFs without data loss. */ +@RunWith(JUnit4.class) +public class ZetaSqlJavaUdfTypeTest extends ZetaSqlTestBase { + @Rule public transient TestPipeline pipeline = TestPipeline.create(); + @Rule public ExpectedException thrown = ExpectedException.none(); + + @Before + public void setUp() { + initialize(); + + // Add BeamJavaUdfCalcRule to planner to enable UDFs. + this.config = + Frameworks.newConfigBuilder(config) + .ruleSets( + ZetaSQLQueryPlanner.getZetaSqlRuleSets( + ImmutableList.of(BeamJavaUdfCalcRule.INSTANCE)) + .toArray(new RuleSet[0])) + .build(); + } + + public static class BooleanIdentityFn implements BeamSqlUdf { + public Boolean eval(Boolean input) { + return input; + } + } + + public static class Int64IdentityFn implements BeamSqlUdf { + public Long eval(Long input) { + return input; + } + } + + public static class StringIdentityFn implements BeamSqlUdf { + public String eval(String input) { + return input; + } + } + + public static class BytesIdentityFn implements BeamSqlUdf { + public byte[] eval(byte[] input) { + return input; + } + } + + public static class DoubleIdentityFn implements BeamSqlUdf { + public Double eval(Double input) { + return input; + } + } + + private void runUdfTypeTest(String query, Object result, Schema.TypeName typeName) + throws NoSuchMethodException { + Class<? extends BeamSqlUdf> udf; + Class<?> javaType; + switch (typeName) { Review comment: nit: The more logic you add to a test the more likely you are to introduce bugs. You might just make these (udf and javaType) arguments to `runUdfTypeTest`. (Those are only used to lookup `Method`, so it might be best to just make Method the argument this class takes.) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
