Add registerSqlOperator() to BeamSqlEnv for custom SQL operators

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,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

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