Repository: ludwig-ai/ludwig
Pull request: Removes empty partitions after dropping rows and splitting datasets
This PR addresses two separate issues: https://github.com/ludwig-ai/ludwig/issues/2324 and https://github.com/ludwig-ai/ludwig/issues/2308.

The issues are addressed by culling empty partitions from the Dask DataFrame at two points: (1) after dropping rows with NaNs (part of the DROP_ROWS missing value strategy) and (2) after splitting the dataset into train/val/test.

In order to maintain/increase performance, we add a persist call at the end of `build_dataset`, which makes it relatively inexpensive to compute the length of partitions repeatedly downstream.

Problem statement:
Metadata mismatch while calculating overall stats
On occassion, I see the following error in `model.evaluate()` while computing overall stats using a sample of the dataframe. I haven't been able to reproduce this error consistently, but it does surface from time to time. 

This was seen while training a small image dataset with only 1 image input feature and a categorical output feature.

```
Full computation of stats only supported for pandas dataframes. Sampling the first 10000 rows of the feature and target dataframes for computing overall stats.
Traceback (most recent call last):
  File "/home/ray/train_ispy2.py", line 145, in <module>
    app()
  File "/home/ray/train_ispy2.py", line 126, in app
    eval_stats, train_stats, _, _ = model.experiment(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/api.py", line 1147, in experiment
    eval_stats, _, _ = self.evaluate(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/api.py", line 912, in evaluate
    overall_stats = calculate_overall_stats(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/models/predictor.py", line 309, in calculate_overall_stats
    feature_df = feature_df.head(n=STATS_SAMPLE_SIZE, npartitions=-1, compute=True)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/dataframe/core.py", line 1219, in head
    return self._head(n=n, npartitions=npartitions, compute=compute, safe=safe)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/dataframe/core.py", line 1253, in _head
    result = result.compute()
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/base.py", line 312, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/base.py", line 600, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ray/util/dask/scheduler.py", line 202, in ray_dask_get
    result = ray_get_unpack(object_refs, progress_bar_actor=pb_actor)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ray/util/dask/scheduler.py", line 503, in ray_get_unpack
    computed_result = get_result(object_refs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ray/util/dask/scheduler.py", line 490, in get_result
    return ray.get(object_refs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ray/_private/client_mode_hook.py", line 105, in wrapper
    return func(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ray/worker.py", line 1809, in get
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(ValueError): ray::dask:('head-200-10000-rename-814aae6336586f7ac43cb63c641e157b', 0) (pid=22528, ip=192.168.13.44)
  At least one of the input arguments for this task could not be computed:
ray.exceptions.RayTaskError: ray::dask:('head-partial-10000-rename-814aae6336586f7ac43cb63c641e157b', 156) (pid=22280, ip=192.168.13.44)
  At least one of the input arguments for this task could not be computed:
ray.exceptions.RayTaskError: ray::dask:('block_to_df-rename-814aae6336586f7ac43cb63c641e157b', 156) (pid=22276, ip=192.168.13.44)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ray/util/dask/scheduler.py", line 431, in dask_task_wrapper
    result = func(*actual_args)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/optimization.py", line 990, in __call__
    return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 149, in get
    result = _execute_task(task, cache)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 119, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 119, in <genexpr>
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 113, in _execute_task
    return [_execute_task(a, cache) for a in arg]
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 113, in <listcomp>
    return [_execute_task(a, cache) for a in arg]
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 119, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 119, in <genexpr>
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 113, in _execute_task
    return [_execute_task(a, cache) for a in arg]
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 113, in <listcomp>
    return [_execute_task(a, cache) for a in arg]
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/core.py", line 119, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/ray/anaconda3/lib/python3.8/site-packages/dask/dataframe/utils.py", line 400, in check_meta
    raise ValueError(
ValueError: Metadata mismatch found in `from_delayed`.
```

Hints from issue discussion:
Update: This might be an issue on Ray's end - https://docs.ray.io/en/latest/_modules/ray/data/dataset.html#Dataset.to_dask 

@arnavgarg1 to create an reproducible script for this error and then create an issue on Ray's GitHub to get this fixed. In the mean time, we might be manually able to override/customize the `dask.from_delayed` call to get past this issue.
@geoffreyangus found that this issue arises because of empty partitions after splitting the dataset. @geoffreyangus will create a follow-up PR to address this

Commit message: Merge branch 'master' into remove-empty-partitions

Patch to review:
diff --git a/ludwig/data/dataframe/dask.py b/ludwig/data/dataframe/dask.py
index 783bbe893ab..a8fbcd86496 100644
--- a/ludwig/data/dataframe/dask.py
+++ b/ludwig/data/dataframe/dask.py
@@ -20,8 +20,11 @@
 import dask
 import dask.array as da
 import dask.dataframe as dd
+import ray
 import ray.data
 from dask.diagnostics import ProgressBar
+from ray.data.block import Block, BlockAccessor
+from ray.util.client.common import ClientObjectRef
 from ray.util.dask import ray_dask_get
 
 from ludwig.data.dataframe.base import DataFrameEngine
@@ -114,6 +117,23 @@ def split(self, df, probabilities):
         slices = df.partitions
         return split_by_slices(slices, n, probabilities)
 
+    def remove_empty_partitions(self, df):
+        # Reference: https://stackoverflow.com/questions/47812785/remove-empty-partitions-in-dask
+        ll = list(df.map_partitions(len).compute())
+        if all([ll_i > 0 for ll_i in ll]):
+            return df
+
+        df_delayed = df.to_delayed()
+        df_delayed_new = list()
+        empty_partition = None
+        for ix, n in enumerate(ll):
+            if n == 0:
+                empty_partition = df.get_partition(ix)
+            else:
+                df_delayed_new.append(df_delayed[ix])
+        df = dd.from_delayed(df_delayed_new, meta=empty_partition)
+        return df
+
     def to_parquet(self, df, path, index=False):
         with ProgressBar():
             df.to_parquet(
@@ -129,7 +149,24 @@ def to_ray_dataset(self, df):
         return from_dask(df)
 
     def from_ray_dataset(self, dataset) -> dd.DataFrame:
-        return dataset.to_dask()
+        """Custom Ray to Dask conversion implementation to pass in meta during dd.DataFrame creation."""
+
+        @dask.delayed
+        def block_to_df(block: Block):
+            block = BlockAccessor.for_block(block)
+            if isinstance(block, (ray.ObjectRef, ClientObjectRef)):
+                raise ValueError(
+                    "Dataset.to_dask() must be used with Dask-on-Ray, please "
+                    "set the Dask scheduler to ray_dask_get (located in "
+                    "ray.util.dask)."
+                )
+            return block.to_pandas()
+
+        # Use first few row from ray dataset to generate meta to infer types even if there are NaNs
+        meta = dataset.limit(100).to_pandas()
+        # meta = meta.dtypes.apply(lambda x: x.name).to_dict()
+        ddf = dd.from_delayed([block_to_df(block) for block in dataset.get_internal_block_refs()], meta=meta)
+        return ddf
 
     @property
     def array_lib(self):
diff --git a/ludwig/data/dataframe/modin.py b/ludwig/data/dataframe/modin.py
index a4cd65372d0..c26547cd919 100644
--- a/ludwig/data/dataframe/modin.py
+++ b/ludwig/data/dataframe/modin.py
@@ -59,6 +59,9 @@ def reduce_objects(self, series, reduce_fn):
     def split(self, df, probabilities):
         return split_by_slices(df.iloc, len(df), probabilities)
 
+    def remove_empty_partitions(self, df):
+        return df
+
     def to_parquet(self, df, path, index=False):
         df.to_parquet(path, engine="pyarrow", index=index)
 
diff --git a/ludwig/data/dataframe/pandas.py b/ludwig/data/dataframe/pandas.py
index 319aa508132..3a63fe115b9 100644
--- a/ludwig/data/dataframe/pandas.py
+++ b/ludwig/data/dataframe/pandas.py
@@ -61,6 +61,9 @@ def reduce_objects(self, series, reduce_fn):
     def split(self, df, probabilities):
         return split_by_slices(df.iloc, len(df), probabilities)
 
+    def remove_empty_partitions(self, df):
+        return df
+
     def to_parquet(self, df, path, index=False):
         df.to_parquet(path, engine="pyarrow", index=index)
 
diff --git a/ludwig/data/preprocessing.py b/ludwig/data/preprocessing.py
index 1b6eb8f2012..f03977d3e6d 100644
--- a/ludwig/data/preprocessing.py
+++ b/ludwig/data/preprocessing.py
@@ -1187,6 +1187,12 @@ def build_dataset(
         col_name_to_dtype[col_name] = col.dtype
     dataset = dataset.astype(col_name_to_dtype)
 
+    # Persist the completed dataset with no NaNs
+    dataset = backend.df_engine.persist(dataset)
+
+    # Remove partitions that are empty after removing NaNs
+    dataset = backend.df_engine.remove_empty_partitions(dataset)
+
     return dataset, metadata
 
 
diff --git a/ludwig/data/split.py b/ludwig/data/split.py
index 45c8f9b414d..275824424f9 100644
--- a/ludwig/data/split.py
+++ b/ludwig/data/split.py
@@ -217,4 +217,7 @@ def split_dataset(
             "Encountered an empty training set while splitting data. Please double check the preprocessing split "
             "configuration."
         )
+
+    # Remove partitions that are empty after splitting
+    datasets = [None if dataset is None else backend.df_engine.remove_empty_partitions(dataset) for dataset in datasets]
     return datasets
diff --git a/tests/integration_tests/test_preprocessing.py b/tests/integration_tests/test_preprocessing.py
index e88285801a9..cc829517891 100644
--- a/tests/integration_tests/test_preprocessing.py
+++ b/tests/integration_tests/test_preprocessing.py
@@ -124,7 +124,6 @@ def test_dask_known_divisions(feature_fn, csv_filename, tmpdir):
     input_features = [feature_fn(os.path.join(tmpdir, "generated_output"))]
     output_features = [category_feature(vocab_size=5, reduce_input="sum")]
 
-    # num_examples=100 and npartitions=2 to ensure the test is not flaky, by having non-empty post-split datasets.
     data_csv = generate_data(input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=100)
     data_df = dd.from_pandas(pd.read_csv(data_csv), npartitions=2)
     assert data_df.known_divisions
@@ -146,6 +145,38 @@ def test_dask_known_divisions(feature_fn, csv_filename, tmpdir):
         )
 
 
+@pytest.mark.distributed
+def test_drop_empty_partitions(csv_filename, tmpdir):
+    import dask.dataframe as dd
+
+    input_features = [image_feature(os.path.join(tmpdir, "generated_output"))]
+    output_features = [category_feature(vocab_size=5, reduce_input="sum")]
+
+    # num_examples and npartitions set such that each post-split DataFrame has >1 samples, but empty partitions.
+    data_csv = generate_data(input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=25)
+    data_df = dd.from_pandas(pd.read_csv(data_csv), npartitions=10)
+
+    config = {
+        "input_features": input_features,
+        "output_features": output_features,
+        "trainer": {
+            "epochs": 2,
+        },
+    }
+
+    backend = "ray"
+    with init_backend(backend):
+        model = LudwigModel(config, backend=backend)
+        train_set, val_set, test_set, _ = model.preprocess(
+            data_df,
+            skip_save_processed_input=True,
+        )
+        for dataset in [train_set, val_set, test_set]:
+            df = dataset.ds.to_dask()
+            for partition in df.partitions:
+                assert len(partition) > 0, "empty partitions found in dataset"
+
+
 @pytest.mark.parametrize("generate_images_as_numpy", [False, True])
 def test_read_image_from_path(tmpdir, csv_filename, generate_images_as_numpy):
     input_features = [image_feature(os.path.join(tmpdir, "generated_output"), save_as_numpy=generate_images_as_numpy)]


Code review comments:
diff_hunk: @@ -129,7 +149,24 @@ def to_ray_dataset(self, df):
         return from_dask(df)
 
     def from_ray_dataset(self, dataset) -> dd.DataFrame:
-        return dataset.to_dask()
+        """Custom Ray to Dask conversion implementation to pass in meta during dd.DataFrame creation."""
path: ludwig/data/dataframe/dask.py
text: @user1:
@author should we revert back to `dataset.to_dask()` if the empty partitions issue doesn't stem from a need to pass in meta?

@author:
Yup, reverted!