Repository: vaexio/vaex

Issue:
[FEATURE-REQUEST] OHE encoder to output columns with boolean dtype
OHE currently outputs int64 cols which is highly uneconomical. It is possible to cast it to int32 or bool of course but this adds another virtual column in the df state and it makes the state transfer slower.

This is probably nothing when dealing with small amount of data but when things get bigger, it helps a lot to save some space by downcasting integer and floats.

[FEATURE-REQUEST] OHE encoder to output columns with boolean dtype
OHE currently outputs int64 cols which is highly uneconomical. It is possible to cast it to int32 or bool of course but this adds another virtual column in the df state and it makes the state transfer slower.

This is probably nothing when dealing with small amount of data but when things get bigger, it helps a lot to save some space by downcasting integer and floats.



Discussion:
For OHE encoding, int64 definitely seems like an overkill. I would second the idea to instead return a boolean dtype or something smaller like uint8. Seems like a good opportunity for a PR! I may look into it this weekend if I get time.
Indeed, I agree, I can't say without digging in why it's an int64, but an int8 sounds like a good option to me. Looking forward to a PR @kmcentush !
This might be useful btw:
```
import numpy as np
x = np.arange(10)
np.where(x > 4, np.uint8(1), np.uint8(0))
```
For OHE encoding, int64 definitely seems like an overkill. I would second the idea to instead return a boolean dtype or something smaller like uint8. Seems like a good opportunity for a PR! I may look into it this weekend if I get time.
Indeed, I agree, I can't say without digging in why it's an int64, but an int8 sounds like a good option to me. Looking forward to a PR @kmcentush !
This might be useful btw:
```
import numpy as np
x = np.arange(10)
np.where(x > 4, np.uint8(1), np.uint8(0))
```

Resolving patch:
diff --git a/packages/vaex-core/vaex/functions.py b/packages/vaex-core/vaex/functions.py
--- a/packages/vaex-core/vaex/functions.py
+++ b/packages/vaex-core/vaex/functions.py
@@ -2561,7 +2561,7 @@ def default_label(properties):
 import vaex.arrow.numpy_dispatch
 
 @register_function()
-def where(condition, x, y):
+def where(condition, x, y, dtype=None):
     # special where support for strings
     # TODO: this should be replaced by an arrow compute function in the future
     if type(x) == str:
@@ -2611,6 +2611,18 @@ def where(condition, x, y):
         choices = vaex.array_types.concat([x, y])
         values = choices.take(indices)
         return values
+
+    # cast x and y
+    if dtype is not None:
+        if np.can_cast(x, dtype):
+            dtype_scalar = np.dtype(dtype)
+            x = dtype_scalar.type(x)
+        if np.can_cast(y, dtype):
+            dtype_scalar = np.dtype(dtype)
+            y = dtype_scalar.type(y)
+
     # default callback is on numpy
+    # where() respects the dtypes of x and y; ex: if x and y are 'uint8', the resulting array will also be 'uint8'
     ar = np.where(condition, x, y)
+
     return ar
diff --git a/packages/vaex-ml/vaex/ml/transformations.py b/packages/vaex-ml/vaex/ml/transformations.py
--- a/packages/vaex-ml/vaex/ml/transformations.py
+++ b/packages/vaex-ml/vaex/ml/transformations.py
@@ -5,7 +5,6 @@
 from .state import HasState
 import traitlets
 from vaex.utils import _ensure_strings_from_expressions
-import numpy as np
 import warnings
 
 sklearn = vaex.utils.optional_import("sklearn", modules=[
@@ -423,6 +422,8 @@ def fit(self, df):
             uniques.append(unique_values)
         self.uniques_ = uniques
 
+        # detect ability to downcast to uint8
+
     def transform(self, df):
         '''Transform a DataFrame with a fitted OneHotEncoder.
 
@@ -431,17 +432,19 @@ def transform(self, df):
         :rtype: DataFrame
         '''
         copy = df.copy()
+        downcast_uint8 = np.can_cast(self.one, np.uint8) and np.can_cast(self.zero, np.uint8)
+        dtype = 'uint8' if downcast_uint8 else None
         # for each feature, add a virtual column for each unique entry
         for i, feature in enumerate(self.features):
             for j, value in enumerate(self.uniques_[i]):
                 str_value = str(value) if value is not None else 'missing'
                 column_name = self.prefix + feature + '_' + str_value
                 if value is None:
-                    copy[column_name] = copy.func.where(copy[feature].ismissing(), self.one, self.zero)
+                    copy[column_name] = copy.func.where(copy[feature].ismissing(), self.one, self.zero, dtype=dtype)
                 elif isinstance(value, np.float) and np.isnan(value):
-                    copy[column_name] = copy.func.where(copy[feature].isnan(), self.one, self.zero)
+                    copy[column_name] = copy.func.where(copy[feature].isnan(), self.one, self.zero, dtype=dtype)
                 else:
-                    copy[column_name] = copy.func.where(copy[feature] == value, self.one, self.zero)
+                    copy[column_name] = copy.func.where(copy[feature] == value, self.one, self.zero, dtype=dtype)
         return copy
 
 


Test patch:
diff --git a/tests/ml/ml_test.py b/tests/ml/ml_test.py
--- a/tests/ml/ml_test.py
+++ b/tests/ml/ml_test.py
@@ -307,6 +307,13 @@ def test_one_hot_encoding_with_na(df_factory):
     assert df_train['y_31.0'].tolist() == [1, 0, 0, 1]
     assert df_train['y_nan'].tolist() == [0, 0, 1, 0]
 
+    assert(str(df_train.x_missing.dtype) == 'uint8')
+    assert(str(df_train.x_Michael.dtype) == 'uint8')
+    assert(str(df_train.x_Reggie.dtype) == 'uint8')
+    assert(str(df_train['y_23.0'].dtype) == 'uint8')
+    assert(str(df_train['y_31.0'].dtype) == 'uint8')
+    assert(str(df_train['y_nan'].dtype) == 'uint8')
+
     df_test = enc.transform(df_test)
     assert df_test.x_missing.tolist() == [0, 0, 1, 1]
     assert df_test.x_Michael.tolist() == [1, 0, 0, 0]
@@ -315,6 +322,13 @@ def test_one_hot_encoding_with_na(df_factory):
     assert df_test['y_31.0'].tolist() == [0, 1, 0, 0]
     assert df_test['y_nan'].tolist() == [0, 0, 1, 1]
 
+    assert(str(df_test.x_missing.dtype) == 'uint8')
+    assert(str(df_test.x_Michael.dtype) == 'uint8')
+    assert(str(df_test.x_Reggie.dtype) == 'uint8')
+    assert(str(df_test['y_23.0'].dtype) == 'uint8')
+    assert(str(df_test['y_31.0'].dtype) == 'uint8')
+    assert(str(df_test['y_nan'].dtype) == 'uint8')
+
 
 def test_maxabs_scaler(df_factory):
     x = np.array([-2.65395789, -7.97116295, -4.76729177, -0.76885033, -6.45609635])