diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java index d84783118bbd..dc93818c01c3 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java @@ -132,6 +132,20 @@ public Map getPipelineOptions() { return connection.getPipelineOptionsMap(); } + /** + * Registers a pre-built {@link + * org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlOperator} into the + * session-scoped operator table chained at the front of the planner's operator table. Use this + * (instead of the schema-function registration path) when the operator must declare a non-fixed + * operand checker (e.g. {@code VARIADIC}) — the schema auto-wrapping path always produces a + * fixed-parameter checker, which breaks SQL routine overload resolution for {@code ANY}-typed + * parameters (the precedence comparison asserts on {@code ANY}). + */ + public void registerSqlOperator( + org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlOperator operator) { + connection.getExtraOperatorTable().add(operator); + } + public String explain(String sqlString) throws ParseException { try { return RelOptUtil.toString(planner.convertToBeamRel(sqlString, QueryParameters.ofNone())); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java index aa6f4d121871..cedfc5dc6d1d 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java @@ -158,7 +158,15 @@ public FrameworkConfig defaultConfig(JdbcConnection connection, Collection pipelineOptionsMap; private @Nullable PipelineOptions pipelineOptions; + /** + * A mutable, session-scoped operator table that callers can populate with custom {@link + * org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlOperator}s after the + * connection (and its planner) is built. {@link CalciteQueryPlanner#defaultConfig} chains this + * table at the front of the validator/converter operator table, so operators added here resolve + * in subsequent SQL. This is the channel for functions that must declare a non-fixed (e.g. {@code + * VARIADIC}) operand checker, which the schema-function auto-wrapping in {@code + * CalciteCatalogReader.toOp} (always fixed-parameter {@code OperandMetadataImpl}) cannot express. + * The list is held by reference in the planner config, so additions made after the planner is + * built are visible. + */ + private final org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.util + .ListSqlOperatorTable + extraOperatorTable = + new org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.util + .ListSqlOperatorTable(); + private JdbcConnection(CalciteConnection connection) throws SQLException { super(connection); this.pipelineOptionsMap = Collections.emptyMap(); } + /** + * The session-scoped, mutable operator table chained at the front of the planner's operator + * table. Add custom {@link + * org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlOperator}s here to make them + * resolvable in subsequent SQL. + */ + public org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.util.ListSqlOperatorTable + getExtraOperatorTable() { + return extraOperatorTable; + } + /** * Wraps and initializes the initial connection created by Calcite. * diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvRegisterOperatorTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvRegisterOperatorTest.java new file mode 100644 index 000000000000..a7ff3584f2ba --- /dev/null +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvRegisterOperatorTest.java @@ -0,0 +1,100 @@ +/* + * 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.impl; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import org.apache.beam.sdk.extensions.sql.impl.rel.BaseRelTest; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlFunction; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlIdentifier; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlKind; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlOperator; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlSyntax; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.type.OperandTypes; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.type.ReturnTypes; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.validate.SqlNameMatcher; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.validate.SqlNameMatchers; +import org.junit.Test; + +/** + * Tests for {@link BeamSqlEnv#registerSqlOperator(SqlOperator)} and the session-scoped operator + * table in {@link JdbcConnection}. + */ +public class BeamSqlEnvRegisterOperatorTest extends BaseRelTest { + + /** A simple custom function for testing, not present in the default operator table. */ + private static final String CUSTOM_FUNCTION_NAME = "BEAM_TEST_CUSTOM_FUNC"; + + private static SqlFunction createCustomFunction() { + return new SqlFunction( + CUSTOM_FUNCTION_NAME, + SqlKind.OTHER_FUNCTION, + ReturnTypes.INTEGER, + null, + OperandTypes.NUMERIC, + SqlFunctionCategory.USER_DEFINED_FUNCTION); + } + + @Test + public void testExtraOperatorTableStartsEmpty() { + // A fresh env should have an empty extra operator table + BeamSqlEnv freshEnv = BeamSqlEnv.readOnly("empty_test", new java.util.HashMap<>()); + assertTrue( + "Extra operator table should be empty on a fresh env", + freshEnv.connection.getExtraOperatorTable().getOperatorList().isEmpty()); + } + + @Test + public void testRegisterSqlOperator_addsToExtraOperatorTable() { + SqlFunction customFunc = createCustomFunction(); + env.registerSqlOperator(customFunc); + + // Verify the operator is in the extra operator table + assertTrue( + "Registered operator should be present in the extra operator table", + env.connection.getExtraOperatorTable().getOperatorList().contains(customFunc)); + assertEquals(1, env.connection.getExtraOperatorTable().getOperatorList().size()); + } + + @Test + public void testRegisterSqlOperator_makesOperatorResolvableByName() { + SqlFunction customFunc = createCustomFunction(); + env.registerSqlOperator(customFunc); + + // Verify the operator is resolvable by name in the extra operator table + SqlNameMatcher nameMatcher = SqlNameMatchers.withCaseSensitive(false); + java.util.List resolvedOps = new java.util.ArrayList<>(); + env.connection + .getExtraOperatorTable() + .lookupOperatorOverloads( + new SqlIdentifier(CUSTOM_FUNCTION_NAME, SqlParserPos.ZERO), + SqlFunctionCategory.USER_DEFINED_FUNCTION, + SqlSyntax.FUNCTION, + resolvedOps, + nameMatcher); + + assertFalse( + "Custom operator should be resolvable by name after registration", resolvedOps.isEmpty()); + assertTrue( + "The resolved operator should be the one we registered", resolvedOps.contains(customFunc)); + } +}