Fixed Large memory requirements for SimpleImputer strategy median #4794

python/cuml/thirdparty_adapters/adapters.py

@@ -295,16 +295,22 @@ def _masked_column_median(arr, masked_value):
     mask = _get_mask(arr, masked_value)
     if arr.size == 0:
         return cp.full(arr.shape[1], cp.nan)
-    arr_sorted = arr.copy()
     if not cp.isnan(masked_value):
+        arr_sorted = arr.copy()
         # If nan is not the missing value, any column with nans should
         # have a median of nan
         nan_cols = cp.any(cp.isnan(arr), axis=0)
         arr_sorted[mask] = cp.nan
+        arr_sorted.sort(axis=0)
     else:
         nan_cols = cp.full(arr.shape[1], False)
-    # nans are always sorted to end of array
-    arr_sorted = cp.sort(arr_sorted, axis=0)
+<<<<<<< HEAD
+        # nans are always sorted to end of array and the sort call 
+        # copies the data
+=======
+        # nans are always sorted to end of array and the sort call copies the data
+>>>>>>> febd55587f1aed5b9728210412472c5cd94d4f6c
+        arr_sorted = cp.sort(arr_sorted, axis=0)
 
     count_missing_values = mask.sum(axis=0)
     # Ignore missing values in determining "halfway" index of sorted