Ensure make_meta doesn't hold ref to data (#9354)

jcrist · web-flow · commit 25afc340ab87 · 2022-08-04T15:42:19.000-07:00
Previously `make_meta(df)` would accidentally hold a reference to the
original backing data (which is potentially large), resulting in higher
than expected memory usage. We now perform a `copy` to ensure that the
view is replaced by a new array, removing these references.
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
@@ -65,12 +65,15 @@ def _(x):
 
 @make_meta_dispatch.register((pd.Series, pd.DataFrame))
 def _(x, index=None):
-    return x.iloc[:0]
+    out = x.iloc[:0].copy(deep=True)
+    # index isn't copied by default in pandas, even if deep=true
+    out.index = out.index.copy(deep=True)
+    return out
 
 
 @make_meta_dispatch.register(pd.Index)
 def _(x, index=None):
-    return x[0:0]
+    return x[0:0].copy(deep=True)
 
 
 meta_object_types: tuple[type, ...] = (pd.Series, pd.DataFrame, pd.Index, pd.MultiIndex)
diff --git a/dask/dataframe/tests/test_utils_dataframe.py b/dask/dataframe/tests/test_utils_dataframe.py
@@ -48,17 +48,36 @@ def test_make_meta():
     assert len(meta) == 0
     assert (meta.dtypes == df.dtypes).all()
     assert isinstance(meta.index, type(df.index))
+    # - ensure no references to original data arrays are kept
+    for col in "abc":
+        meta_pointer = meta[col].values.__array_interface__["data"][0]
+        df_pointer = df[col].values.__array_interface__["data"][0]
+        assert meta_pointer != df_pointer
+    meta_pointer = meta.index.values.__array_interface__["data"][0]
+    df_pointer = df.index.values.__array_interface__["data"][0]
+    assert meta_pointer != df_pointer
 
     # Pandas series
     meta = make_meta(df.a)
     assert len(meta) == 0
     assert meta.dtype == df.a.dtype
     assert isinstance(meta.index, type(df.index))
+    # - ensure no references to original data arrays are kept
+    meta_pointer = meta.values.__array_interface__["data"][0]
+    df_pointer = df.a.values.__array_interface__["data"][0]
+    assert meta_pointer != df_pointer
+    meta_pointer = meta.index.values.__array_interface__["data"][0]
+    df_pointer = df.index.values.__array_interface__["data"][0]
+    assert meta_pointer != df_pointer
 
     # Pandas index
     meta = make_meta(df.index)
     assert isinstance(meta, type(df.index))
     assert len(meta) == 0
+    # - ensure no references to original data arrays are kept
+    meta_pointer = meta.values.__array_interface__["data"][0]
+    df_pointer = df.index.values.__array_interface__["data"][0]
+    assert meta_pointer != df_pointer
 
     # Dask object
     ddf = dd.from_pandas(df, npartitions=2)
