ENH: Lock in aggregation functions

Author: zbrookle

dataframe_sql/grammar/sql.grammar

@@ -42,7 +42,7 @@ JOIN_TYPE.5: "INNER"i | "CROSS"i | /FULL\sOUTER/i | /LEFT\sOUTER/i |  /RIGHT\sOU
                | "CASE"i (when_then)+ "ELSE"i expression_math "END"i -> case_expression
                | "CAST"i "(" expression_math "AS"i TYPENAME ")" -> as_type
                | "CAST"i "(" literal "AS"i TYPENAME ")" -> literal_cast
-               | aggregation "(" expression_math ")" -> sql_aggregation
+               | AGGREGATION "(" expression_math ")" -> sql_aggregation
                | "RANK"i "(" ")" rank_form -> rank_expression
                | "DENSE_RANK"i "(" ")" rank_form -> dense_rank_expression
 
@@ -54,9 +54,12 @@ order: expression_math ["ASC"i] -> order_asc
           | expression_math "DESC"i -> order_desc
 
 column_name: [NAME "."] NAME
-?expression_product: expression
-                  | expression_product "*" expression -> expression_mul
-                  | expression_product "/" expression -> expression_div
+?expression_product: expression_parens
+                  | expression_product "*" expression_parens -> expression_mul
+                  | expression_product "/" expression_parens -> expression_div
+
+?expression_parens: expression
+                  | "(" expression_parens "*" expression ")"
 
 ?expression: [NAME "."] (NAME | STAR) -> column_name
             | literal
@@ -79,7 +82,7 @@ TYPENAME:  "object"i
          | "datetime64"i
          | "timestamp"i
          | "category"i
-?aggregation: NAME -> aggregation_name
+AGGREGATION.8: "sum"i | "avg"i | "min"i | "max"i
 alias: NAME -> alias_string
 _window_name: NAME
 limit_count: integer -> limit_count

dataframe_sql/parsing/sql_parser.py

@@ -554,14 +554,6 @@ def where_expr(self, truth_value_dataframe):
         """
         return Token("where_expr", truth_value_dataframe[0])
 
-    def aggregation_name(self, aggregation_name):
-        """
-        Returns the function name tree
-        :param aggregation_name:
-        :return:
-        """
-        return aggregation_name[0].value
-
     def alias_string(self, name: List[str]):
         """
         Returns an alias token_or_tree with the name extracted
@@ -918,21 +910,14 @@ def aggregate(self, function_name_list_form):
         """
         return "".join(function_name_list_form)
 
-    def aggregation_name(self, tokens):
-        """
-        Extracts function name from token_or_tree
-        :param tokens:
-        :return:
-        """
-        return tokens[0].value
-
-    def sql_aggregation(self, aggregation_expr):
+    def sql_aggregation(self, aggregation_expr: list):
         """
         Handles presence of aggregation in an sql_object
         :param aggregation_expr: Function sql_object
         :return:
         """
-        aggregate_name = aggregation_expr[0]
+        aggregate_token: Token = aggregation_expr[0]
+        aggregate_name: str = aggregate_token.value
         column = aggregation_expr[1]
         table = self.dataframe_map[column.table]
         column_true_name = self.column_name_map[column.table][column.name]

dataframe_sql/tests/pandas_sql_functionality_test.py

@@ -1410,29 +1410,29 @@ def test_order_of_operations_no_parens():
     tm.assert_frame_equal(pandas_frame, my_frame)
 
 
-# def test_order_of_operations_with_parens():
-#     """
-#     Test math parentheses
-#     :return:
-#     """
-#
-#     my_frame = query(
-#         "select 20 * (avocado_id + 3) / (20 + avocado_id) as my_math from " "avocado"
-#     )
-#
-#     pandas_frame = AVOCADO.copy()[["avocado_id"]]
-#     pandas_frame["my_math"] = (
-#         20 * (pandas_frame["avocado_id"] + 3) / (20 + pandas_frame["avocado_id"])
-#     )
-#
-#     pandas_frame = pandas_frame.drop(columns=["avocado_id"])
-#
-#     tm.assert_frame_equal(pandas_frame, my_frame)
+def test_order_of_operations_with_parens():
+    """
+    Test math parentheses
+    :return:
+    """
+
+    my_frame = query(
+        "select 20 * (avocado_id + 3) / (20 + avocado_id) as my_math from " "avocado"
+    )
+
+    pandas_frame = AVOCADO.copy()[["avocado_id"]]
+    pandas_frame["my_math"] = (
+        20 * (pandas_frame["avocado_id"] + 3) / (20 + pandas_frame["avocado_id"])
+    )
+
+    pandas_frame = pandas_frame.drop(columns=["avocado_id"])
+
+    tm.assert_frame_equal(pandas_frame, my_frame)
 
 
 if __name__ == "__main__":
     register_env_tables()
 
-    # test_order_of_operations_with_parens()
+    test_sum()
 
     remove_env_tables()

dataframe_sql/tests/sql_execution_plan_test.py

@@ -462,6 +462,9 @@ def test_having_one_condition():
         "select min(temp) from forest_fires having min(temp) > 2",
         show_execution_plan=True,
     )
+
+    print(plan)
+
     assert (
         plan == "FOREST_FIRES.loc[:, ['temp']].assign(__=1).groupby(['__'])"
         ".agg(**{'_col0': ('temp', 'min')}).reset_index(drop=True)"
@@ -1087,6 +1090,6 @@ def test_timestamps():
 if __name__ == "__main__":
     register_env_tables()
 
-    test_using_math()
+    test_having_one_condition()
 
     remove_env_tables()