get if_else working

improve result comparison

Author: JohnMount

data_algebra/db_model.py

@@ -44,6 +44,18 @@ def _db_is_bad_expr(dbmodel, expression):
     )
 
 
+def _db_if_else_expr(dbmodel, expression):
+    if_expr = dbmodel.expr_to_sql(expression.args[0], want_inline_parens=True)
+    x_expr = dbmodel.expr_to_sql(expression.args[1], want_inline_parens=True)
+    y_expr = dbmodel.expr_to_sql(expression.args[2], want_inline_parens=True)
+    return (
+        "CASE" +\
+            " WHEN " + if_expr + " THEN " + x_expr +\
+            " WHEN NOT " + if_expr + " THEN " + y_expr +\
+            " ELSE NULL END")
+
+
+
 def _db_neg_expr(dbmodel, expression):
     subexpr = dbmodel.expr_to_sql(expression.args[0], want_inline_parens=True)
     return "( -" + subexpr + " )"
@@ -53,6 +65,7 @@ def _db_neg_expr(dbmodel, expression):
     "is_null": _db_is_null_expr,
     "is_bad": _db_is_bad_expr,
     "neg": _db_neg_expr,
+    "if_else": _db_if_else_expr,
 }
 
 

data_algebra/expr_rep.py

@@ -1,8 +1,6 @@
 from typing import Union
 import collections
 
-import pandas
-
 import data_algebra.util
 import data_algebra.env
 
@@ -43,6 +41,15 @@ def __uop_expr__(self, op, *, params=None):
             raise TypeError("op is supposed to be a string")
         return Expression(op, (self,), params=params)
 
+    def __triop_expr__(self, op, x, y):
+        if not isinstance(op, str):
+            raise TypeError("op is supposed to be a string")
+        if not isinstance(x, Term):
+            x = Value(x)
+        if not isinstance(y, Term):
+            y = Value(y)
+        return Expression(op, (self, x, y), inline=False)
+
     # tree re-write
 
     def replace_view(self, view):
@@ -289,6 +296,9 @@ def is_null(self):
     def is_bad(self):
         return self.__uop_expr__("is_bad")
 
+    def if_else(self, x, y):
+        return self.__triop_expr__("if_else", x, y)
+
 
 class Value(Term):
     def __init__(self, value):
@@ -337,15 +347,12 @@ def get_column_names(self, columns_seen):
 }
 
 
-pandas_eval_env = {
-    "is_null": lambda x: pandas.isnull(x),
-    "is_bad": data_algebra.util.is_bad,
-}
-
-
 pd_formatters = {
     "is_bad": lambda expr: "@is_bad(" + expr.args[0].to_pandas() + ")",
     "is_null": lambda expr: "@is_null(" + expr.args[0].to_pandas() + ")",
+    "if_else": lambda expr: "@if_else(" + expr.args[0].to_pandas() +\
+                            ", " + expr.args[1].to_pandas() +\
+                            ", " + expr.args[2].to_pandas()+ ")",
     "neg": lambda expr: "-" + expr.args[0].to_pandas(want_inline_parens=True),
 }
 

data_algebra/pandas_model.py

@@ -1,11 +1,21 @@
+
+import numpy
 import pandas
 
 import data_algebra
+import data_algebra.util
 import data_algebra.data_model
 import data_algebra.expr_rep
 import data_algebra.data_ops
 
 
+pandas_eval_env = {
+    "is_null": lambda x: pandas.isnull(x),
+    "is_bad": data_algebra.util.is_bad,
+    "if_else": lambda c, x, y: numpy.where(c, x, y)
+}
+
+
 class PandasModel(data_algebra.data_model.DataModel):
     def __init__(self):
         data_algebra.data_model.DataModel.__init__(self)
@@ -52,7 +62,7 @@ def extend_step(self, op, *, data_map, eval_env):
                 op_src = opk.to_pandas()
                 res[k] = res.eval(
                     op_src,
-                    local_dict=data_algebra.expr_rep.pandas_eval_env,
+                    local_dict=pandas_eval_env,
                     global_dict=eval_env,
                 )
         else:

data_algebra/util.py

@@ -107,8 +107,11 @@ def equivalent_frames(
             if not all([ca_null[i] == cb_null[i] for i in range(a.shape[0])]):
                 return False
         if can_convert_v_to_numeric(ca):
+            ca = numpy.asarray(ca, dtype=float)
+            cb = numpy.asarray(cb, dtype=float)
             dif = ca - cb
-            if dif.abs().max() > float_tol:
+            dif = numpy.asarray([abs(d) for d in dif if not pandas.isnull(d)])
+            if dif.max() > float_tol:
                 return False
         else:
             if not all([ca[i] == cb[i] for i in range(a.shape[0])]):

tests/test_if_else.py

@@ -0,0 +1,45 @@
+
+import sqlite3
+
+import pandas
+
+from data_algebra.data_ops import *
+from data_algebra.SQLite import SQLiteModel
+import data_algebra.util
+
+
+def test_if_else():
+    d = pandas.DataFrame({
+        'a': [True, False],
+        'b': [1 ,2],
+        'c': [3, 4]
+    })
+
+    ops = TableDescription('d', ['a', 'b', 'c']). \
+        extend({'d': 'a.if_else(b, c)'})
+
+    expect = pandas.DataFrame({
+        'a': [True, False],
+        'b': [1, 2],
+        'c': [3, 4],
+        'd': [1, 4],
+        })
+
+    res_pandas = ops.transform(d)
+
+    assert data_algebra.util.equivalent_frames(res_pandas, expect)
+
+    db_model = SQLiteModel()
+
+    ops_sql = ops.to_sql(db_model)
+
+    conn = sqlite3.connect(":memory:")
+    db_model.prepare_connection(conn)
+
+    db_model.insert_table(conn, d, 'd')
+
+    res_db = db_model.read_query(conn, ops_sql)
+
+    conn.close()
+
+    assert data_algebra.util.equivalent_frames(res_db, expect)