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

Author: damccorm

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;
 
@@ -132,6 +133,18 @@ public Map<String, String> getPipelineOptions() {
     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()));

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<RuleS
         .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();
   }

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;
 
@@ -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.
    *

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));
+  }
+}