Merge pull request #12 from jingchen2222/feat/support_window_union

Feat/support window union
4paradigm · May 25, 2021 · 86e7be0 · 86e7be0
2 parents 739652f + a6825e8
commit 86e7be0
Show file tree

Hide file tree

Showing 7 changed files with 355 additions and 4 deletions.
diff --git a/zetasql/parser/ast_node_kind.h b/zetasql/parser/ast_node_kind.h
@@ -289,6 +289,7 @@ enum ASTNodeKind {
   AST_TVF_SCHEMA_COLUMN,
   AST_TYPE_PARAMETER_LIST,
   AST_UNARY_EXPRESSION,
+  AST_UNION_TABLE_REFERENCE_LIST,
   AST_UNNEST_EXPRESSION,
   AST_UNNEST_EXPRESSION_WITH_OPT_ALIAS_AND_OFFSET,
   AST_UNTIL_CLAUSE,

diff --git a/zetasql/parser/bison_parser.y b/zetasql/parser/bison_parser.y
@@ -1291,6 +1291,9 @@ using zetasql::ASTDropStatement;
 %type <node> table_path_expression
 %type <node> table_path_expression_base
 %type <node> table_primary
+%type <node> union_table_reference
+%type <node> union_table_reference_list
+%type <node> opt_union_table_reference_list
 %type <node> table_subquery
 %type <node> templated_parameter_type
 %type <node> terminated_statement
@@ -6668,17 +6671,18 @@ opt_exclude_current_time:
     };
   | /* Nothing */ { $$ = false; }
   ;
+
 window_specification:
     identifier
       {
         $$ = MAKE_NODE(ASTWindowSpecification, @$, {$1});
       }
-    | "(" opt_identifier opt_partition_by_clause opt_order_by_clause
+    | "(" opt_identifier opt_union_table_reference_list opt_partition_by_clause opt_order_by_clause
           opt_window_frame_clause opt_exclude_current_time opt_instance_not_in_window ")"
       {
-        auto *window_spec = MAKE_NODE(ASTWindowSpecification, @$, {$2, $3, $4, $5});
-        window_spec->set_is_exclude_current_time($6);
-        window_spec->set_is_instance_not_in_window($7);
+        auto *window_spec = MAKE_NODE(ASTWindowSpecification, @$, {$2, $3, $4, $5, $6});
+        window_spec->set_is_exclude_current_time($7);
+        window_spec->set_is_instance_not_in_window($8);
         $$ = window_spec;
       }
    ;
@@ -8219,6 +8223,54 @@ execute_using_argument_list:
     }
   ;
 
+union_table_reference:
+  maybe_dashed_path_expression
+  | "(" query ")"
+    {
+      zetasql::ASTQuery* query = $2;
+      query->set_is_nested(true);
+      $$ = MAKE_NODE(ASTTableSubquery, @$, {
+          $2, nullptr, nullptr, nullptr, nullptr});
+    }
+  | "(" query ")" "AS" identifier
+    {
+      auto* alias = MAKE_NODE(ASTAlias, @$, {$5});
+      zetasql::ASTQuery* query = $2;
+      query->set_is_nested(true);
+      $$ = MAKE_NODE(ASTTableSubquery, @$, {
+          $2, alias, nullptr, nullptr, nullptr});
+    }
+  | "(" query ")" identifier
+    {
+      auto* alias = MAKE_NODE(ASTAlias, @$, {$4});
+      zetasql::ASTQuery* query = $2;
+      query->set_is_nested(true);
+      $$ = MAKE_NODE(ASTTableSubquery, @$, {
+          $2, alias, nullptr, nullptr, nullptr});
+    }
+  ;
+
+union_table_reference_list:
+  union_table_reference
+    {
+      $$ = MAKE_NODE(ASTUnionTableReferenceList, @$, {$1});
+    }
+  | union_table_reference_list "," union_table_reference
+    {
+      $$ = WithEndLocation(WithExtraChildren($1, {$3}), @$);
+    }
+  ;
+
+opt_union_table_reference_list:
+  KW_UNION union_table_reference_list
+    {
+      $$ = $2;
+    }
+  | /* Nothing */
+    {
+      $$ = nullptr;
+    }
+  ;
 opt_execute_using_clause:
   KW_USING execute_using_argument_list
     {

diff --git a/zetasql/parser/parse_tree.cc b/zetasql/parser/parse_tree.cc
@@ -323,6 +323,7 @@ static absl::flat_hash_map<ASTNodeKind, std::string> CreateNodeNamesMap() {
   map[AST_TVF_SCHEMA] = "TVFSchema";
   map[AST_TVF] = "TVF";
   map[AST_TYPE_PARAMETER_LIST] = "TypeParameterList";
+  map[AST_UNION_TABLE_REFERENCE_LIST] = "UnionTableReferenceList";
   map[AST_UNARY_EXPRESSION] = "UnaryExpression";
   map[AST_UNNEST_EXPRESSION] = "UnnestExpression";
   map[AST_UNNEST_EXPRESSION_WITH_OPT_ALIAS_AND_OFFSET] =

diff --git a/zetasql/parser/parse_tree_manual.h b/zetasql/parser/parse_tree_manual.h
@@ -2820,6 +2820,7 @@ class ASTWindowSpecification final : public ASTNode {
   zetasql_base::StatusOr<VisitResult> Accept(
       NonRecursiveParseTreeVisitor* visitor) const override;
 
+  const ASTUnionTableReferenceList* union_table_references() const { return union_table_references_; }
   const ASTPartitionBy* partition_by() const { return partition_by_; }
   const ASTOrderBy* order_by() const { return order_by_; }
   const ASTWindowFrame* window_frame() const { return window_frame_; }
@@ -2836,13 +2837,15 @@ class ASTWindowSpecification final : public ASTNode {
   void InitFields() final {
     FieldLoader fl(this);
     fl.AddOptional(&base_window_name_, AST_IDENTIFIER);
+    fl.AddOptional(&union_table_references_, AST_UNION_TABLE_REFERENCE_LIST);
     fl.AddOptional(&partition_by_, AST_PARTITION_BY);
     fl.AddOptional(&order_by_, AST_ORDER_BY);
     fl.AddOptional(&window_frame_, AST_WINDOW_FRAME);
   }
 
   // All are optional, can be NULL.
   const ASTIdentifier* base_window_name_ = nullptr;
+  const ASTUnionTableReferenceList* union_table_references_ = nullptr;
   const ASTPartitionBy* partition_by_ = nullptr;
   const ASTOrderBy* order_by_ = nullptr;
   const ASTWindowFrame* window_frame_ = nullptr;
@@ -3792,6 +3795,28 @@ class ASTTableClause final : public ASTNode {
   const ASTTVF* tvf_ = nullptr;
 };
 
+class ASTUnionTableReferenceList final : public ASTNode {
+ public:
+  static constexpr ASTNodeKind kConcreteNodeKind = AST_UNION_TABLE_REFERENCE_LIST;
+
+  ASTUnionTableReferenceList() : ASTNode(kConcreteNodeKind) {}
+  void Accept(ParseTreeVisitor* visitor, void* data) const override;
+  zetasql_base::StatusOr<VisitResult> Accept(
+      NonRecursiveParseTreeVisitor* visitor) const override;
+
+  const absl::Span<const ASTNode* const>& table_references() const {
+    return table_references_;
+  }
+
+ private:
+  void InitFields() final {
+    FieldLoader fl(this);
+    fl.AddRestAsRepeated(&table_references_);
+  }
+
+  absl::Span<const ASTNode* const> table_references_;
+};
+
 // This represents a clause of form "MODEL <target>", where <target> is a model
 // name.
 class ASTModelClause final : public ASTNode {