FIX: showstopper bug with NA handling and data-independent builders

njsmith · njsmith · commit b0ee5efcc577 · 2013-08-20T16:33:55.000+01:00
Some formulas, like "~ 1", don't actually depend on the passed-in data. So given a formula like "y ~ 1", what we do is build the y matrix for the LHS, and then build a 1x1 matrix for the RHS, and then broadcast the RHS matrix to match the left. However, this broadcasting logic was broadcasting to the *original* size of the LHS matrix, *before* NA removal. One result: expressions like dmatrices("y ~ 1") were returning matrices with different numbers of rows, whenever y had missing values. Another result: dmatrix(..., result_type="dataframe") was totally broken in the presence of NAs, because RHS-only formulas still have an invisible zero-column LHS which is calculated and then discarded. But this was calculated with the wrong number of rows, and then when we tried to convert it to a zero-column DataFrame, its shape didn't match the index that we tried to put on it. (Because NA removal *was* correctly affecting the index.) Fixes pydatagh-22.
diff --git a/patsy/build.py b/patsy/build.py
@@ -918,13 +918,19 @@ def build_design_matrices(builders, data,
     # Handle NAs
     values = evaluator_to_values.values()
     is_NAs = evaluator_to_isNAs.values()
+    # num_rows is None iff evaluator_to_values (and associated sets like
+    # 'values') are empty, i.e., we have no actual evaluators involved
+    # (formulas like "~ 1").
     if return_type == "dataframe" and num_rows is not None:
         if pandas_index is None:
             pandas_index = np.arange(num_rows)
         values.append(pandas_index)
         is_NAs.append(np.zeros(len(pandas_index), dtype=bool))
     origins = [evaluator.factor.origin for evaluator in evaluator_to_values]
     new_values = NA_action.handle_NA(values, is_NAs, origins)
+    # NA_action may have changed the number of rows.
+    if num_rows is not None:
+        num_rows = new_values[0].shape[0]
     if return_type == "dataframe" and num_rows is not None:
         pandas_index = new_values.pop()
     evaluator_to_values = dict(zip(evaluator_to_values, new_values))
diff --git a/patsy/test_highlevel.py b/patsy/test_highlevel.py
@@ -645,12 +645,34 @@ def raise_patsy_error(x):
 def test_dmatrix_NA_action():
     data = {"x": [1, 2, 3, np.nan], "y": [np.nan, 20, 30, 40]}
 
-    mat = dmatrix("x + y", data=data)
-    assert np.array_equal(mat, [[1, 2, 20],
-                                [1, 3, 30]])
-    assert_raises(PatsyError, dmatrix, "x + y", data=data, NA_action="raise")
-
-    lmat, rmat = dmatrices("y ~ x", data=data)
-    assert np.array_equal(lmat, [[20], [30]])
-    assert np.array_equal(rmat, [[1, 2], [1, 3]])
-    assert_raises(PatsyError, dmatrices, "y ~ x", data=data, NA_action="raise")
+    for return_type in ["matrix", "dataframe"]:
+        mat = dmatrix("x + y", data=data, return_type=return_type)
+        assert np.array_equal(mat, [[1, 2, 20],
+                                    [1, 3, 30]])
+        if return_type == "dataframe":
+            assert mat.index.equals([1, 2])
+        assert_raises(PatsyError, dmatrix, "x + y", data=data,
+                      return_type=return_type,
+                      NA_action="raise")
+
+        lmat, rmat = dmatrices("y ~ x", data=data, return_type=return_type)
+        assert np.array_equal(lmat, [[20], [30]])
+        assert np.array_equal(rmat, [[1, 2], [1, 3]])
+        if return_type == "dataframe":
+            assert lmat.index.equals([1, 2])
+            assert rmat.index.equals([1, 2])
+        assert_raises(PatsyError,
+                      dmatrices, "y ~ x", data=data, return_type=return_type,
+                      NA_action="raise")
+
+        # Initial release for the NA handling code had problems with
+        # non-data-dependent matrices like "~ 1".
+        lmat, rmat = dmatrices("y ~ 1", data=data, return_type=return_type)
+        assert np.array_equal(lmat, [[20], [30], [40]])
+        assert np.array_equal(rmat, [[1], [1], [1]])
+        if return_type == "dataframe":
+            assert lmat.index.equals([1, 2, 3])
+            assert rmat.index.equals([1, 2, 3])
+        assert_raises(PatsyError,
+                      dmatrices, "y ~ 1", data=data, return_type=return_type,
+                      NA_action="raise")
