Repository: modin-project/modin

Issue:
.query() seems to corrupt something

### System information
- **OS**: Python 3.6.9 :: Anaconda, Inc.
- **Modin installed**:  PyPi (`pip install modin[Ray]`)
- **Modin version**:  0.6.2
- **Python version **: Python 3.6.9 :: Anaconda, Inc.
- **Exact command to reproduce**:

```
import modin.pandas as pd
print(pd.__version__)
# '0.6.2'
print(pd.__pandas_version__)
# 0.25.1

d = pd.DataFrame({"x": [-1, 0, 1, None], 
                  "y": [1, 2, None, 3]})
d['z'] = d.eval('x / y')
d = d.query('z >= 0')
d = d.reset_index(drop=True)
d.columns = ['a', 'b', 'c']
print(d)
```


### Describe the problem

Printing throws an error.

### Source code / logs

Pretty sure I am running `ray` as I get the message:

```
WARNING: Not monitoring node memory since `psutil` is not installed. Install this with `pip install psutil` (or ray[debug]) to enable debugging of memory-related crashes.
```

Error msg:

```
Traceback (most recent call last):
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/ray/pandas_on_ray/frame/partition.py", line 33, in get
    return ray.get(self.oid)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/ray/worker.py", line 2247, in get
    raise value
ray.exceptions.RayTaskError: ray_worker (pid=14936, host=JAMMacMini-4.local)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/base/frame/data.py", line 287, in apply_idx_objs
    df.columns = cols
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/pandas/core/generic.py", line 5192, in __setattr__
    return object.__setattr__(self, name, value)
  File "pandas/_libs/properties.pyx", line 67, in pandas._libs.properties.AxisProperty.__set__
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/pandas/core/generic.py", line 690, in _set_axis
    self._data.set_axis(axis, labels)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/pandas/core/internals/managers.py", line 183, in set_axis
    "values have {new} elements".format(old=old_len, new=new_len)
ValueError: Length mismatch: Expected axis has 3 elements, new values have 2 elements
During handling of the above exception, another exception occurred:
ray_worker (pid=14936, host=JAMMacMini-4.local)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/ray/pandas_on_ray/frame/partition.py", line 206, in deploy_ray_func
    result = func(partition.copy(), **kwargs)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/base/frame/data.py", line 287, in apply_idx_objs
    df.columns = cols
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/pandas/core/generic.py", line 5192, in __setattr__
    return object.__setattr__(self, name, value)
  File "pandas/_libs/properties.pyx", line 67, in pandas._libs.properties.AxisProperty.__set__
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/pandas/core/generic.py", line 690, in _set_axis
    self._data.set_axis(axis, labels)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/pandas/core/internals/managers.py", line 183, in set_axis
    "values have {new} elements".format(old=old_len, new=new_len)
ValueError: Length mismatch: Expected axis has 3 elements, new values have 2 elements
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 3326, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-5-a2acebf4323b>", line 7, in <module>
    print(d)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/pandas/base.py", line 3323, in __str__
    return repr(self)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/pandas/dataframe.py", line 94, in __repr__
    result = repr(self._build_repr_df(num_rows, num_cols))
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/pandas/base.py", line 62, in _build_repr_df
    head_front = head.to_pandas()
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/backends/pandas/query_compiler.py", line 57, in to_pandas
    return self._modin_frame.to_pandas()
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/base/frame/data.py", line 1186, in to_pandas
    df = self._frame_mgr_cls.to_pandas(self._partitions)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/base/frame/partition_manager.py", line 264, in to_pandas
    retrieved_objects = [[obj.to_pandas() for obj in part] for part in partitions]
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/base/frame/partition_manager.py", line 264, in <listcomp>
    retrieved_objects = [[obj.to_pandas() for obj in part] for part in partitions]
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/base/frame/partition_manager.py", line 264, in <listcomp>
    retrieved_objects = [[obj.to_pandas() for obj in part] for part in partitions]
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/ray/pandas_on_ray/frame/partition.py", line 81, in to_pandas
    dataframe = self.get()
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/ray/pandas_on_ray/frame/partition.py", line 35, in get
    handle_ray_task_error(e)
  File "/Users/johnmount/opt/anaconda3/envs/aiAcademy/lib/python3.6/site-packages/modin/engines/ray/utils.py", line 8, in handle_ray_task_error
    raise getattr(builtins, s.split(":")[0])("".join(s.split(":")[1:]))
ValueError:  Length mismatch Expected axis has 3 elements, new values have 2 elements
```



Discussion:
Hi @JohnMount, thanks for the report!

I can reproduce this issue and confirm it is an issue with `query`, and it is related to some metadata optimizations we have internally. `query` should not be changing the partitioning, but it is. It is only happening after a new column is added, which is why it was missed by the tests. I will debug this.

Until it is fixed, there is a temporary workaround that you can add to the code to force Modin to recompute the partitioning metadata that is the source of this issue:

```python
import modin.pandas as pd 
print(pd.__version__) 
# '0.6.2' 
print(pd.__pandas_version__) 
# 0.25.1 
 
d = pd.DataFrame({"x": [-1, 0, 1, None],  
                  "y": [1, 2, None, 3]}) 
d['z'] = d.eval('x / y') 
d = d.query('z >= 0') 
# This is the line to add to reset the column partition metadata
d._query_compiler._modin_frame._column_widths_cache = None 
d = d.reset_index(drop=True) 
d.columns = ['a', 'b', 'c'] 
print(d) 
```

Recomputing this does take computation, but it should be very little and only noticeable on smaller datasets. Thanks again for posting this!
Thank you for the `modin` package and looking into his.  I am sorry the example was as long as it was, but it looked like it was some multi-stage issue (and it took me a while to cut my true example down to this).  Also thank you for the work-around.

My actual goal is to use `modin` as an optional back-end for my own [`data_algebra`](https://github.com/WinVector/data_algebra) query system.  To this end I have made the `pandas` package a variable so I can re-base my own package on `pandas` or `modin` at wll.  I figure with my odd-luck my test suite might generate a few more issues here and there.

Working with `modin` is an neat thing I will be looking into for a while, so I have no specific schedule on getting an integration to work.
Thanks @JohnMount, please do forward any additional issues. Let me know if you have any other questions about Modin!

Resolving patch:
diff --git a/modin/engines/base/frame/data.py b/modin/engines/base/frame/data.py
--- a/modin/engines/base/frame/data.py
+++ b/modin/engines/base/frame/data.py
@@ -788,7 +788,7 @@ def filter_full_axis(self, axis, func):
             A new dataframe.
         """
         new_partitions = self._frame_mgr_cls.map_axis_partitions(
-            axis, self._partitions, func
+            axis, self._partitions, func, keep_partitioning=True
         )
         if axis == 0:
             new_index = self.index
@@ -835,7 +835,10 @@ def _apply_full_axis(
             A new dataframe.
         """
         new_partitions = self._frame_mgr_cls.map_axis_partitions(
-            axis, self._partitions, self._build_mapreduce_func(axis, func)
+            axis,
+            self._partitions,
+            self._build_mapreduce_func(axis, func),
+            keep_partitioning=True,
         )
         # Index objects for new object creation. This is shorter than if..else
         if new_columns is None:
diff --git a/modin/engines/base/frame/partition_manager.py b/modin/engines/base/frame/partition_manager.py
--- a/modin/engines/base/frame/partition_manager.py
+++ b/modin/engines/base/frame/partition_manager.py
@@ -111,7 +111,7 @@ def lazy_map_partitions(cls, partitions, map_func):
         )
 
     @classmethod
-    def map_axis_partitions(cls, axis, partitions, map_func):
+    def map_axis_partitions(cls, axis, partitions, map_func, keep_partitioning=False):
         """Applies `map_func` to every partition.
 
         Note: This method should be used in the case that `map_func` relies on
@@ -127,7 +127,10 @@ def map_axis_partitions(cls, axis, partitions, map_func):
         # Since we are already splitting the DataFrame back up after an
         # operation, we will just use this time to compute the number of
         # partitions as best we can right now.
-        num_splits = cls._compute_num_partitions()
+        if keep_partitioning:
+            num_splits = len(partitions) if axis == 0 else len(partitions.T)
+        else:
+            num_splits = cls._compute_num_partitions()
         preprocessed_map_func = cls.preprocess_func(map_func)
         partitions = (
             cls.column_partitions(partitions)


Test patch:
diff --git a/modin/pandas/test/test_dataframe.py b/modin/pandas/test/test_dataframe.py
--- a/modin/pandas/test/test_dataframe.py
+++ b/modin/pandas/test/test_dataframe.py
@@ -3764,6 +3764,22 @@ def test_query(self, data, funcs):
             modin_result = modin_df.query(funcs)
             df_equals(modin_result, pandas_result)
 
+    def test_query_after_insert(self):
+        modin_df = pd.DataFrame({"x": [-1, 0, 1, None], "y": [1, 2, None, 3]})
+        modin_df["z"] = modin_df.eval("x / y")
+        modin_df = modin_df.query("z >= 0")
+        modin_result = modin_df.reset_index(drop=True)
+        modin_result.columns = ["a", "b", "c"]
+
+        pandas_df = pd.DataFrame({"x": [-1, 0, 1, None], "y": [1, 2, None, 3]})
+        pandas_df["z"] = pandas_df.eval("x / y")
+        pandas_df = pandas_df.query("z >= 0")
+        pandas_result = pandas_df.reset_index(drop=True)
+        pandas_result.columns = ["a", "b", "c"]
+
+        df_equals(modin_result, pandas_result)
+        df_equals(modin_df, pandas_df)
+
     @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
     @pytest.mark.parametrize("axis", axis_values, ids=axis_keys)
     @pytest.mark.parametrize(