Repository: matthewwardrop/formulaic

Issue:
No way to retrieve and reuse drop_index when na_action='drop'
Thank you for writing this library.
I have added pull request that enables this, I would appreciate your insight as I am not sure what is the best way to handle this.
Also maybe I missed something.
https://github.com/matthewwardrop/formulaic/pull/194

Hints from issue comments:
@matthewwardrop: @andrejmuhic I've merged in support for this as described in #197 :). Let me know if this doesn't work for you.@andrejmuhic: @matthewwardrop Thank you for this. I think this will work nicely.
I will test it out soon. I am not a big fan of passing results by modifying arguments in python but your solution seems more pragmatic.

Solution patch:
diff --git a/docsite/docs/guides/missing_data.ipynb b/docsite/docs/guides/missing_data.ipynb
new file mode 100644
index 00000000..338fad9d
--- /dev/null
+++ b/docsite/docs/guides/missing_data.ipynb
@@ -0,0 +1,416 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Sooner or later, you will encounter datasets with null values, and it is \n",
+    "important to know how their presence will impact your modeling. Formulaic\n",
+    "model matrix materialization procedures allow you to specify how you want nulls\n",
+    "to be handled. You can either:\n",
+    "\n",
+    "* Automatically drop null rows from the dataset (the default).\n",
+    "* Ignore nulls, and allow them to propagate naturally.\n",
+    "* Raise an exception when null values are encountered.\n",
+    "\n",
+    "You can specify the desired behaviour by passing an `NAAction` enum value (or\n",
+    "string value thereof) to the materialization methods (`model_matrix`, and \n",
+    "`*.get_model_matrix()`). Examples of each of these approaches is show below."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Dropping null rows (`NAAction.DROP`, or `\"drop\"`)\n",
+    "\n",
+    "This the default behaviour, and will result in any row with a null in any column\n",
+    "that is being used by the materialization being dropped from the resulting \n",
+    "dataset. For example:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>Intercept</th>\n",
+       "      <th>c</th>\n",
+       "      <th>C[T.b]</th>\n",
+       "      <th>C[T.c]</th>\n",
+       "      <th>C[T.d]</th>\n",
+       "      <th>C[T.e]</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>2.0</td>\n",
+       "      <td>1</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>5.0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   Intercept    c  C[T.b]  C[T.c]  C[T.d]  C[T.e]\n",
+       "0        1.0  1.0       0       0       0       0\n",
+       "1        1.0  2.0       1       0       0       0\n",
+       "4        1.0  5.0       0       0       0       1"
+      ]
+     },
+     "execution_count": 36,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from formulaic import model_matrix\n",
+    "from formulaic.materializers import NAAction\n",
+    "from pandas import DataFrame, Categorical\n",
+    "\n",
+    "df = DataFrame({\n",
+    "    \"c\": [1, 2, None, 4, 5],\n",
+    "    \"C\": Categorical([\"a\", \"b\", \"c\", None, \"e\"], categories=[\"a\", \"b\", \"c\", \"d\", \"e\"])\n",
+    "})\n",
+    "\n",
+    "model_matrix(\"c + C\", df, na_action=NAAction.DROP)\n",
+    "# Equivlent to:\n",
+    "#   * model_matrix(\"c + C\", df)\n",
+    "#   * model_matrix(\"c + C\", df, na_action=\"drop\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can also specify additional rows to drop using the `drop_rows` argument:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>Intercept</th>\n",
+       "      <th>c</th>\n",
+       "      <th>C[T.b]</th>\n",
+       "      <th>C[T.c]</th>\n",
+       "      <th>C[T.d]</th>\n",
+       "      <th>C[T.e]</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>2.0</td>\n",
+       "      <td>1</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   Intercept    c  C[T.b]  C[T.c]  C[T.d]  C[T.e]\n",
+       "1        1.0  2.0       1       0       0       0"
+      ]
+     },
+     "execution_count": 24,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_matrix(\"c + C\", df, drop_rows={0, 4})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note that the set passed to `drop_rows` is expected to be mutable, as it will be\n",
+    "updated with the indices of rows dropped automatically also; which can be useful\n",
+    "if you need to keep track of this information outside of the materialization\n",
+    "procedure."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{0, np.int64(2), np.int64(3), 4}"
+      ]
+     },
+     "execution_count": 25,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "drop_rows = {0, 4}\n",
+    "model_matrix(\"c + C\", df, drop_rows=drop_rows)\n",
+    "drop_rows"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Ignore nulls (`NAAction.IGNORE`, or `\"ignore\"`)\n",
+    "\n",
+    "If your modeling toolkit can handle the presence of nulls, or you otherwise want\n",
+    "to keep them in the dataset, you can pass `na_action = \"ignore\"` to the \n",
+    "materialization methods. This will allow null values to remain in columns, and \n",
+    "take no action to prevent the propagation of nulls."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>Intercept</th>\n",
+       "      <th>c</th>\n",
+       "      <th>C[T.b]</th>\n",
+       "      <th>C[T.c]</th>\n",
+       "      <th>C[T.d]</th>\n",
+       "      <th>C[T.e]</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>2.0</td>\n",
+       "      <td>1</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>NaN</td>\n",
+       "      <td>0</td>\n",
+       "      <td>1</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>4.0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>1.0</td>\n",
+       "      <td>5.0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   Intercept    c  C[T.b]  C[T.c]  C[T.d]  C[T.e]\n",
+       "0        1.0  1.0       0       0       0       0\n",
+       "1        1.0  2.0       1       0       0       0\n",
+       "2        1.0  NaN       0       1       0       0\n",
+       "3        1.0  4.0       0       0       0       0\n",
+       "4        1.0  5.0       0       0       0       1"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_matrix(\"c + C\", df, na_action=\"ignore\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note the `NaN` in the `c` column, and that `NaN` does NOT appear in the dummy\n",
+    "coding of C on row 3, consistent with standard implementations of dummy coding.\n",
+    "This could result in misleading model estimates, so care should be taken.\n",
+    "\n",
+    "You can combine this with `drop_rows`, as described above, to manually filter\n",
+    "out the null values you are concerned about."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Raise for null values (`NAAction.RAISE` or `\"raise\"`)\n",
+    "\n",
+    "If you are unwilling to risk the perils of dropping or ignoring null values, you\n",
+    "can instead opt to raise an exception whenever a null value is found. This can\n",
+    "prevent yourself from accidentally biasing your model, but also makes your code\n",
+    "more brittle. For example:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Error encountered while checking for nulls in `C`: `C` contains null values after evaluation.\n"
+     ]
+    }
+   ],
+   "source": [
+    "try:\n",
+    "    model_matrix(\"c + C\", df, na_action=\"raise\")\n",
+    "except Exception as e:\n",
+    "    print(e)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "As with ignoring nulls above, you can combine this raising behaviour with\n",
+    "`drop_rows` to manually filter out the null values that you feel you can safely\n",
+    "ignore, and then raise if any additional null values make it into your data."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/docsite/mkdocs.yml b/docsite/mkdocs.yml
index 98121674..1ae03a36 100644
--- a/docsite/mkdocs.yml
+++ b/docsite/mkdocs.yml
@@ -10,6 +10,7 @@ nav:
     - Introduction: guides/index.md
     - Quickstart: guides/quickstart.ipynb
     - How it works: guides/formulae.ipynb
+    - Handling Missing Data: guides/missing_data.ipynb
     - Formula Grammar: guides/grammar.md
     - Model Specs: guides/model_specs.ipynb
     - Transforms: guides/transforms.ipynb
diff --git a/formulaic/formula.py b/formulaic/formula.py
index eb422dc6..6dd450cc 100644
--- a/formulaic/formula.py
+++ b/formulaic/formula.py
@@ -220,6 +220,7 @@ def get_model_matrix(
         self,
         data: Any,
         context: Optional[Mapping[str, Any]] = None,
+        drop_rows: Optional[Set[int]] = None,
         **spec_overrides: Any,
     ) -> Union[ModelMatrix, Structured[ModelMatrix]]:
         """
@@ -230,13 +231,17 @@ def get_model_matrix(
             data: The data for which to build the model matrices.
             context: An additional mapping object of names to make available in
                 when evaluating formula term factors.
+            drop_rows: An optional set of row indices to drop from the model
+                matrix. If specified, it will also be updated during
+                materialization with any additional rows dropped due to null
+                values.
             spec_overrides: Any `ModelSpec` attributes to set/override. See
                 `ModelSpec` for more details.
         """
         from .model_spec import ModelSpec
 
         return ModelSpec.from_spec(self, **spec_overrides).get_model_matrix(
-            data, context=context
+            data, context=context, drop_rows=drop_rows
         )
 
     def differentiate(  # pylint: disable=redefined-builtin
diff --git a/formulaic/materializers/base.py b/formulaic/materializers/base.py
index 59a93894..deec86a3 100644
--- a/formulaic/materializers/base.py
+++ b/formulaic/materializers/base.py
@@ -157,6 +157,7 @@ def nrows(self) -> int:
     def get_model_matrix(
         self,
         spec: Union[FormulaSpec, ModelMatrix, ModelMatrices, ModelSpec, ModelSpecs],
+        drop_rows: Optional[Set[int]] = None,
         **spec_overrides: Any,
     ) -> Union[ModelMatrix, ModelMatrices]:
         from formulaic import ModelSpec
@@ -175,7 +176,7 @@ def get_model_matrix(
 
         # Step 1: Evaluate all factors and cache the results, keeping track of
         # which rows need dropping (if `self.config.na_action == 'drop'`).
-        drop_rows: Set[int] = set()
+        drop_rows: Set[int] = drop_rows if drop_rows is not None else set()
         for factor in factors:
             self._evaluate_factor(factor, factor_evaluation_model_spec, drop_rows)
         drop_rows: Sequence[int] = sorted(drop_rows)
diff --git a/formulaic/model_spec.py b/formulaic/model_spec.py
index f8c8ce0a..00a3ffef 100644
--- a/formulaic/model_spec.py
+++ b/formulaic/model_spec.py
@@ -311,6 +311,7 @@ def get_model_matrix(
         self,
         data: Any,
         context: Optional[Mapping[str, Any]] = None,
+        drop_rows: Optional[Set[int]] = None,
         **attr_overrides: Any,
     ) -> ModelMatrix:
         """
@@ -321,6 +322,10 @@ def get_model_matrix(
             data: The data for which to build the model matrices.
             context: An additional mapping object of names to make available in
                 when evaluating formula term factors.
+            drop_rows: An optional set of row indices to drop from the model
+                matrix. If specified, it will also be updated during
+                materialization with any additional rows dropped due to null
+                values.
             attr_overrides: Any `ModelSpec` attributes to override before
                 constructing model matrices. This is shorthand for first
                 running `ModelSpec.update(**attr_overrides)`.
@@ -333,7 +338,7 @@ def get_model_matrix(
             materializer = FormulaMaterializer.for_materializer(self.materializer)
         return materializer(
             data, context=context, **(self.materializer_params or {})
-        ).get_model_matrix(self)
+        ).get_model_matrix(self, drop_rows=drop_rows)
 
     def get_linear_constraints(self, spec: LinearConstraintSpec) -> LinearConstraints:
         """
@@ -410,6 +415,7 @@ def get_model_matrix(
         self,
         data: Any,
         context: Optional[Mapping[str, Any]] = None,
+        drop_rows: Optional[Set[int]] = None,
         **attr_overrides: Any,
     ) -> ModelMatrices:
         """
@@ -423,6 +429,10 @@ def get_model_matrix(
             data: The data for which to build the model matrices.
             context: An additional mapping object of names to make available in
                 when evaluating formula term factors.
+            drop_rows: An optional set of row indices to drop from the model
+                matrix. If specified, it will also be updated during
+                materialization with any additional rows dropped due to null
+                values.
             attr_overrides: Any `ModelSpec` attributes to override before
                 constructing model matrices. This is shorthand for first
                 running `ModelSpec.from_spec(model_specs, **attr_overrides)`.
@@ -431,7 +441,7 @@ def get_model_matrix(
 
         if attr_overrides:
             return ModelSpec.from_spec(self, **attr_overrides).get_model_matrix(
-                data, context=context
+                data, context=context, drop_rows=drop_rows
             )
 
         # Check whether we can generate model matrices jointly (i.e. all
@@ -469,7 +479,9 @@ def get_model_matrix(
         return cast(
             ModelMatrices,
             self._map(
-                lambda model_spec: model_spec.get_model_matrix(data, context=context),
+                lambda model_spec: model_spec.get_model_matrix(
+                    data, context=context, drop_rows=drop_rows
+                ),
                 as_type=ModelMatrices,
             ),
         )
diff --git a/formulaic/sugar.py b/formulaic/sugar.py
index b977e039..844f060e 100644
--- a/formulaic/sugar.py
+++ b/formulaic/sugar.py
@@ -1,4 +1,4 @@
-from typing import Any, Mapping, Union
+from typing import Any, Mapping, Optional, Set, Union
 
 from .formula import FormulaSpec
 from .model_matrix import ModelMatrices, ModelMatrix
@@ -11,6 +11,7 @@ def model_matrix(
     data: Any,
     *,
     context: Union[int, Mapping[str, Any]] = 0,
+    drop_rows: Optional[Set[int]] = None,
     **spec_overrides: Any,
 ) -> Union[ModelMatrix, ModelMatrices]:
     """
@@ -40,6 +41,9 @@ def model_matrix(
             means that all variables in the caller's scope should be made
             accessible when interpreting and evaluating formulae). Otherwise, a
             mapping from variable name to value is expected.
+        drop_rows: An optional set of row indices to drop from the model matrix.
+            If specified, it will also be updated during materialization with
+            any additional rows dropped due to null values.
         spec_overrides: Any `ModelSpec` attributes to set/override. See
             `ModelSpec` for more details.
 
@@ -49,5 +53,5 @@ def model_matrix(
     """
     _context = capture_context(context + 1) if isinstance(context, int) else context
     return ModelSpec.from_spec(spec, **spec_overrides).get_model_matrix(
-        data, context=_context
+        data, context=_context, drop_rows=drop_rows
     )


Test patch:
diff --git a/tests/materializers/test_pandas.py b/tests/materializers/test_pandas.py
index dc3da5d3..9e532906 100644
--- a/tests/materializers/test_pandas.py
+++ b/tests/materializers/test_pandas.py
@@ -480,3 +480,18 @@ def test_transform_state_with_inconsistent_formatting(self, data):
     def test_nested_transform_state(self, data):
         ms = PandasMaterializer(data).get_model_matrix("bs(bs(a))").model_spec
         assert {"bs(a)", "bs(bs(a))"}.issubset(ms.transform_state)
+
+    def test_drop_rows(self, data, data_with_nulls):
+        drop_rows = {0, 1}
+        mm = PandasMaterializer(data).get_model_matrix("a", drop_rows=drop_rows)
+        assert mm.shape == (1, 2)
+        assert list(mm.index) == [2]
+        assert drop_rows == {0, 1}
+
+        drop_rows = {0, 1}
+        mm = PandasMaterializer(data_with_nulls).get_model_matrix(
+            "a", drop_rows=drop_rows
+        )
+        assert mm.shape == (0, 2)
+        assert not list(mm.index)
+        assert drop_rows == {0, 1, 2}