From e5c1967554aeb848a8430830a6e58a30b6a15683 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Wed, 22 Jul 2026 14:28:12 +0000 Subject: [PATCH] Add registerSqlOperator() to BeamSqlEnv for custom SQL operators Backports commit f53e1abe from dataflow/beam-spark. Adds a mutable, session-scoped ListSqlOperatorTable (extraOperatorTable) to JdbcConnection, chained at the front of the planner's operator table in CalciteQueryPlanner.defaultConfig. This allows registering custom SqlOperators (e.g. with VARIADIC operand checkers) that cannot be expressed through the schema-function auto-wrapping path. Adds BeamSqlEnvRegisterOperatorTest with tests verifying: - Operators are added to the extra operator table - Table starts empty on fresh env - Registered operators are resolvable by name --- .../sdk/extensions/sql/impl/BeamSqlEnv.java | 14 +++ .../sql/impl/CalciteQueryPlanner.java | 10 +- .../extensions/sql/impl/JdbcConnection.java | 28 +++++ .../impl/BeamSqlEnvRegisterOperatorTest.java | 100 ++++++++++++++++++ 4 files changed, 151 insertions(+), 1 deletion(-) create mode 100644 sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvRegisterOperatorTest.java 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)); + } +}