From e5c1967554aeb848a8430830a6e58a30b6a15683 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Wed, 22 Jul 2026 14:28:12 +0000 Subject: [PATCH 1/2] 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)); + } +} From 222f3db65ef313bc4a2171db975e46138e3cee3a Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Mon, 27 Jul 2026 20:42:53 +0000 Subject: [PATCH 2/2] format imports --- .../sdk/extensions/sql/impl/BeamSqlEnv.java | 17 +++++------ .../extensions/sql/impl/JdbcConnection.java | 30 ++++++++----------- 2 files changed, 20 insertions(+), 27 deletions(-) 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 dc93818c01c3..8d8d48dd5774 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 @@ -49,6 +49,7 @@ import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptUtil; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.schema.Function; 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.tools.RuleSet; import org.checkerframework.checker.nullness.qual.Nullable; @@ -133,16 +134,14 @@ public Map getPipelineOptions() { } /** - * 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}). + * Registers a custom {@link SqlOperator} into the current session's operator table. This allows + * registering functions with non-fixed operand types (e.g. VARIADIC), which the schema-function + * auto-wrapping mechanism cannot express. + * + *

Only safe to call before the first SQL query is planned in this environment (otherwise the + * parser/validator operator table may have already been built without it). */ - public void registerSqlOperator( - org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlOperator operator) { + public void registerSqlOperator(SqlOperator operator) { connection.getExtraOperatorTable().add(operator); } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcConnection.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcConnection.java index 4a77a5993ea3..cbdb1e999534 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcConnection.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcConnection.java @@ -31,6 +31,8 @@ import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalciteSchema; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.schema.Schema; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.schema.SchemaPlus; +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.util.ListSqlOperatorTable; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.checkerframework.checker.nullness.qual.Nullable; @@ -53,20 +55,15 @@ public class JdbcConnection extends CalciteConnectionWrapper { /** * 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. + * 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 final ListSqlOperatorTable extraOperatorTable = new ListSqlOperatorTable(); private JdbcConnection(CalciteConnection connection) throws SQLException { super(connection); @@ -75,12 +72,9 @@ private JdbcConnection(CalciteConnection connection) throws SQLException { /** * 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. + * table. Add custom {@link 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() { + public ListSqlOperatorTable getExtraOperatorTable() { return extraOperatorTable; }