[GitHub] [beam] ibzib commented on a change in pull request #13934: [BEAM-11778] Create a wrapper for ZetaSQL catalog and refactor accordingly.

[GitBox <gi...@apache.org>]

ibzib commented on a change in pull request #13934:
URL: https://github.com/apache/beam/pull/13934#discussion_r583280073



##########
File path: sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCatalog.java
##########
@@ -0,0 +1,417 @@
+/*
+ * 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.zetasql;
+
+import com.google.common.collect.ImmutableList;
+import com.google.zetasql.Analyzer;
+import com.google.zetasql.AnalyzerOptions;
+import com.google.zetasql.Function;
+import com.google.zetasql.FunctionArgumentType;
+import com.google.zetasql.FunctionSignature;
+import com.google.zetasql.SimpleCatalog;
+import com.google.zetasql.TVFRelation;
+import com.google.zetasql.TableValuedFunction;
+import com.google.zetasql.TypeFactory;
+import com.google.zetasql.ZetaSQLBuiltinFunctionOptions;
+import com.google.zetasql.ZetaSQLFunctions;
+import com.google.zetasql.ZetaSQLType;
+import com.google.zetasql.resolvedast.ResolvedNode;
+import com.google.zetasql.resolvedast.ResolvedNodes;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.extensions.sql.impl.JavaUdfLoader;
+import org.apache.beam.sdk.extensions.sql.impl.SqlConversionException;
+import org.apache.beam.sdk.extensions.sql.impl.utils.TVFStreamingUtils;
+import org.apache.beam.sdk.extensions.sql.udf.ScalarFn;
+import org.apache.beam.sdk.extensions.sql.zetasql.translation.UserFunctionDefinitions;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.adapter.java.JavaTypeFactory;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataType;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.SchemaPlus;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+
+/**
+ * Catalog for registering tables and functions. Populates a {@link SimpleCatalog} based on a {@link
+ * SchemaPlus}.
+ */
+public class BeamZetaSqlCatalog {
+  // ZetaSQL function group identifiers. Different function groups may have divergent translation
+  // paths.
+  public static final String PRE_DEFINED_WINDOW_FUNCTIONS = "pre_defined_window_functions";
+  public static final String USER_DEFINED_SQL_FUNCTIONS = "user_defined_functions";
+  public static final String USER_DEFINED_JAVA_SCALAR_FUNCTIONS =
+      "user_defined_java_scalar_functions";
+  /**
+   * Same as {@link Function}.ZETASQL_FUNCTION_GROUP_NAME. Identifies built-in ZetaSQL functions.
+   */
+  public static final String ZETASQL_FUNCTION_GROUP_NAME = "ZetaSQL";
+
+  private static final ImmutableList<String> PRE_DEFINED_WINDOW_FUNCTION_DECLARATIONS =
+      ImmutableList.of(
+          // TODO: support optional function argument (for window_offset).
+          "CREATE FUNCTION TUMBLE(ts TIMESTAMP, window_size STRING) AS (1);",
+          "CREATE FUNCTION TUMBLE_START(window_size STRING) RETURNS TIMESTAMP AS (null);",
+          "CREATE FUNCTION TUMBLE_END(window_size STRING) RETURNS TIMESTAMP AS (null);",
+          "CREATE FUNCTION HOP(ts TIMESTAMP, emit_frequency STRING, window_size STRING) AS (1);",
+          "CREATE FUNCTION HOP_START(emit_frequency STRING, window_size STRING) "
+              + "RETURNS TIMESTAMP AS (null);",
+          "CREATE FUNCTION HOP_END(emit_frequency STRING, window_size STRING) "
+              + "RETURNS TIMESTAMP AS (null);",
+          "CREATE FUNCTION SESSION(ts TIMESTAMP, session_gap STRING) AS (1);",
+          "CREATE FUNCTION SESSION_START(session_gap STRING) RETURNS TIMESTAMP AS (null);",
+          "CREATE FUNCTION SESSION_END(session_gap STRING) RETURNS TIMESTAMP AS (null);");
+
+  /** The top-level Calcite schema, which may contain sub-schemas. */
+  private final SchemaPlus calciteSchema;
+  /**
+   * The top-level ZetaSQL catalog, which may contain nested catalogs for qualified table and
+   * function references.
+   */
+  private final SimpleCatalog zetaSqlCatalog;
+
+  private final JavaTypeFactory typeFactory;
+
+  private final JavaUdfLoader javaUdfLoader = new JavaUdfLoader();
+  private final Map<List<String>, ResolvedNodes.ResolvedCreateFunctionStmt> sqlScalarUdfs =
+      new HashMap<>();
+  /** User-defined table valued functions. */
+  private final Map<List<String>, ResolvedNode> sqlUdtvfs = new HashMap<>();
+
+  private final Map<List<String>, UserFunctionDefinitions.JavaScalarFunction> javaScalarUdfs =
+      new HashMap<>();
+
+  private BeamZetaSqlCatalog(
+      SchemaPlus calciteSchema, SimpleCatalog zetaSqlCatalog, JavaTypeFactory typeFactory) {
+    this.calciteSchema = calciteSchema;
+    this.zetaSqlCatalog = zetaSqlCatalog;
+    this.typeFactory = typeFactory;
+  }
+
+  /** Return catalog pre-populated with builtin functions. */
+  static BeamZetaSqlCatalog create(
+      SchemaPlus topLevelSchema,

Review comment:
       done

##########
File path: sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLPlannerImpl.java
##########
@@ -93,79 +78,28 @@
 
   public RelRoot rel(String sql, QueryParameters params) {
     RelOptCluster cluster = RelOptCluster.create(planner, new RexBuilder(typeFactory));
-    QueryTrait trait = new QueryTrait();
-    SqlAnalyzer analyzer =
-        new SqlAnalyzer(trait, defaultSchemaPlus, (JavaTypeFactory) cluster.getTypeFactory());
-
     AnalyzerOptions options = SqlAnalyzer.getAnalyzerOptions(params, defaultTimezone);
+    BeamZetaSqlCatalog catalog =
+        BeamZetaSqlCatalog.create(
+            defaultSchemaPlus,
+            new SimpleCatalog(defaultSchemaPlus.getName()),

Review comment:
       done (moved to `create` out of a possibly misguided aversion to doing anything in a constructor)




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org