Chore: clean up ANALYZE implementation (#4607)

* Chore: clean up ANALYZE implementation

* PR feedback

Author: georgesittas

sqlglot/expressions.py

@@ -2329,7 +2329,7 @@ class LoadData(Expression):
 
 
 class Partition(Expression):
-    arg_types = {"expressions": True}
+    arg_types = {"expressions": True, "subpartition": False}
 
 
 class PartitionRange(Expression):
@@ -4728,7 +4728,7 @@ class Analyze(Expression):
     }
 
 
-class Statistics(Expression):
+class AnalyzeStatistics(Expression):
     arg_types = {
         "kind": True,
         "option": False,
@@ -4737,7 +4737,7 @@ class Statistics(Expression):
     }
 
 
-class Histogram(Expression):
+class AnalyzeHistogram(Expression):
     arg_types = {
         "this": True,
         "expressions": True,
@@ -4746,29 +4746,20 @@ class Histogram(Expression):
     }
 
 
-class Sample(Expression):
-    arg_types = {
-        "kind": True,
-        "sample": True,
-    }
+class AnalyzeSample(Expression):
+    arg_types = {"kind": True, "sample": True}
 
 
 class AnalyzeListChainedRows(Expression):
-    arg_types = {
-        "expression": False,
-    }
+    arg_types = {"expression": False}
 
 
 class AnalyzeDelete(Expression):
-    arg_types = {
-        "kind": False,
-    }
+    arg_types = {"kind": False}
 
 
 class AnalyzeWith(Expression):
-    arg_types = {
-        "expressions": True,
-    }
+    arg_types = {"expressions": True}
 
 
 class AnalyzeValidate(Expression):

sqlglot/generator.py

@@ -114,6 +114,8 @@ class Generator(metaclass=_Generator):
         **JSON_PATH_PART_TRANSFORMS,
         exp.AllowedValuesProperty: lambda self,
         e: f"ALLOWED_VALUES {self.expressions(e, flat=True)}",
+        exp.AnalyzeColumns: lambda self, e: self.sql(e, "this"),
+        exp.AnalyzeWith: lambda self, e: self.expressions(e, prefix="WITH ", sep=" "),
         exp.ArrayContainsAll: lambda self, e: self.binary(e, "@>"),
         exp.ArrayOverlaps: lambda self, e: self.binary(e, "&&"),
         exp.AutoRefreshProperty: lambda self, e: f"AUTO REFRESH {self.sql(e, 'this')}",
@@ -198,6 +200,7 @@ class Generator(metaclass=_Generator):
         exp.TransientProperty: lambda *_: "TRANSIENT",
         exp.Union: lambda self, e: self.set_operations(e),
         exp.UnloggedProperty: lambda *_: "UNLOGGED",
+        exp.UsingData: lambda self, e: f"USING DATA {self.sql(e, 'this')}",
         exp.Uuid: lambda *_: "UUID()",
         exp.UppercaseColumnConstraint: lambda *_: "UPPERCASE",
         exp.VarMap: lambda self, e: self.func("MAP", e.args["keys"], e.args["values"]),
@@ -1558,7 +1561,8 @@ def national_sql(self, expression: exp.National, prefix: str = "N") -> str:
         return f"{prefix}{string}"
 
     def partition_sql(self, expression: exp.Partition) -> str:
-        return f"PARTITION({self.expressions(expression, flat=True)})"
+        partition_keyword = "SUBPARTITION" if expression.args.get("subpartition") else "PARTITION"
+        return f"{partition_keyword}({self.expressions(expression, flat=True)})"
 
     def properties_sql(self, expression: exp.Properties) -> str:
         root_properties = []
@@ -4656,12 +4660,12 @@ def unpivotcolumns_sql(self, expression: exp.UnpivotColumns) -> str:
 
         return f"NAME {name} VALUE {values}"
 
-    def sample_sql(self, expression: exp.Sample) -> str:
+    def analyzesample_sql(self, expression: exp.AnalyzeSample) -> str:
         kind = self.sql(expression, "kind")
         sample = self.sql(expression, "sample")
         return f"SAMPLE {sample} {kind}"
 
-    def statistics_sql(self, expression: exp.Statistics) -> str:
+    def analyzestatistics_sql(self, expression: exp.AnalyzeStatistics) -> str:
         kind = self.sql(expression, "kind")
         option = self.sql(expression, "option")
         option = f" {option}" if option else ""
@@ -4671,7 +4675,7 @@ def statistics_sql(self, expression: exp.Statistics) -> str:
         columns = f" {columns}" if columns else ""
         return f"{kind}{option} STATISTICS{this}{columns}"
 
-    def histogram_sql(self, expression: exp.Histogram) -> str:
+    def analyzehistogram_sql(self, expression: exp.AnalyzeHistogram) -> str:
         this = self.sql(expression, "this")
         columns = self.expressions(expression)
         inner_expression = self.sql(expression, "expression")
@@ -4680,21 +4684,11 @@ def histogram_sql(self, expression: exp.Histogram) -> str:
         update_options = f" {update_options} UPDATE" if update_options else ""
         return f"{this} HISTOGRAM ON {columns}{inner_expression}{update_options}"
 
-    def usingdata_sql(self, expression: exp.UsingData) -> str:
-        data = self.sql(expression, "this")
-        return f"USING DATA {data}"
-
-    def analyzecolumns_sql(self, expression: exp.AnalyzeColumns) -> str:
-        return self.sql(expression, "this")
-
     def analyzedelete_sql(self, expression: exp.AnalyzeDelete) -> str:
         kind = self.sql(expression, "kind")
         kind = f" {kind}" if kind else ""
         return f"DELETE{kind} STATISTICS"
 
-    def analyzewith_sql(self, expression: exp.AnalyzeWith) -> str:
-        return self.expressions(expression, prefix="WITH ", sep=" ")
-
     def analyzelistchainedrows_sql(self, expression: exp.AnalyzeListChainedRows) -> str:
         inner_expression = self.sql(expression, "expression")
         return f"LIST CHAINED ROWS{inner_expression}"

sqlglot/parser.py

@@ -1352,6 +1352,8 @@ class Parser(metaclass=_Parser):
         "VALIDATE": lambda self: self._parse_analyze_validate(),
     }
 
+    PARTITION_KEYWORDS = {"PARTITION", "SUBPARTITION"}
+
     AMBIGUOUS_ALIAS_TOKENS = (TokenType.LIMIT, TokenType.OFFSET)
 
     OPERATION_MODIFIERS: t.Set[str] = set()
@@ -2984,11 +2986,13 @@ def _parse_cache(self) -> exp.Cache:
         )
 
     def _parse_partition(self) -> t.Optional[exp.Partition]:
-        if not self._match(TokenType.PARTITION):
+        if not self._match_texts(self.PARTITION_KEYWORDS):
             return None
 
         return self.expression(
-            exp.Partition, expressions=self._parse_wrapped_csv(self._parse_assignment)
+            exp.Partition,
+            subpartition=self._prev.text.upper() == "SUBPARTITION",
+            expressions=self._parse_wrapped_csv(self._parse_assignment),
         )
 
     def _parse_value(self) -> t.Optional[exp.Tuple]:
@@ -7145,44 +7149,41 @@ def _parse_analyze(self) -> exp.Analyze | exp.Command:
             else:
                 options.append(self._prev.text.upper())
 
-        kind = None
-        mode = None
         this: t.Optional[exp.Expression] = None
-        partition: t.Optional[exp.Expression] = None
         inner_expression: t.Optional[exp.Expression] = None
 
+        kind = self._curr and self._curr.text.upper()
+
         if self._match(TokenType.TABLE) or self._match(TokenType.INDEX):
-            kind = self._prev.text.upper()
             this = self._parse_table_parts()
         elif self._match_text_seq("TABLES"):
-            kind = self._prev.text.upper()
             if self._match_set((TokenType.FROM, TokenType.IN)):
                 kind = f"{kind} {self._prev.text.upper()}"
                 this = self._parse_table(schema=True, is_db_reference=True)
         elif self._match_text_seq("DATABASE"):
-            kind = self._prev.text.upper()
             this = self._parse_table(schema=True, is_db_reference=True)
         elif self._match_text_seq("CLUSTER"):
-            kind = self._prev.text.upper()
             this = self._parse_table()
         # Try matching inner expr keywords before fallback to parse table.
         elif self._match_texts(self.ANALYZE_EXPRESSION_PARSERS):
+            kind = None
             inner_expression = self.ANALYZE_EXPRESSION_PARSERS[self._prev.text.upper()](self)
         else:
             # Empty kind  https://prestodb.io/docs/current/sql/analyze.html
+            kind = None
             this = self._parse_table_parts()
 
-        try:
-            partition = self._parse_partition()
-        except ParseError:
+        partition = self._try_parse(self._parse_partition)
+        if not partition and self._match_texts(self.PARTITION_KEYWORDS):
             return self._parse_as_command(start)
 
         # https://docs.starrocks.io/docs/sql-reference/sql-statements/cbo_stats/ANALYZE_TABLE/
-        if self._match_text_seq("WITH", "SYNC", "MODE", advance=False) or self._match_text_seq(
-            "WITH", "ASYNC", "MODE", advance=False
+        if self._match_text_seq("WITH", "SYNC", "MODE") or self._match_text_seq(
+            "WITH", "ASYNC", "MODE"
         ):
-            mode = f"WITH {self._next.text.upper()} MODE"
-            self._advance(3)
+            mode = f"WITH {self._tokens[self._index-2].text.upper()} MODE"
+        else:
+            mode = None
 
         if self._match_texts(self.ANALYZE_EXPRESSION_PARSERS):
             inner_expression = self.ANALYZE_EXPRESSION_PARSERS[self._prev.text.upper()](self)
@@ -7200,7 +7201,7 @@ def _parse_analyze(self) -> exp.Analyze | exp.Command:
         )
 
     # https://spark.apache.org/docs/3.5.1/sql-ref-syntax-aux-analyze-table.html
-    def _parse_analyze_statistics(self) -> exp.Statistics:
+    def _parse_analyze_statistics(self) -> exp.AnalyzeStatistics:
         this = None
         kind = self._prev.text.upper()
         option = self._prev.text.upper() if self._match_text_seq("DELTA") else None
@@ -7221,14 +7222,14 @@ def _parse_analyze_statistics(self) -> exp.Statistics:
             sample = self._parse_number()
             expressions = [
                 self.expression(
-                    exp.Sample,
+                    exp.AnalyzeSample,
                     sample=sample,
                     kind=self._prev.text.upper() if self._match(TokenType.PERCENT) else None,
                 )
             ]
 
         return self.expression(
-            exp.Statistics, kind=kind, option=option, this=this, expressions=expressions
+            exp.AnalyzeStatistics, kind=kind, option=option, this=this, expressions=expressions
         )
 
     # https://docs.oracle.com/en/database/oracle/oracle-database/21/sqlrf/ANALYZE.html
@@ -7257,24 +7258,21 @@ def _parse_analyze_columns(self) -> t.Optional[exp.AnalyzeColumns]:
         this = self._prev.text.upper()
         if self._match_text_seq("COLUMNS"):
             return self.expression(exp.AnalyzeColumns, this=f"{this} {self._prev.text.upper()}")
-        self.raise_error("Expecting COLUMNS")
         return None
 
     def _parse_analyze_delete(self) -> t.Optional[exp.AnalyzeDelete]:
         kind = self._prev.text.upper() if self._match_text_seq("SYSTEM") else None
         if self._match_text_seq("STATISTICS"):
             return self.expression(exp.AnalyzeDelete, kind=kind)
-        self.raise_error("Expecting STATISTICS")
         return None
 
     def _parse_analyze_list(self) -> t.Optional[exp.AnalyzeListChainedRows]:
         if self._match_text_seq("CHAINED", "ROWS"):
             return self.expression(exp.AnalyzeListChainedRows, expression=self._parse_into())
-        self.raise_error("Expecting CHAINED ROWS")
         return None
 
     # https://dev.mysql.com/doc/refman/8.4/en/analyze-table.html
-    def _parse_analyze_histogram(self) -> exp.Histogram:
+    def _parse_analyze_histogram(self) -> exp.AnalyzeHistogram:
         this = self._prev.text.upper()
         expression: t.Optional[exp.Expression] = None
         expressions = []
@@ -7305,7 +7303,7 @@ def _parse_analyze_histogram(self) -> exp.Histogram:
                 expression = self.expression(exp.UsingData, this=self._parse_string())
 
         return self.expression(
-            exp.Histogram,
+            exp.AnalyzeHistogram,
             this=this,
             expressions=expressions,
             expression=expression,

tests/dialects/test_oracle.py

@@ -659,6 +659,7 @@ def test_analyze(self):
         self.validate_identity("ANALYZE TABLE tbl")
         self.validate_identity("ANALYZE INDEX ndx")
         self.validate_identity("ANALYZE TABLE db.tbl PARTITION(foo = 'foo', bar = 'bar')")
+        self.validate_identity("ANALYZE TABLE db.tbl SUBPARTITION(foo = 'foo', bar = 'bar')")
         self.validate_identity("ANALYZE INDEX db.ndx PARTITION(foo = 'foo', bar = 'bar')")
         self.validate_identity("ANALYZE INDEX db.ndx PARTITION(part1)")
         self.validate_identity("ANALYZE CLUSTER db.cluster")