Repository: andersbogsnes/ml_tooling

Issue:
ENH: Datasets should check cached `.x` and `.y` attributes when dumping data
When dumping data from one Dataset to another, the dataset always fetches a new copy of the data and dumps it. This can potentially take time and is not necessary if the data is already cached

**Describe the solution you'd like**
All data related functionality, such as `make_prediction` or  `dump_data` should check for `x` and `y` and copy those instead of getting a fresh dataset. Should be able to disable this behaviour with a `use_cache` flag

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454?src=pr&el=h1) Report
> Merging [#454](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454?src=pr&el=desc) into [develop](https://codecov.io/gh/andersbogsnes/ml_tooling/commit/9d950c9b31856c0771094f9ee46e6d34566c6e10?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `100%`.

[![Impacted file tree graph](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454/graphs/tree.svg?width=650&token=ucKtL3BlxO&height=150&src=pr)](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454?src=pr&el=tree)

```diff
@@          Coverage Diff           @@
##           develop   #454   +/-   ##
======================================
  Coverage      100%   100%           
======================================
  Files           55     55           
  Lines         1412   1412           
  Branches       164    164           
======================================
  Hits          1412   1412
```


| [Impacted Files](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/ml\_tooling/data/sql.py](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454/diff?src=pr&el=tree#diff-c3JjL21sX3Rvb2xpbmcvZGF0YS9zcWwucHk=) | `100% <100%> (ø)` | :arrow_up: |
| [src/ml\_tooling/baseclass.py](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454/diff?src=pr&el=tree#diff-c3JjL21sX3Rvb2xpbmcvYmFzZWNsYXNzLnB5) | `100% <100%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454?src=pr&el=footer). Last update [9d950c9...9baf40e](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/454?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5bacac4..04dce69 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ human friendly manner
 - Dataset now verifies that `load_training_data` and `load_prediction_data` do not return empty
 - Added a missing data visualization to `Dataset.plot`
 - FillNA now accepts a `is_nan`flag which adds a flag indicating that a value was missing
+- `Model.make_prediction` now accepts a `use_cache`flag to score everything in cached `.x`
 
 # v0.10.2
 - Fixed type inferences from data to sql in _load_data
diff --git a/src/ml_tooling/baseclass.py b/src/ml_tooling/baseclass.py
index 239a12e..e62a5ff 100644
--- a/src/ml_tooling/baseclass.py
+++ b/src/ml_tooling/baseclass.py
@@ -227,6 +227,7 @@ class Model:
         *args,
         proba: bool = False,
         use_index: bool = False,
+        use_cache: bool = False,
         **kwargs,
     ) -> pd.DataFrame:
         """
@@ -246,6 +247,10 @@ class Model:
         use_index: bool
             Whether the index from the prediction data should be used for the result.
 
+        use_cache: bool
+            Whether or not to use the cached data in dataset to make predictions.
+            Useful for seeing probability distributions of the model
+
         Returns
         -------
         pd.DataFrame
@@ -255,7 +260,9 @@ class Model:
             raise MLToolingError(
                 f"{self.estimator_name} does not have a `predict_proba` method"
             )
-        x = data._load_prediction_data(*args, **kwargs)
+
+        x = data.x if use_cache else data._load_prediction_data(*args, **kwargs)
+
         try:
             if proba:
                 data = self.estimator.predict_proba(x)
diff --git a/src/ml_tooling/data/sql.py b/src/ml_tooling/data/sql.py
index d17aed8..03ef3d1 100644
--- a/src/ml_tooling/data/sql.py
+++ b/src/ml_tooling/data/sql.py
@@ -107,7 +107,7 @@ class SQLDataset(Dataset, metaclass=abc.ABCMeta):
         with self.create_connection() as conn:
             return super()._load_prediction_data(*args, conn=conn, **kwargs)
 
-    def _dump_data(self) -> pd.DataFrame:
+    def _dump_data(self, use_cache=False) -> pd.DataFrame:
         """
         Reads the underlying SQL table and returns a DataFrame
 
@@ -116,6 +116,7 @@ class SQLDataset(Dataset, metaclass=abc.ABCMeta):
         pd.DataFrame
 
         """
+
         logger.info(f"Dumping data from {self.table}")
         logger.debug(f"Dumping data from {self.engine}/{self.table.name}")
         stmt = sa.select([self.table])


Test patch:
diff --git a/tests/test_baseclass.py b/tests/test_baseclass.py
index 0e82975..77b3022 100644
--- a/tests/test_baseclass.py
+++ b/tests/test_baseclass.py
@@ -53,85 +53,6 @@ class TestBaseClass:
         ):
             Model({})
 
-    def test_make_prediction_errors_when_model_is_not_fitted(self, train_iris_dataset):
-        with pytest.raises(MLToolingError, match="You haven't fitted the estimator"):
-            model = Model(LinearRegression())
-            model.make_prediction(train_iris_dataset, 5)
-
-    def test_make_prediction_with_regression_sqldataset_works_as_expected(
-        self, boston_sqldataset, loaded_boston_db
-    ):
-        dataset = boston_sqldataset(loaded_boston_db, schema=None)
-        dataset.create_train_test(stratify=False)
-        model = Model(LinearRegression())
-        model.train_estimator(dataset)
-
-        result = model.make_prediction(dataset, 0)
-
-        assert result.shape == (1, 1)
-        assert result.columns.tolist() == ["Prediction"]
-
-    def test_make_prediction_with_classification_sqldataset_works_as_expected(
-        self, iris_sqldataset, loaded_iris_db
-    ):
-        dataset = iris_sqldataset(loaded_iris_db, schema=None)
-        dataset.create_train_test()
-        model = Model(LogisticRegression(solver="lbfgs"))
-        model.train_estimator(dataset)
-
-        result = model.make_prediction(dataset, 0, proba=True)
-
-        assert result.shape == (1, 2)
-        assert result.columns.tolist() == ["Probability Class 0", "Probability Class 1"]
-
-    def test_make_prediction_errors_if_asked_for_proba_without_predict_proba_method(
-        self, train_iris_dataset
-    ):
-        with pytest.raises(
-            MLToolingError, match="LinearRegression does not have a `predict_proba`"
-        ):
-            model = Model(LinearRegression())
-            model.train_estimator(train_iris_dataset)
-            model.make_prediction(train_iris_dataset, 5, proba=True)
-
-    @pytest.mark.parametrize("use_index, expected_index", [(False, 0), (True, 5)])
-    def test_make_prediction_returns_prediction_if_proba_is_false(
-        self,
-        classifier: Model,
-        use_index: bool,
-        expected_index: int,
-        train_iris_dataset,
-    ):
-        results = classifier.make_prediction(
-            train_iris_dataset, 5, proba=False, use_index=use_index
-        )
-        assert isinstance(results, pd.DataFrame)
-        assert 2 == results.ndim
-        assert np.all((results == 1) | (results == 0))
-        assert np.all(np.sum(results, axis=1) == 0)
-        assert results.index == pd.RangeIndex(
-            start=expected_index, stop=expected_index + 1, step=1
-        )
-
-    @pytest.mark.parametrize("use_index, expected_index", [(False, 0), (True, 5)])
-    def test_make_prediction_returns_proba_if_proba_is_true(
-        self,
-        classifier: Model,
-        use_index: bool,
-        expected_index: int,
-        train_iris_dataset,
-    ):
-        results = classifier.make_prediction(
-            train_iris_dataset, 5, proba=True, use_index=use_index
-        )
-        assert isinstance(results, pd.DataFrame)
-        assert 2 == results.ndim
-        assert np.all((results <= 1) & (results >= 0))
-        assert np.all(np.sum(results, axis=1) == 1)
-        assert results.index == pd.RangeIndex(
-            start=expected_index, stop=expected_index + 1, step=1
-        )
-
     def test_default_metric_getter_works_as_expected_classifier(self):
         rf = Model(RandomForestClassifier(n_estimators=10))
         assert rf.config.CLASSIFIER_METRIC == "accuracy"
@@ -723,3 +644,92 @@ class TestGridSearch:
             _, _ = classifier.gridsearch(
                 train_iris_dataset, param_grid={"clf__penalty": ["l1", "l2"]}
             )
+
+
+class TestMakePrediction:
+    def test_make_prediction_errors_when_model_is_not_fitted(self, train_iris_dataset):
+        model = Model(LinearRegression())
+        with pytest.raises(MLToolingError, match="You haven't fitted the estimator"):
+            model.make_prediction(train_iris_dataset, 5)
+
+    def test_make_prediction_with_regression_sqldataset_works_as_expected(
+        self, boston_sqldataset, loaded_boston_db
+    ):
+        dataset = boston_sqldataset(loaded_boston_db, schema=None)
+        dataset.create_train_test(stratify=False)
+        model = Model(LinearRegression())
+        model.train_estimator(dataset)
+
+        result = model.make_prediction(dataset, 0)
+
+        assert result.shape == (1, 1)
+        assert result.columns.tolist() == ["Prediction"]
+
+    def test_make_prediction_with_classification_sqldataset_works_as_expected(
+        self, iris_sqldataset, loaded_iris_db
+    ):
+        dataset = iris_sqldataset(loaded_iris_db, schema=None)
+        dataset.create_train_test()
+        model = Model(LogisticRegression(solver="lbfgs"))
+        model.train_estimator(dataset)
+
+        result = model.make_prediction(dataset, 0, proba=True)
+
+        assert result.shape == (1, 2)
+        assert result.columns.tolist() == ["Probability Class 0", "Probability Class 1"]
+
+    def test_make_prediction_errors_if_asked_for_proba_without_predict_proba_method(
+        self, train_iris_dataset
+    ):
+        with pytest.raises(
+            MLToolingError, match="LinearRegression does not have a `predict_proba`"
+        ):
+            model = Model(LinearRegression())
+            model.train_estimator(train_iris_dataset)
+            model.make_prediction(train_iris_dataset, 5, proba=True)
+
+    @pytest.mark.parametrize("use_index, expected_index", [(False, 0), (True, 5)])
+    def test_make_prediction_returns_prediction_if_proba_is_false(
+        self,
+        classifier: Model,
+        use_index: bool,
+        expected_index: int,
+        train_iris_dataset,
+    ):
+        results = classifier.make_prediction(
+            train_iris_dataset, 5, proba=False, use_index=use_index
+        )
+        assert isinstance(results, pd.DataFrame)
+        assert 2 == results.ndim
+        assert np.all((results == 1) | (results == 0))
+        assert np.all(np.sum(results, axis=1) == 0)
+        assert results.index == pd.RangeIndex(
+            start=expected_index, stop=expected_index + 1, step=1
+        )
+
+    @pytest.mark.parametrize("use_index, expected_index", [(False, 0), (True, 5)])
+    def test_make_prediction_returns_proba_if_proba_is_true(
+        self,
+        classifier: Model,
+        use_index: bool,
+        expected_index: int,
+        train_iris_dataset,
+    ):
+        results = classifier.make_prediction(
+            train_iris_dataset, 5, proba=True, use_index=use_index
+        )
+        assert isinstance(results, pd.DataFrame)
+        assert 2 == results.ndim
+        assert np.all((results <= 1) & (results >= 0))
+        assert np.all(np.sum(results, axis=1) == 1)
+        assert results.index == pd.RangeIndex(
+            start=expected_index, stop=expected_index + 1, step=1
+        )
+
+    def test_make_prediction_uses_cache_if_set(self, classifier, iris_df):
+        mock_dataset = MagicMock(spec=Dataset)
+        mock_dataset.x = iris_df.head(5).drop(columns="target")
+
+        result = classifier.make_prediction(mock_dataset, use_cache=True)
+        assert len(result) == 5
+        mock_dataset._load_prediction_data.assert_not_called()