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

Author: damccorm

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java

@@ -132,6 +132,20 @@ public Map<String, String> 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()));

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java

@@ -158,7 +158,17 @@ 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

@@ -51,11 +51,39 @@ 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
+   * 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.
    *

sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnvRegisterOperatorTest.java

@@ -0,0 +1,102 @@
+/*
+ * 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));
+  }
+}