Tests on Postgres (#42)

Author: amol-

.github/workflows/tests.yml

@@ -27,6 +27,13 @@ jobs:
         run: |
           uv sync --dev
 
+      - name: Set up PostgreSQL
+        uses: harmon758/postgresql-action@v1
+        with:
+          postgresql db: mustelatestdb
+          postgresql user: mustelatestuser
+          postgresql password: mustelatestpassword
+
       - name: Run Test Suite
         run: |
           uv run pytest -v --tb=short --disable-warnings --maxfail=1 --cov=mustela

pyproject.toml

@@ -42,6 +42,9 @@ test = [
   "ibis-framework[duckdb]>=5.1.0",
   "pytest-cov>=5.0.0",
   "pytest>=8.3.2",
+  "sqlalchemy",
+  "psycopg2",
+  "duckdb",
 ]
 
 [tool.uv]
@@ -65,6 +68,9 @@ dev-dependencies = [
     "pydot",
     "onnxruntime",
     "onnxscript",
+    "sqlalchemy",
+    "psycopg2",
+    "duckdb",
 ]
 
 

src/mustela/translation/optimizer.py

@@ -174,8 +174,9 @@ def fold_case(self, expr: ibis.Value | ibis.Deferred) -> ibis.Value:
             else:
                 return op.default.to_expr()
         elif len(op.cases) == 1 and results_are_literals and possible_values == {1, 0}:
-            # results are 1 or 0, we can fold it to a boolean
-            # expression.
+            # results are 1 or 0, we can fold it to a boolean expression.
+            # FIXME: This doesn't work on postgresql so we need to disable it for the moment.
+            return expr
             if op.results[0].value == 1:
                 return (op.cases[0].to_expr()).cast("float64")
             else:

src/mustela/translation/steps/onehotencoder.py

@@ -32,11 +32,9 @@ def process(self) -> None:
             raise ValueError("OneHotEncoder: input expression not found")
 
         casted_variables = [
-            self._optimizer.fold_cast(
-                typing.cast(ibis.expr.types.BooleanValue, (input_expr == cat)).cast(
-                    "float64"
-                )
-            ).name(self.variable_unique_short_alias("onehot"))
+            ibis.ifelse(input_expr == cat, 1, 0)
+            .cast("float64")
+            .name(self.variable_unique_short_alias("oh"))
             for cat in cats
         ]
 

src/mustela/translation/steps/softmax.py

@@ -53,7 +53,7 @@ def compute_softmax(
         if isinstance(data, VariablesGroup):
             data = NumericVariablesGroup(data)
             max_value = ibis.greatest(*data.values()).name(
-                translator.variable_unique_short_alias("sfmmax")
+                translator.variable_unique_short_alias("sfmx")
             )
             translator.preserve(max_value)
 

src/mustela/translation/steps/trees/classifier.py

@@ -103,20 +103,20 @@ def build_classifier(
         else:
             ordered_features = typing.cast(list[ibis.Value], [input_expr])
         ordered_features = [
-            feature.name(self.variable_unique_short_alias("tclass"))
+            feature.name(self.variable_unique_short_alias("tcl"))
             for feature in ordered_features
         ]
         ordered_features = self.preserve(*ordered_features)
 
         def build_tree_case(node: dict) -> dict[str | int, ibis.Expr]:
             # Leaf node, return the votes
             if node["mode"] == "LEAF":
-                votes = {}
-                for clslabel in classlabels:
-                    # We can assume missing class = weight 0
-                    # The optimizer will remove this if both true and false have 0.
-                    votes[clslabel] = ibis.literal(node["weight"].get(clslabel, 0))
-                return votes
+                # We can assume missing class = weight 0
+                # The optimizer will remove this if both true and false have 0.
+                return {
+                    clslabel: ibis.literal(node["weight"].get(clslabel, 0.0))
+                    for clslabel in classlabels
+                }
 
             # Branch node, build a CASE statement
             feature_expr = ordered_features[node["feature_id"]]

src/mustela/translation/steps/trees/regressor.py

@@ -46,7 +46,7 @@ def build_regressor(self, input_expr: VariablesGroup | ibis.Expr) -> ibis.Expr:
         else:
             ordered_features = typing.cast(list[ibis.Value], [input_expr])
         ordered_features = [
-            feature.name(self.variable_unique_short_alias("tclass"))
+            feature.name(self.variable_unique_short_alias("tcl"))
             for feature in ordered_features
         ]
         ordered_features = self.preserve(*ordered_features)

tests/test_pipeline_e2e.py

@@ -1,6 +1,7 @@
 import sqlite3
 
 import duckdb
+import sqlalchemy
 import numpy as np
 import pandas as pd
 import pytest
@@ -41,7 +42,7 @@ def diabetes_data(self):
         df = pd.concat([X, y], axis=1)
         return df, feature_names
 
-    @pytest.fixture(params=["duckdb", "sqlite"])
+    @pytest.fixture(params=["duckdb", "sqlite", "postgres"])
     def db_connection(self, request):
         dialect = request.param
         if dialect == "duckdb":
@@ -52,13 +53,19 @@ def db_connection(self, request):
             conn = sqlite3.connect(":memory:")
             yield conn, dialect
             conn.close()
+        elif dialect == "postgres":
+            try:
+                conn = sqlalchemy.create_engine("postgresql://mustelatestuser:mustelatestpassword@localhost:5432/mustelatestdb")
+            except (sqlalchemy.exc.OperationalError, ImportError):
+                pytest.skip("Postgres database not available")
+            yield conn, dialect
+            conn.dispose()
 
     def execute_sql(self, sql, conn, dialect, data):
         if dialect == "duckdb":
             conn.execute("CREATE TABLE data AS SELECT * FROM data")
-            # print(conn.execute("SELECT * FROM data").fetchdf())
             result = conn.execute(sql).fetchdf()
-        elif dialect == "sqlite":
+        elif dialect in ("sqlite", "postgres"):
             data.to_sql("data", conn, index=False, if_exists="replace")
             result = pd.read_sql(sql, conn)
         return result