Support TVF

Author: Matts966

.github/workflows/test.yml

@@ -79,9 +79,9 @@ jobs:
         with:
           commit-message: Update Snapshot
           reviewers: Matts966
-          branch: snapshot-testing-${{ steps.extract_branch.outputs.branch }}
+          branch: snapshot-${{ steps.extract_branch.outputs.branch }}
           delete-branch: true
           title: 'Updated Snapshot :tada:'
           body: |
             Updated Snapshot :tada:
-            Auto-generated by [create-pull-request][https://github.com/peter-evans/create-pull-request]
+            Auto-generated by [create-pull-request](https://github.com/peter-evans/create-pull-request)

alphasql/BUILD

@@ -99,6 +99,7 @@ cc_binary(
         "@com_google_zetasql//zetasql/public:evaluator_table_iterator",
         "@com_google_zetasql//zetasql/public:language_options",
         "@com_google_zetasql//zetasql/public:simple_catalog",
+        "@com_google_zetasql//zetasql/public:templated_sql_tvf",
         "@com_google_zetasql//zetasql/public:type",
         "@com_google_zetasql//zetasql/public:type_cc_proto",
         "@com_google_zetasql//zetasql/public:value",

alphasql/alphacheck.cc

@@ -43,6 +43,7 @@
 #include "zetasql/public/type.h"
 #include "zetasql/public/value.h"
 #include "zetasql/public/templated_sql_function.h"
+#include "zetasql/public/templated_sql_tvf.h"
 #include "zetasql/resolved_ast/resolved_ast.h"
 
 #include "alphasql/common_lib.h"
@@ -216,6 +217,19 @@ absl::Status check(const std::string &sql, const ASTStatement *statement,
     }
     break;
   }
+  case RESOLVED_CREATE_TABLE_FUNCTION_STMT: {
+    auto *create_table_function_stmt =
+        resolved_statement->GetAs<ResolvedCreateTableFunctionStmt>();
+    std::cout
+        << "Create Table Function Statement analyzed, adding function to catalog..."
+        << std::endl;
+    catalog->AddOwnedTableValuedFunction(new TemplatedSQLTVF(
+      create_table_function_stmt->name_path(),
+      create_table_function_stmt->signature(),
+      create_table_function_stmt->argument_name_list(),
+      ParseResumeLocation::FromString(create_table_function_stmt->code())));
+    break;
+  }
   // TODO: DROP PROCEDURE Support?
   case RESOLVED_CREATE_PROCEDURE_STMT: {
     auto *create_procedure_stmt =
@@ -484,13 +498,25 @@ int main(int argc, char *argv[]) {
       status = zetasql::UpdateErrorLocationPayloadWithFilenameIfNotPresent(
           status, sql_file_path);
       std::cerr << "ERROR: " << status << std::endl;
-      std::cout << "catalog:" << std::endl;
+      std::cout << "tables:" << std::endl;
       // For deterministic output
       auto table_names = catalog->table_names();
       std::sort(table_names.begin(), table_names.end());
       for (const std::string &table_name : table_names) {
         std::cout << "\t" << table_name << std::endl;
       }
+      // Too many outputs
+      /* auto function_names = catalog->function_names(); */
+      /* std::sort(function_names.begin(), function_names.end()); */
+      /* for (const std::string &function_name : function_names) { */
+      /*   std::cout << "\t" << function_name << std::endl; */
+      /* } */
+      std::cout << "tvfs:" << std::endl;
+      auto table_function_names = catalog->table_valued_function_names();
+      std::sort(table_function_names.begin(), table_function_names.end());
+      for (const std::string &table_function_name : table_function_names) {
+        std::cout << "\t" << table_function_name << std::endl;
+      }
       return 1;
     }
   }

alphasql/identifier_resolver.cc

@@ -138,12 +138,6 @@ void IdentifierResolver::visitASTCreateTableStatement(
   visitASTChildren(node, data);
 }
 
-// TODO:(Matts966) Check if this node is callee or caller and implement
-// correctly void IdentifierResolver::visitASTTVF(const ASTTVF* node, void*
-// data) {
-//   function_information.called.insert(node->name()->ToIdentifierVector());
-// }
-
 // Check INSERT and UPDATE statement to emit warnings for side effects.
 void IdentifierResolver::visitASTInsertStatement(const ASTInsertStatement *node,
                                                  void *data) {
@@ -231,6 +225,12 @@ void IdentifierResolver::visitASTDropFunctionStatement(
   visitASTChildren(node, data);
 }
 
+void IdentifierResolver::visitASTTVF(const ASTTVF* node, void* data) {
+  identifier_information.function_information.called.insert(
+      node->name()->ToIdentifierVector());
+  visitASTChildren(node, data);
+}
+
 void IdentifierResolver::visitASTFunctionCall(const ASTFunctionCall *node,
                                               void *data) {
   identifier_information.function_information.called.insert(

alphasql/identifier_resolver.h

@@ -101,6 +101,7 @@ class IdentifierResolver : public DefaultParseTreeVisitor {
   // Functions
   void visitASTDropFunctionStatement(const ASTDropFunctionStatement *node,
                                      void *data) override;
+  void visitASTTVF(const ASTTVF* node, void* data) override;
   void visitASTFunctionCall(const ASTFunctionCall *node, void *data) override;
   void visitASTFunctionDeclaration(const ASTFunctionDeclaration *node,
                                    void *data) override;

samples/tvf/call_tvf.sql

@@ -0,0 +1,5 @@
+CREATE OR REPLACE TABLE `dataset.tvf_table` AS
+SELECT
+  *
+FROM
+  `dataset.tvf`(0);

samples/tvf/define_tvf.sql

@@ -0,0 +1,13 @@
+CREATE OR REPLACE TABLE FUNCTION `dataset.tvf`(input int64)
+AS
+  SELECT
+    *
+      EXCEPT (
+        test_code
+      )
+  FROM
+    UNNEST(ARRAY[STRUCT("1" AS id, "apple" AS product_name, 120 AS price, date("2022-01-11") AS purchase_date,
+      1 AS test_code), ("2", "banana", 100, date("2022-01-11"), 1), ("3", "orange", 80, date("2022-01-12"),
+      1)])
+  WHERE
+    test_code = input;