Repository: modin-project/modin

Issue:
BUG: `virtual_partition.wait` failed if there are duplicate refs
### Modin version checks

- [X] I have checked that this issue has not already been reported.

- [X] I have confirmed this bug exists on the latest released version of Modin.

- [X] I have confirmed this bug exists on the main branch of Modin. (In order to do this you can follow [this guide](https://modin.readthedocs.io/en/stable/getting_started/installation.html#installing-from-the-github-master-branch).)


### Reproducible Example

```python
import numpy as np
import modin.pandas as pd

frame_c = pd.DataFrame(np.zeros((100, 20), dtype=np.float32, order="C"))
frame_c = [frame_c] * 20
df = pd.concat(frame_c)
partitions = df._query_compiler._modin_frame._partitions.flatten()
if hasattr(partitions[0], "wait"):
    all(map(lambda partition: partition.wait(), partitions))
print(df)
```


### Issue Description

`ray.wait` doesn't work with duplicated refs, although it is common for pandas operations to use the same objects.

### Expected Behavior

The producer must work without errors.

### Error Logs

<details>

```python-traceback

Traceback (most recent call last):
  File "test_wait.py", line 9, in <module>
    all(map(lambda partition: partition.wait(), partitions))
  File "test_wait.py", line 9, in <lambda>
    all(map(lambda partition: partition.wait(), partitions))
  File "C:\projects\modin\modin\core\execution\ray\implementations\pandas_on_ray\partitioning\virtual_partition.py", line 472, in wait
    ray.wait(futures, num_returns=len(futures))
  File "C:\Users\amyachev\Miniconda3\envs\modin\lib\site-packages\ray\_private\client_mode_hook.py", line 105, 
in wrapper
    return func(*args, **kwargs)
  File "C:\Users\amyachev\Miniconda3\envs\modin\lib\site-packages\ray\worker.py", line 1982, in wait
    raise ValueError("Wait requires a list of unique object refs.")
ValueError: Wait requires a list of unique object refs.

```

</details>


### Installed Versions

<details>

INSTALLED VERSIONS
------------------
commit           : 08a41d678490314db6f2ee95fe434fbaef90f93b
python           : 3.8.13.final.0
python-bits      : 64
OS               : Windows
OS-release       : 10
Version          : 10.0.19044
machine          : AMD64
processor        : Intel64 Family 6 Model 140 Stepping 1, GenuineIntel
byteorder        : little
LC_ALL           : None
LANG             : en_US.UTF-8
LOCALE           : English_United States.1252

Modin dependencies
------------------
modin            : 0.15.0+159.g08a41d67
ray              : 1.12.1
dask             : 2022.01.1
distributed      : 2022.01.1
hdk              : None

pandas dependencies
-------------------
pandas           : 1.4.4
numpy            : 1.23.0
pytz             : 2022.1
dateutil         : 2.8.2
setuptools       : 59.8.0
pip              : 22.1.2
Cython           : None
pytest           : 7.1.2
hypothesis       : None
sphinx           : 5.0.2
blosc            : None
feather          : 0.4.1
xlsxwriter       : None
lxml.etree       : 4.9.0
html5lib         : None
pymysql          : None
psycopg2         : 2.9.3
jinja2           : 3.1.2
IPython          : 8.4.0
pandas_datareader: None
bs4              : 4.11.1
bottleneck       : None
brotli           :
fastparquet      : None
fsspec           : 2022.5.0
gcsfs            : None
markupsafe       : 2.1.1
matplotlib       : 3.2.2
numba            : None
numexpr          : 2.7.3
odfpy            : None
openpyxl         : 3.0.9
pandas_gbq       : 0.17.6
pyarrow          : 8.0.0
pyreadstat       : None
pyxlsb           : None
s3fs             : 2022.5.0
scipy            : 1.8.1
snappy           : None
sqlalchemy       : 1.4.39
tables           : 3.7.0
tabulate         : None
xarray           : 2022.3.0
xlrd             : 2.0.1
xlwt             : None
zstandard        : None
</details>


Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/modin-project/modin/pull/5058?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project) Report
> Merging [#5058](https://codecov.io/gh/modin-project/modin/pull/5058?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project) (ce3314b) into [master](https://codecov.io/gh/modin-project/modin/commit/c89f8ba6aaa575ed44f381ad838c8e39050bc102?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project) (c89f8ba) will **decrease** coverage by `57.37%`.
> The diff coverage is `0.00%`.

> :exclamation: Current head ce3314b differs from pull request most recent head a27db11. Consider uploading reports for the commit a27db11 to get more accurate results

```diff
@@             Coverage Diff             @@
##           master    #5058       +/-   ##
===========================================
- Coverage   84.98%   27.61%   -57.38%     
===========================================
  Files         253      253               
  Lines       19124    19127        +3     
===========================================
- Hits        16252     5281    -10971     
- Misses       2872    13846    +10974     
```


| [Impacted Files](https://codecov.io/gh/modin-project/modin/pull/5058?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project) | Coverage Δ | |
|---|---|---|
| [modin/core/execution/ray/common/utils.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vY29yZS9leGVjdXRpb24vcmF5L2NvbW1vbi91dGlscy5weQ==) | `0.00% <0.00%> (-90.91%)` | :arrow_down: |
| [...mentations/pandas\_on\_ray/partitioning/partition.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vY29yZS9leGVjdXRpb24vcmF5L2ltcGxlbWVudGF0aW9ucy9wYW5kYXNfb25fcmF5L3BhcnRpdGlvbmluZy9wYXJ0aXRpb24ucHk=) | `0.00% <0.00%> (-90.76%)` | :arrow_down: |
| [...ns/pandas\_on\_ray/partitioning/partition\_manager.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vY29yZS9leGVjdXRpb24vcmF5L2ltcGxlbWVudGF0aW9ucy9wYW5kYXNfb25fcmF5L3BhcnRpdGlvbmluZy9wYXJ0aXRpb25fbWFuYWdlci5weQ==) | `0.00% <0.00%> (-67.40%)` | :arrow_down: |
| [...ns/pandas\_on\_ray/partitioning/virtual\_partition.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vY29yZS9leGVjdXRpb24vcmF5L2ltcGxlbWVudGF0aW9ucy9wYW5kYXNfb25fcmF5L3BhcnRpdGlvbmluZy92aXJ0dWFsX3BhcnRpdGlvbi5weQ==) | `0.00% <0.00%> (-91.18%)` | :arrow_down: |
| [modin/pandas/window.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vcGFuZGFzL3dpbmRvdy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [modin/distributed/\_\_init\_\_.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vZGlzdHJpYnV0ZWQvX19pbml0X18ucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [modin/experimental/core/\_\_init\_\_.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vZXhwZXJpbWVudGFsL2NvcmUvX19pbml0X18ucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [modin/core/execution/dask/\_\_init\_\_.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vY29yZS9leGVjdXRpb24vZGFzay9fX2luaXRfXy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [modin/experimental/pandas/\_\_init\_\_.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vZXhwZXJpbWVudGFsL3BhbmRhcy9fX2luaXRfXy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [modin/experimental/sklearn/\_\_init\_\_.py](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project#diff-bW9kaW4vZXhwZXJpbWVudGFsL3NrbGVhcm4vX19pbml0X18ucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| ... and [165 more](https://codecov.io/gh/modin-project/modin/pull/5058/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project) | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=modin-project)


Patch:
diff --git a/docs/release_notes/release_notes-0.16.0.rst b/docs/release_notes/release_notes-0.16.0.rst
index 978b3650..5adac7b2 100644
--- a/docs/release_notes/release_notes-0.16.0.rst
+++ b/docs/release_notes/release_notes-0.16.0.rst
@@ -57,6 +57,7 @@ Key Features and Updates
   * FIX-#4090: Fixed check if the index is trivial (#4936)
   * FIX-#5042: Fix series __getitem__ with invalid strings (#5048)
   * FIX-#4691: Fix binary operations between virtual partitions (#5049)  
+  * FIX-#5045: Fix ray virtual_partition.wait with duplicate object refs (#5058)
 * Performance enhancements
   * PERF-#4182: Add cell-wise execution for binary ops, fix bin ops for empty dataframes (#4391)
   * PERF-#4288: Improve perf of `groupby.mean` for narrow data (#4591)
diff --git a/modin/core/execution/ray/common/utils.py b/modin/core/execution/ray/common/utils.py
index acba8870..9b51e916 100644
--- a/modin/core/execution/ray/common/utils.py
+++ b/modin/core/execution/ray/common/utils.py
@@ -246,3 +246,26 @@ def deserialize(obj):
         return dict(zip(obj.keys(), RayWrapper.materialize(list(obj.values()))))
     else:
         return obj
+
+
+def wait(obj_ids):
+    """
+    Wrap ``ray.wait`` to handle duplicate object references.
+
+    ``ray.wait`` assumes a list of unique object references: see
+    https://github.com/modin-project/modin/issues/5045
+
+    Parameters
+    ----------
+    obj_ids : List[ObjectIDType]
+        The object IDs to wait on.
+
+    Returns
+    -------
+    Tuple[List[ObjectIDType], List[ObjectIDType]]
+        A list of object IDs that are ready, and a list of object IDs remaining (this
+        is the same as for ``ray.wait``). Unlike ``ray.wait``, the order of these IDs is not
+        guaranteed.
+    """
+    unique_ids = list(set(obj_ids))
+    return ray.wait(unique_ids, num_returns=len(unique_ids))
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
index 876a1059..6561fcda 100644
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
@@ -16,7 +16,7 @@
 import ray
 from ray.util import get_node_ip_address
 import uuid
-from modin.core.execution.ray.common.utils import deserialize, ObjectIDType
+from modin.core.execution.ray.common.utils import deserialize, ObjectIDType, wait
 from modin.core.execution.ray.common import RayWrapper
 
 from modin.core.dataframe.pandas.partitioning.partition import PandasDataframePartition
@@ -197,7 +197,7 @@ class PandasOnRayDataframePartition(PandasDataframePartition):
     def wait(self):
         """Wait completing computations on the object wrapped by the partition."""
         self.drain_call_queue()
-        ray.wait([self._data])
+        wait([self._data])
 
     def __copy__(self):
         """
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
index f74b15f3..91a720bc 100644
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
@@ -16,13 +16,12 @@
 import inspect
 import threading
 
-import ray
-
 from modin.config import ProgressBar
 from modin.core.execution.ray.generic.partitioning import (
     GenericRayDataframePartitionManager,
 )
 from modin.core.execution.ray.common import RayWrapper
+from modin.core.execution.ray.common.utils import wait
 from .virtual_partition import (
     PandasOnRayDataframeColumnPartition,
     PandasOnRayDataframeRowPartition,
@@ -131,7 +130,7 @@ class PandasOnRayDataframePartitionManager(GenericRayDataframePartitionManager):
         blocks = [
             block for partition in partitions for block in partition.list_of_blocks
         ]
-        ray.wait(blocks, num_returns=len(blocks))
+        wait(blocks)
 
 
 def _make_wrapped_method(name: str):
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py
index bc359c06..fc7aea7a 100644
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py
@@ -20,7 +20,7 @@ from ray.util import get_node_ip_address
 from modin.core.dataframe.pandas.partitioning.axis_partition import (
     PandasDataframeAxisPartition,
 )
-from modin.core.execution.ray.common.utils import deserialize
+from modin.core.execution.ray.common.utils import deserialize, wait
 from .partition import PandasOnRayDataframePartition
 from modin.utils import _inherit_docstrings
 
@@ -469,7 +469,7 @@ class PandasOnRayDataframeVirtualPartition(PandasDataframeAxisPartition):
         """Wait completing computations on the object wrapped by the partition."""
         self.drain_call_queue()
         futures = self.list_of_blocks
-        ray.wait(futures, num_returns=len(futures))
+        wait(futures)
 
     def add_to_apply_calls(self, func, *args, length=None, width=None, **kwargs):
         """


Test patch:
diff --git a/modin/test/storage_formats/pandas/test_internals.py b/modin/test/storage_formats/pandas/test_internals.py
index d553660b..92c91027 100644
--- a/modin/test/storage_formats/pandas/test_internals.py
+++ b/modin/test/storage_formats/pandas/test_internals.py
@@ -19,6 +19,7 @@ from modin.pandas.test.utils import (
 )
 from modin.config import NPartitions, Engine
 
+import numpy as np
 import pandas
 import pytest
 
@@ -414,3 +415,21 @@ def test_virtual_partition_apply_not_returning_pandas_dataframe(
 
     apply_result = partition.apply(lambda df: 1).get()
     assert apply_result == 1
+
+
+@pytest.mark.skipif(
+    Engine.get() != "Ray",
+    reason="Only ray.wait() does not take duplicate object refs.",
+)
+def test_virtual_partition_dup_object_ref():
+    # See https://github.com/modin-project/modin/issues/5045
+    frame_c = pd.DataFrame(np.zeros((100, 20), dtype=np.float32, order="C"))
+    frame_c = [frame_c] * 20
+    df = pd.concat(frame_c)
+    partition = df._query_compiler._modin_frame._partitions.flatten()[0]
+    obj_refs = partition.list_of_blocks
+    assert len(obj_refs) != len(
+        set(obj_refs)
+    ), "Test setup did not contain duplicate objects"
+    # The below call to wait() should not crash
+    partition.wait()