From 67c36e77b2474b36a1cf25ed6dc3d1a61022be17 Mon Sep 17 00:00:00 2001 From: Konstantin Malanchev Date: Mon, 24 Feb 2025 16:05:15 -0500 Subject: [PATCH] Fix meta derivation for from_lists --- src/nested_dask/core.py | 4 ++-- tests/nested_dask/test_nestedframe.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/nested_dask/core.py b/src/nested_dask/core.py index 9ca77de..c1659b4 100644 --- a/src/nested_dask/core.py +++ b/src/nested_dask/core.py @@ -9,7 +9,7 @@ import pyarrow as pa from dask.dataframe.dask_expr._collection import new_collection from nested_pandas.series.dtype import NestedDtype -from nested_pandas.series.packer import pack, pack_flat +from nested_pandas.series.packer import pack, pack_flat, pack_lists from pandas._libs import lib from pandas._typing import AnyAll, Axis, IndexLabel from pandas.api.extensions import no_default @@ -414,7 +414,7 @@ def from_lists(cls, df, base_columns=None, list_columns=None, name="nested"): meta = npd.NestedFrame(df[base_columns]._meta) - nested_meta = pack(df[list_columns]._meta, name) + nested_meta = pack_lists(df[list_columns]._meta, name) meta = meta.join(nested_meta) return df.map_partitions( diff --git a/tests/nested_dask/test_nestedframe.py b/tests/nested_dask/test_nestedframe.py index 3546eab..5698b89 100644 --- a/tests/nested_dask/test_nestedframe.py +++ b/tests/nested_dask/test_nestedframe.py @@ -8,6 +8,7 @@ import pytest from nested_dask.datasets import generate_data from nested_pandas.series.dtype import NestedDtype +from pandas.testing import assert_frame_equal dask.config.set({"dataframe.convert-string": False}) @@ -205,6 +206,7 @@ def test_from_lists(): ndf_comp = ndf.compute() assert list(ndf.columns) == list(ndf_comp.columns) assert list(ndf["nested"].nest.fields) == list(ndf["nested"].nest.fields) + assert_frame_equal(ndf_comp.iloc[:0], ndf.meta) # Check with just list_columns ndf = nd.NestedFrame.from_lists(nf, list_columns=["e", "f"])