Skip to content

Preserve weights through pickling and deepcopy - #313

Open
vahid-ahmadi wants to merge 1 commit into
mainfrom
fix/weights-metadata
Open

Preserve weights through pickling and deepcopy#313
vahid-ahmadi wants to merge 1 commit into
mainfrom
fix/weights-metadata

Conversation

@vahid-ahmadi

Copy link
Copy Markdown

Partially addresses #300 (the pickle half — see the note at the end).

Problem

weights was a plain instance attribute, so it was not part of the pickle state:

pickle.loads(pickle.dumps(MicroSeries([1, 2], weights=[3, 4]))).sum()
# AttributeError: 'MicroSeries' object has no attribute 'weights'

That breaks to_pickle/read_pickle and anything shipping a weighted object across a process boundary — joblib, multiprocessing, Dask.

MicroDataFrame had a second, quieter failure underneath it. Its weighted aggregations are installed as per-instance closures by override_df_functions, which only runs in __init__ — a path unpickling skips. So even once weights were restored, an unpickled frame's sum was no longer overridden and fell through to pandas' unweighted implementation:

d = MicroDataFrame(pd.DataFrame({"x": [1, 2, 3]}), weights=[1, 2, 3])
pickle.loads(pickle.dumps(d)).sum()["x"]   # -> 6, expected 14 (no error)

Change

  • _metadata = ["weights"] on both classes. pandas includes _metadata attributes in the pickle state, so weights round-trip.
  • MicroDataFrame.__setstate__ reinstalls the aggregation overrides after unpickling.

Verification

Three new tests; all three fail on main (AttributeError for the Series, 6 != 14 for the frame) and pass with the change. Also verified by hand:

before after
pickle round-trip (MicroSeries) AttributeError 14.0
pickle round-trip (MicroDataFrame sum) 6 (silently unweighted) 14.0
to_pickle/read_pickle 6 14.0
copy.deepcopy 14.0 14.0
groupby(...).sum() after unpickle unweighted 14.0

Non-default indexes covered. make test — 65 passed on pandas 3.0.5 and pandas 2.3.3.

Scope

This is deliberately the low-risk half of #300. Declaring _constructor/_constructor_sliced — which is what would stop sort_values(ascending=False), fillna and sample from returning weightless plain pandas objects — is a bigger change: every pandas-internal construction would then run MicroDataFrame.__init__, including override_df_functions, so it needs a performance pass and a decision about how weights should be reindexed for row-changing operations. I've left #300 open for that, and noted there that the per-instance-closure design is what makes it awkward.

weights was a plain instance attribute, so it was absent from the pickle
state:

    pickle.loads(pickle.dumps(MicroSeries([1, 2], weights=[3, 4]))).sum()
    # AttributeError: 'MicroSeries' object has no attribute 'weights'

Declaring _metadata = ["weights"] fixes it: pandas includes _metadata
attributes in the pickle state.

MicroDataFrame needed one more step. Its weighted aggregations are
installed as per-instance closures by override_df_functions, which only
runs in __init__ — a path unpickling skips. So an unpickled frame kept
its weights but mdf.sum() fell through to the unweighted pandas
implementation and returned 6 instead of 14, with no error. __setstate__
now reinstalls them.

This covers the pickle half of #300; the _constructor rework that would
stop sort_values/fillna/sample from dropping weights is still open there.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant