Repository: ray-project/ray
Pull request: [data/preprocessors] Fix StandardScaler to handle NaN stats
<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

When `StandardScaler.stats_` has a column with `NaN` then short circuit transform for that columns.

## Related issue number

Closes #51243

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
    - [ ] I've added any new APIs to the API Reference. For example, if I added a
           method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(


Problem statement:
[Data] `StandardScaler.transform()` cannot handle non-`NAN` columns when `StandardScaler.stats_` is `NAN` for that column
### What happened + What you expected to happen

`StandardScaler.transform()` should ideally return `NAN` for columns which have `StandardScaler.stats_` as None. But the current implementation leads to -

```
TypeError: unsupported operand type(s) for -: 'float' and 'NoneType'
```

### Versions / Dependencies

-

### Reproduction script

```
import pandas as pd
import ray
from ray.data.preprocessors import StandardScaler


df_sample = pd.DataFrame({'A': [-1, 0], 'B': [None, None]})

ds_sample = ray.data.from_pandas(df_sample)
pp = StandardScaler(columns=['A', 'B'])
pp = pp.fit(ds_sample)

# `None` mean/std for C.
print(f"preprocessor stats: {pp.stats_}")


df_transform = pd.DataFrame({'A': [0, 0], 'B': [0, None]})

ds_transform = ray.data.from_pandas(df_transform)
print(pp.transform(ds_transform).show()) # Errors
```

### Issue Severity

High: It blocks me from completing my task.

Hints from issue discussion:

Commit message: Fix StandardScaler to handle NAN stats

Signed-off-by: Praveen Gorthy <praveeng@anyscale.com>

Patch to review:
diff --git a/python/ray/data/preprocessors/scaler.py b/python/ray/data/preprocessors/scaler.py
index 1a9ba9a95682f..103edaa5db27f 100644
--- a/python/ray/data/preprocessors/scaler.py
+++ b/python/ray/data/preprocessors/scaler.py
@@ -95,6 +95,10 @@ def column_standard_scaler(s: pd.Series):
             s_mean = self.stats_[f"mean({s.name})"]
             s_std = self.stats_[f"std({s.name})"]
 
+            if s_std is None and s_mean is None:
+                s[:] = None
+                return s
+
             # Handle division by zero.
             # TODO: extend this to handle near-zero values.
             if s_std == 0:
diff --git a/python/ray/data/tests/preprocessors/test_scaler.py b/python/ray/data/tests/preprocessors/test_scaler.py
index d67ad369057f9..a510a9d1e9500 100644
--- a/python/ray/data/tests/preprocessors/test_scaler.py
+++ b/python/ray/data/tests/preprocessors/test_scaler.py
@@ -1,5 +1,6 @@
 import pandas as pd
 import pytest
+import numpy as np
 
 import ray
 from ray.data.preprocessor import PreprocessorNotFittedException
@@ -263,33 +264,50 @@ def test_standard_scaler():
     col_a = [-1, 0, 1, 2]
     col_b = [1, 1, 5, 5]
     col_c = [1, 1, 1, None]
-    in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
-    ds = ray.data.from_pandas(in_df)
+    col_d = [None, None, None, None]
+    sample_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c, "D": col_d})
+    ds = ray.data.from_pandas(sample_df)
 
-    scaler = StandardScaler(["B", "C"])
+    scaler = StandardScaler(["B", "C", "D"])
 
     # Transform with unfitted preprocessor.
     with pytest.raises(PreprocessorNotFittedException):
         scaler.transform(ds)
 
     # Fit data.
-    scaler.fit(ds)
+    scaler = scaler.fit(ds)
     assert scaler.stats_ == {
         "mean(B)": 3.0,
         "mean(C)": 1.0,
+        "mean(D)": None,
         "std(B)": 2.0,
         "std(C)": 0.0,
+        "std(D)": None,
     }
 
     # Transform data.
-    transformed = scaler.transform(ds)
+    in_col_a = [-1, 0, 1, 2]
+    in_col_b = [1, 1, 5, 5]
+    in_col_c = [1, 1, 1, None]
+    in_col_d = [0, None, None, None]
+    in_df = pd.DataFrame.from_dict(
+        {"A": in_col_a, "B": in_col_b, "C": in_col_c, "D": in_col_d}
+    )
+    in_ds = ray.data.from_pandas(in_df)
+    transformed = scaler.transform(in_ds)
     out_df = transformed.to_pandas()
 
     processed_col_a = col_a
     processed_col_b = [-1.0, -1.0, 1.0, 1.0]
     processed_col_c = [0.0, 0.0, 0.0, None]
+    processed_col_d = [np.nan, np.nan, np.nan, np.nan]
     expected_df = pd.DataFrame.from_dict(
-        {"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
+        {
+            "A": processed_col_a,
+            "B": processed_col_b,
+            "C": processed_col_c,
+            "D": processed_col_d,
+        }
     )
 
     pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
@@ -298,8 +316,9 @@ def test_standard_scaler():
     pred_col_a = [1, 2, 3]
     pred_col_b = [3, 5, 7]
     pred_col_c = [0, 1, 2]
+    pred_col_d = [None, None, None]
     pred_in_df = pd.DataFrame.from_dict(
-        {"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
+        {"A": pred_col_a, "B": pred_col_b, "C": pred_col_c, "D": pred_col_d}
     )
 
     pred_out_df = scaler.transform_batch(pred_in_df)
@@ -307,11 +326,13 @@ def test_standard_scaler():
     pred_processed_col_a = pred_col_a
     pred_processed_col_b = [0.0, 1.0, 2.0]
     pred_processed_col_c = [-1.0, 0.0, 1.0]
+    pred_processed_col_d = [None, None, None]
     pred_expected_df = pd.DataFrame.from_dict(
         {
             "A": pred_processed_col_a,
             "B": pred_processed_col_b,
             "C": pred_processed_col_c,
+            "D": pred_processed_col_d,
         }
     )
 


Code review comments:
diff_hunk: @@ -95,6 +95,10 @@ def column_standard_scaler(s: pd.Series):
             s_mean = self.stats_[f"mean({s.name})"]
             s_std = self.stats_[f"std({s.name})"]
 
+            if s_std is None and s_mean is None:
+                s[:] = None
path: python/ray/data/preprocessors/scaler.py
text: @user1:
I think this should be set to np.nan instead?

```python
In [5]: from sklearn.preprocessing import StandardScaler
   ...: data = [[0, 0], [0, 0], [1, None], [1, 1]]
   ...: scaler = StandardScaler()
   ...: print(scaler.fit(data))
   ...: print(scaler.mean_)
   ...: print(scaler.transform(data))
   ...: print(scaler.transform([[2, 2]]))
StandardScaler()
[0.5        0.33333333]
[[-1.         -0.70710678]
 [-1.         -0.70710678]
 [ 1.                 nan]
 [ 1.          1.41421356]]
[[3.         3.53553391]]
```