Supported MAX/MIN/MAXIF/MINIF

Author: dx-tang

forms/core/forms.py

@@ -278,6 +278,27 @@ def compute_formula(self, formula_str: str, num_formulas: int = -1, **kwargs) ->
             print(f"An error occurred: {e}")
             traceback.print_exception(*sys.exc_info())
 
+    def print_sql_strings(self, formula_str: str, num_formulas: int = -1, **kwargs):
+        try:
+            root = parse_formula_str(formula_str)
+            validate(FunctionExecutor.DB_EXECUTOR, self.num_rows, self.num_columns, root)
+            root = PlanRewriter(self.db_config).rewrite_plan(root)
+
+            if num_formulas <= 0:
+                num_formulas = self.num_rows
+            exec_context = DBExecContext(
+                self.connection, self.cursor, self.base_table, START_ROW_ID, START_ROW_ID + num_formulas
+            )
+            executor = DBExecutor(self.db_config, exec_context, self.metrics_tracker)
+            sql_strings = executor.get_sql_strings(root)
+            executor.clean_up()
+
+            for s in sql_strings:
+                print(s) 
+        except FormSException as e:
+            print(f"An error occurred: {e}")
+            traceback.print_exception(*sys.exc_info())
+
     def print_workbook(self, num_rows=10, keep_original_labels=False):
         order_by_clause = ", ".join(self.db_config.order_key)
         query = f"SELECT * FROM {self.db_config.table_name} ORDER BY {order_by_clause} LIMIT {num_rows}"

forms/executor/dbexecutor/dbexecutor.py

@@ -49,6 +49,21 @@ def __init__(
         self.exec_context = exec_context
         self.metrics_tracker = metrics_tracker
 
+    def get_sql_strings(self, formula_plan: PlanNode) -> list:
+        exec_tree = from_plan_to_execution_tree(formula_plan, self.exec_context.base_table)
+        scheduler = Scheduler(exec_tree)
+        sql_strings = []
+        exec_subtree = scheduler.next_subtree()
+        is_root_subtree = not scheduler.has_next_subtree()
+        intermediate_table_name = (
+                exec_tree.intermediate_table_name if isinstance(exec_tree, DBFuncExecNode) else ""
+                )
+        sql_str = translate(
+            exec_subtree, self.exec_context, intermediate_table_name, is_root_subtree
+        ).as_string(self.exec_context.conn)
+        sql_strings.append(sql_str)
+        return sql_strings
+
     def execute_formula_plan(self, formula_plan: PlanNode) -> pd.DataFrame:
         exec_tree = from_plan_to_execution_tree(formula_plan, self.exec_context.base_table)
         scheduler = Scheduler(exec_tree, self.db_config.enable_pipelining)

forms/executor/dbexecutor/translation.py

@@ -242,6 +242,18 @@ def translate_if_function(
     subtree: DBFuncExecNode, exec_context: DBExecContext, base_table: Composable
 ) -> Composable:
     children = subtree.children
+    if isinstance(children[1], DBLitExecNode) and children[1].literal == '"NULL"':
+        return sql.SQL(
+            """CASE
+                WHEN {condition}
+                THEN NULL 
+                ELSE {false_result}
+           END"""
+           ).format(
+                condition=translate_window_clause(children[0], exec_context, base_table),
+                false_result=translate_window_clause(children[2], exec_context, base_table),
+           )       
+
     return sql.SQL(
         """CASE
                 WHEN {condition}
@@ -276,6 +288,14 @@ def translate_aggregate_functions(
                 function=sql.SQL(subtree.function.value),
                 agg_column=sql.SQL("+").join(sql.Identifier(col) for col in child.cols),
             )
+        elif subtree.function == Function.MAX or subtree.function == Function.MIN:
+            row_func = "MAX" if subtree.function == Function.MAX else "MIN"
+            col_func = "GREATEST" if subtree.function == Function.MAX else "LEAST"
+            agg_sql = sql.SQL("""{row_func}({col_func}({agg_column}))""").format(
+                row_func=sql.SQL(row_func),
+                col_func=sql.SQL(col_func),
+                agg_column=sql.SQL(",").join(sql.Identifier(col) for col in child.cols),
+            )
         elif subtree.function == Function.COUNT:
             agg_sql = sql.SQL("""{num_columns} * {function}({agg_column})""").format(
                 num_columns=sql.Literal(len(child.cols)),
@@ -326,7 +346,30 @@ def translate_aggregate_if_functions(
                     )
                     for i in range(len(input_child.cols))
                 )
-            )
+            ) 
+        elif subtree.function == Function.MAXIF or subtree.function == Function.MINIF:
+            row_func = "MAX" if subtree.function == Function.MAXIF else "MIN"
+            col_func = "GREATEST" if subtree.function == Function.MAXIF else "LEAST"
+            agg_sql = sql.SQL("""{row_func}({col_func}({agg_expression}))""").format(
+                row_func=sql.SQL(row_func),
+                col_func=sql.SQL(col_func),
+                agg_expression=sql.SQL(",").join(
+                    sql.SQL(
+                        """ 
+                            CASE
+                                WHEN {input_col}{condition}
+                                THEN {output_col}
+                                ELSE NULL
+                            END
+                        """
+                    ).format(
+                        input_col=sql.Identifier(input_cols[i]),
+                        condition=sql.SQL(condition),
+                        output_col=sql.Identifier(output_cols[i]),
+                    )
+                    for i in range(len(input_child.cols))
+                )
+            )            
         elif subtree.function == Function.COUNTIF:
             agg_sql = sql.SQL("""SUM({agg_expression})""").format(
                 agg_expression=sql.SQL("+").join(

forms/utils/functions.py

@@ -34,6 +34,8 @@ class Function(Enum):
     SUMIF = "sumif"
     COUNTIF = "countif"
     AVERAGEIF = "averageif"
+    MAXIF = "maxif"
+    MINIF = "minif"
 
     # Text Functions
     CONCAT = "concat"
@@ -234,9 +236,13 @@ def from_function_str(function_str: str) -> Function:
     Function.DIVIDE,
     # Aggregation Functions
     Function.SUM,
+    Function.MAX,
+    Function.MIN,
     Function.COUNT,
     Function.AVG,
     Function.SUMIF,
+    Function.MAXIF,
+    Function.MINIF,
     Function.COUNTIF,
     Function.AVERAGEIF,
     # Control Function
@@ -253,8 +259,8 @@ def from_function_str(function_str: str) -> Function:
     Function.INDEX,
 }
 
-DB_AGGREGATE_FUNCTIONS = {Function.SUM, Function.COUNT, Function.AVG}
-DB_AGGREGATE_IF_FUNCTIONS = {Function.SUMIF, Function.COUNTIF, Function.AVERAGEIF}
+DB_AGGREGATE_FUNCTIONS = {Function.SUM, Function.MAX, Function.MIN, Function.COUNT, Function.AVG}
+DB_AGGREGATE_IF_FUNCTIONS = {Function.SUMIF, Function.MAXIF, Function.MINIF, Function.COUNTIF, Function.AVERAGEIF}
 DB_CELL_REFERENCE_FUNCTIONS = ARITHMETIC_FUNCTIONS | {Function.IF} | COMPARISON_FUNCTIONS
 DB_LOOKUP_FUNCTIONS = {
     Function.VLOOKUP,

tests/db_tests/test_formula_to_sql.py

@@ -148,3 +148,6 @@ def test_if_function(get_wb):
     computed_df = wb.compute_formula("=IF(A1 < 3, B1, C1)")
     expected_df = pd.DataFrame({"row_id": [1, 2, 3, 4], "A": [2, 2, 4, 5]})
     assert np.array_equal(computed_df.values, expected_df.values, equal_nan=True)
+
+
+    

tests/db_tests/test_translation.py

@@ -0,0 +1,45 @@
+#  Copyright 2022-2023 The FormS Authors.
+#
+#  Licensed 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.
+
+import pytest
+import os
+import pandas as pd
+import numpy as np
+
+from forms.core.forms import from_db
+
+
+@pytest.fixture(scope="module")
+def get_wb():
+    wb = from_db(
+        host=os.getenv("POSTGRES_HOST"),
+        port=int(os.getenv("POSTGRES_PORT")),
+        username=os.getenv("POSTGRES_USER"),
+        password=os.getenv("POSTGRES_PASSWORD"),
+        db_name=os.getenv("POSTGRES_DB"),
+        table_name=os.getenv("POSTGRES_TEST_TABLE"),
+        primary_key=[os.getenv("POSTGRES_PRIMARY_KEY")],
+        order_key=[os.getenv("POSTGRES_ORDER_KEY")],
+        enable_rewriting=False,
+    )
+
+    # Yield the object to be used in tests
+    yield wb
+    # Close the DBWorkbook
+    wb.close()
+
+# Check out the SQL strings 
+def test_get_sql_strings(get_wb):
+    wb = get_wb
+    wb.print_sql_strings('=MAX(B1:C2)')