Repository: bodo-ai/Bodo
Pull request #484: BSE-4790: Fix parallel get data scope
Author: login: IsaacWarren
Created: 2025-05-12T22:11:18Z | Merged: 2025-05-14T20:07:18Z

## Changes included in this PR
Fixes an issue where parallel gets could try to read an entry from the result registry that's already been deleted. Resolved by attaching the source lazymanager to the lazy plan. This introduced some complications in sending the plan to the workers because the managers need to be removed first so collection isn't triggered.
<!-- Include a brief description of the changes presented in this PR and any extra context that might be helpful for reviewers. -->

## Testing strategy
Unittests locally/PR CI
<!-- 
Before requesting review, verify that your changes pass PR CI by adding "[run ci]" to your commit message (or add a new blank commit with that message) or explain why CI is not necessary (e.g. docs changes). 

Briefly mention how this change is tested e.g. "new unit tests added". To pass automated coverage checks, ensure that you have added `# pragma: no cover` to jitted functions. 

Ensure that newly added tests work locally on 3 ranks using both SPMD and spawn mode (default) when applicable. For example:

SPMD mode: 
  `export BODO_SPAWN_MODE=0;
  mpiexec -n 3 pytest -svW ignore bodo/tests/test_dataframe.py::my_new_test`

Spawn mode (default mode): 
  `export BODO_NUM_WORKERS=3;
  pytest -svW ignore bodo/tests/test_dataframe.py::my_new_test`
-->

## User facing changes
Fix the title bug
<!-- Mention any changes to user facing APIs here and ensure that the documentation is up to date in Bodo/docs/docs -->

## Checklist
- [ ] Pipelines passed before requesting review. To run CI you must include `[run CI]` in your commit message.
- [x] I am familiar with the [Contributing Guide](https://github.com/bodo-ai/Bodo/blob/main/CONTRIBUTING.md) 
- [x] I have installed + ran pre-commit hooks.


Diff:
diff --git a/bodo/pandas/_plan.cpp b/bodo/pandas/_plan.cpp
index dc9aba9fab..7ad4f6007b 100644
--- a/bodo/pandas/_plan.cpp
+++ b/bodo/pandas/_plan.cpp
@@ -1,5 +1,6 @@
 #include "_plan.h"
 #include <arrow/python/pyarrow.h>
+#include <fmt/format.h>
 #include <utility>
 
 #include "_executor.h"
@@ -640,13 +641,13 @@ BodoDataFrameParallelScanFunctionData::CreatePhysicalOperator(
     // bodo.spawn.worker creates a new module with new empty registry.
 
     // Import Python sys module
-    PyObject *sys_module = PyImport_ImportModule("sys");
+    PyObjectPtr sys_module = PyImport_ImportModule("sys");
     if (!sys_module) {
         throw std::runtime_error("Failed to import sys module");
     }
 
     // Get sys.modules dictionary
-    PyObject *modules_dict = PyObject_GetAttrString(sys_module, "modules");
+    PyObjectPtr modules_dict = PyObject_GetAttrString(sys_module, "modules");
     if (!modules_dict) {
         Py_DECREF(sys_module);
         throw std::runtime_error("Failed to get sys.modules");
@@ -655,21 +656,17 @@ BodoDataFrameParallelScanFunctionData::CreatePhysicalOperator(
     // Get __main__ module
     PyObject *main_module = PyDict_GetItemString(modules_dict, "__main__");
     if (!main_module) {
-        Py_DECREF(modules_dict);
-        Py_DECREF(sys_module);
         throw std::runtime_error("Failed to get __main__ module");
     }
 
     // Get RESULT_REGISTRY[result_id]
-    PyObject *result_registry =
+    PyObjectPtr result_registry =
         PyObject_GetAttrString(main_module, "RESULT_REGISTRY");
     PyObject *df = PyDict_GetItemString(result_registry, result_id.c_str());
     if (!df) {
-        throw std::runtime_error("Result ID not found in result registry");
+        throw std::runtime_error(fmt::format(
+            "Result ID {} not found in result registry", result_id.c_str()));
     }
-    Py_DECREF(result_registry);
-    Py_DECREF(modules_dict);
-    Py_DECREF(sys_module);
 
     return std::make_shared<PhysicalReadPandas>(df, selected_columns);
 }
diff --git a/bodo/pandas/base.py b/bodo/pandas/base.py
index a28b303ca1..2f42197869 100644
--- a/bodo/pandas/base.py
+++ b/bodo/pandas/base.py
@@ -7,6 +7,7 @@
 from bodo.pandas.utils import (
     BODO_NONE_DUMMY,
     LazyPlan,
+    LazyPlanDistributedArg,
     arrow_to_empty_df,
     check_args_fallback,
     wrap_plan,
@@ -32,7 +33,11 @@ def from_pandas(df):
     res_id = None
     if bodo.dataframe_library_run_parallel:
         res_id = bodo.spawn.utils.scatter_data(df)
-        plan = LazyPlan("LogicalGetPandasReadParallel", empty_df, res_id)
+        plan = LazyPlan(
+            "LogicalGetPandasReadParallel",
+            empty_df,
+            LazyPlanDistributedArg(None, res_id),
+        )
     else:
         plan = LazyPlan("LogicalGetPandasReadSeq", empty_df, df)
 
diff --git a/bodo/pandas/frame.py b/bodo/pandas/frame.py
index b351ee0b18..0274a16caf 100644
--- a/bodo/pandas/frame.py
+++ b/bodo/pandas/frame.py
@@ -16,6 +16,7 @@
 from bodo.pandas.utils import (
     BodoLibNotImplementedException,
     LazyPlan,
+    LazyPlanDistributedArg,
     arrow_to_empty_df,
     check_args_fallback,
     get_lazy_manager_class,
@@ -64,12 +65,16 @@ def _plan(self):
                     # If the plan has been executed but the results are still
                     # distributed then re-use those results as is.
                     res_id = self._mgr._md_result_id
+                    mgr = self._mgr
                 else:
                     # The data has been collected and is no longer distributed
                     # so we need to re-distribute the results.
                     res_id = bodo.spawn.utils.scatter_data(self)
+                    mgr = None
                 self._source_plan = LazyPlan(
-                    "LogicalGetPandasReadParallel", empty_data, res_id
+                    "LogicalGetPandasReadParallel",
+                    empty_data,
+                    LazyPlanDistributedArg(mgr, res_id),
                 )
             else:
                 self._source_plan = LazyPlan(
diff --git a/bodo/pandas/series.py b/bodo/pandas/series.py
index 072c83d010..19b4f18fe1 100644
--- a/bodo/pandas/series.py
+++ b/bodo/pandas/series.py
@@ -14,6 +14,7 @@
 from bodo.pandas.utils import (
     BodoLibFallbackWarning,
     LazyPlan,
+    LazyPlanDistributedArg,
     arrow_to_empty_df,
     check_args_fallback,
     get_lazy_single_manager_class,
@@ -57,12 +58,16 @@ def _plan(self):
                         # If the plan has been executed but the results are still
                         # distributed then re-use those results as is.
                         res_id = self._mgr._md_result_id
+                        mgr = self._mgr
                     else:
                         # The data has been collected and is no longer distributed
                         # so we need to re-distribute the results.
                         res_id = bodo.spawn.utils.scatter_data(self)
+                        mgr = None
                     self._source_plan = LazyPlan(
-                        "LogicalGetPandasReadParallel", empty_data, res_id
+                        "LogicalGetPandasReadParallel",
+                        empty_data,
+                        LazyPlanDistributedArg(mgr, res_id),
                     )
                 else:
                     self._source_plan = LazyPlan(
diff --git a/bodo/pandas/utils.py b/bodo/pandas/utils.py
index cc5721debd..19206a1433 100644
--- a/bodo/pandas/utils.py
+++ b/bodo/pandas/utils.py
@@ -763,3 +763,21 @@ def arrow_to_empty_df(arrow_schema):
         {field.name: _empty_pd_array(field.type) for field in arrow_schema}
     )
     return _reconstruct_pandas_index(empty_df, arrow_schema)
+
+
+class LazyPlanDistributedArg:
+    """
+    Class to hold the arguments for a LazyPlan that are distributed on the workers.
+    """
+
+    def __init__(self, mgr, res_id: str):
+        self.mgr = mgr
+        self.res_id = res_id
+
+    def __reduce__(self):
+        """
+        This method is used to serialize the object for distribution.
+        We can't send the manager to the workers without triggering collection
+        so we just send the result ID instead.
+        """
+        return (str, (self.res_id,))
diff --git a/bodo/tests/test_df_lib/test_end_to_end.py b/bodo/tests/test_df_lib/test_end_to_end.py
index 16063e621f..ab738e5414 100644
--- a/bodo/tests/test_df_lib/test_end_to_end.py
+++ b/bodo/tests/test_df_lib/test_end_to_end.py
@@ -644,6 +644,27 @@ def test_parquet_read_partitioned_filter(datapath):
     )
 
 
+def test_parquet_read_shape_head(datapath):
+    """
+    Test to catch a case where the original manager goes out of scope
+    causing the parallel get to become invalid.
+    """
+    path = datapath("dataframe_library/df1.parquet")
+
+    def bodo_impl():
+        df = bd.read_parquet(path)
+        return df.shape, df.head(4)
+
+    def pd_impl():
+        df = pd.read_parquet(path)
+        return df.shape, df.head(4)
+
+    bdf_shape, bdf_head = bodo_impl()
+    pdf_shape, pdf_head = pd_impl()
+    assert bdf_shape == pdf_shape
+    _test_equal(bdf_head, pdf_head)
+
+
 def test_project_after_filter(datapath):
     """Test creating a plan with a Projection on top of a filter works"""
     bodo_df1 = bd.read_parquet(datapath("dataframe_library/df1.parquet"))
diff --git a/bodo/tests/test_parquet_read.py b/bodo/tests/test_parquet_read.py
index 6130eabe79..8554a3ffdd 100644
--- a/bodo/tests/test_parquet_read.py
+++ b/bodo/tests/test_parquet_read.py
@@ -1626,7 +1626,8 @@ def impl2():
 
 @pytest.mark.slow
 @pytest.mark.skipif(
-    bodo.test_dataframe_library_enabled, reason="TODO [BSE-4790]: Fix limit pushdown."
+    bodo.test_dataframe_library_enabled,
+    reason="[BSE-4764] Dictionary columns with non-string value type not supported yet.",
 )
 def test_read_pq_head_only(datapath, memory_leak_check):
     """


Reviews:
user: login: DrTodd13
body: Can instead of storing res_id in the _plan case, can you store (self._mgr, res_id)?  Then, you don't need this parameter that is [] almost everywhere?  Then, extract res_id from the tuple in the right places?
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T22:49:17Z
user: login: scott-routledge2
body: Thanks @IsaacWarren . Can we create a test for this?
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T13:46:37Z
user: login: DrTodd13
body: LGTM.  Thanks!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-14T19:12:18Z
user: login: IsaacWarren
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T19:14:28Z

Review comments:
diff_hunk: @@ -32,7 +33,11 @@ def from_pandas(df):
     res_id = None
     if bodo.dataframe_library_run_parallel:
         res_id = bodo.spawn.utils.scatter_data(df)
-        plan = LazyPlan("LogicalGetPandasReadParallel", empty_df, res_id)
+        plan = LazyPlan(
+            "LogicalGetPandasReadParallel",
+            empty_df,
+            LazyPlanDistributedArg(None, res_id),
path: bodo/pandas/base.py
user: login: DrTodd13
body: Just curious why this is "None" here.  When does distributed data created with scatter_data get collected?  Is res_id supposed to be put inside a lazy manager that is then responsible for it?  In short, how does res_id not leaking here?
created_at: 2025-05-14T19:10:59Z
author_association: COLLABORATOR
diff_hunk: @@ -32,7 +33,11 @@ def from_pandas(df):
     res_id = None
     if bodo.dataframe_library_run_parallel:
         res_id = bodo.spawn.utils.scatter_data(df)
-        plan = LazyPlan("LogicalGetPandasReadParallel", empty_df, res_id)
+        plan = LazyPlan(
+            "LogicalGetPandasReadParallel",
+            empty_df,
+            LazyPlanDistributedArg(None, res_id),
path: bodo/pandas/base.py
user: login: IsaacWarren
body: It is, that's what the ticket I've been talking about where objects are getting leaked is coming from, I'm fixing this in a followup
created_at: 2025-05-14T19:14:28Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-05-13T19:46:47Z
body: ## [Codecov](https://app.codecov.io/gh/bodo-ai/Bodo/pull/484?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bodo-ai) Report
Attention: Patch coverage is `27.27273%` with `8 lines` in your changes missing coverage. Please review.
> Project coverage is 66.06%. Comparing base [(`18406a4`)](https://app.codecov.io/gh/bodo-ai/Bodo/commit/18406a4aa2d383ddba61b276172b1e8b3f46a145?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bodo-ai) to head [(`81a243b`)](https://app.codecov.io/gh/bodo-ai/Bodo/commit/81a243b7c3e07d3415ddd4950b001a8facd5bdec?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=bodo-ai).
> Report is 4 commits behind head on main.

:x: Your project status has failed because the head coverage (66.06%) is below the target coverage (80.00%). You can increase the head coverage or adjust the [target](https://docs.codecov.com/docs/commit-status#target) coverage.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #484      +/-   ##
==========================================
+ Coverage   65.42%   66.06%   +0.63%     
==========================================
  Files         176      176              
  Lines       63691    63721      +30     
  Branches     8915     8919       +4     
==========================================
+ Hits        41672    42098     +426     
+ Misses      19397    18999     -398     
- Partials     2622     2624       +2     
```

</details>
author_association: NONE
user: login: IsaacWarren
created_at: 2025-05-13T21:53:24Z
body: > Can instead of storing res_id in the _plan case, can you store (self._mgr, res_id)? Then, you don't need this parameter that is [] almost everywhere? Then, extract res_id from the tuple in the right places?

I spent a day trying to get this to work. I can try to clean it up in my followup to fix scattered data never being deleted
author_association: CONTRIBUTOR
user: login: ehsantn
created_at: 2025-05-14T13:49:37Z
body: @IsaacWarren I'm wondering if the code can be simpler and avoid changing all other code. Why can't the manager be just attached to the parallel read node and nothing else? Let's discuss sometime.
author_association: COLLABORATOR
user: login: IsaacWarren
created_at: 2025-05-14T20:07:08Z
body: Failing tests are fixed in another PR, merging
author_association: CONTRIBUTOR