This is an automated email from the ASF dual-hosted git repository.
damccorm 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 24346193cc8 Add registerSqlOperator() to BeamSqlEnv for custom SQL
operators (#39432)
24346193cc8 is described below
commit 24346193cc83be3ea6070de5d8ca2dd68eefd8d0
Author: Danny McCormick <[email protected]>
AuthorDate: Tue Jul 28 09:18:28 2026 -0400
Add registerSqlOperator() to BeamSqlEnv for custom SQL operators (#39432)
* 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
* format imports
---
.../beam/sdk/extensions/sql/impl/BeamSqlEnv.java | 13 +++
.../extensions/sql/impl/CalciteQueryPlanner.java | 10 ++-
.../sdk/extensions/sql/impl/JdbcConnection.java | 22 +++++
.../sql/impl/BeamSqlEnvRegisterOperatorTest.java | 100 +++++++++++++++++++++
4 files changed, 144 insertions(+), 1 deletion(-)
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 d84783118bb..8d8d48dd577 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.jdbc.CalcitePre
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;
@@ -132,6 +133,18 @@ public class BeamSqlEnv {
return connection.getPipelineOptionsMap();
}
+ /**
+ * 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.
+ *
+ * <p>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(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 aa6f4d12187..cedfc5dc6d1 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 class CalciteQueryPlanner implements QueryPlanner {
.ruleSets(ruleSets.toArray(new RuleSet[0]))
.costFactory(BeamCostModel.FACTORY)
.typeSystem(connection.getTypeFactory().getTypeSystem())
- .operatorTable(SqlOperatorTables.chain(opTab0, catalogReader))
+ .operatorTable(
+ SqlOperatorTables.chain(
+ // Session-scoped custom operators (e.g. registered scalar
Python UDFs that must
+ // declare a non-fixed VARIADIC operand checker). Chained
first and held by
+ // reference, so operators added to the connection's table
after the planner is
+ // built
+ // still resolve. Placing it ahead of the catalogReader avoids
the duplicate
+ // fixed-parameter overload that schema auto-wrapping would
otherwise create.
+ connection.getExtraOperatorTable(), opTab0, catalogReader))
.sqlToRelConverterConfig(sqlToRelConfig)
.build();
}
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 9c54f059f21..cbdb1e99953 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.CalciteCon
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;
@@ -51,11 +53,31 @@ public class JdbcConnection extends
CalciteConnectionWrapper {
private Map<String, String> pipelineOptionsMap;
private @Nullable PipelineOptions pipelineOptions;
+ /**
+ * A mutable, session-scoped operator table that callers can populate with
custom {@link
+ * 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 ListSqlOperatorTable extraOperatorTable = new
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 SqlOperator}s here to make them resolvable in
subsequent SQL.
+ */
+ public 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 00000000000..a7ff3584f2b
--- /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<SqlOperator> 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));
+ }
+}
