Build frames from the manager instead of rebuilding them - #174
Draft
marco-mariotti wants to merge 2 commits into
Draft
Build frames from the manager instead of rebuilding them#174marco-mariotti wants to merge 2 commits into
marco-mariotti wants to merge 2 commits into
Conversation
Every pandas operation that returns a new frame rebuilds it through _constructor, and PyRanges makes that path expensive: pandas builds the frame, PyRanges.__new__ builds a second one just to look at its columns, __init__ builds a third, and RangeFrame.drop/reindex then build a fourth by wrapping the result again. A drop(columns=...) costs four DataFrame constructions and eight manager copies where one of each will do. Three changes, in the order pandas walks them: - _constructor_from_mgr builds straight from the block manager, which is what pandas does for its own frames, instead of routing through PyRanges(DataFrame(...)). The manager carries the columns in axes[0], so the required-column check that decides whether this is still a PyRanges needs no frame at all. - __new__ reads the columns off the frame pandas hands it rather than building another one, and no longer builds an empty frame it immediately discards. - RangeFrame.drop/drop_and_return/reindex return what pandas built when it is already our class, instead of rebuilding it. Metadata-only operations get about twice as fast (drop 141 -> 55 us, reindex 129 -> 44 us, head 134 -> 41 us, PyRanges(df) 26 -> 12 us; flat from 10^4 to 10^8 rows, since none of this scales with the data). Operations whose cost is data movement, such as copy and boolean masking, are unchanged. Building from a manager never calls __init__, so the loci accessor is now built on first use instead. That also fixes .loci raising AttributeError on a PyRanges that came back from pickle. Along the way, PyRanges(array, columns=[...]) silently returned a column-less DataFrame, because __new__ checked the columns of a frame built without them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
pandas 2 rebuilt a subclassed frame as the subclass; pandas 3 rebuilds it as a
plain DataFrame unless _constructor says otherwise, so rf.head() and rf[cols]
had quietly started returning DataFrames while rf.drop() still returned a
RangeFrame, because that one rebuilds its result by hand.
Move _constructor_from_mgr up to RangeFrame and have each class name the columns
it cannot do without, so both get the same rule from one implementation: keep the
class while the frame still has them, hand back a plain DataFrame when it does
not. RangeFrame now survives head, slicing, masking and copy as well.
Dropping Start or End is where the two halves of that rule meet. A RangeFrame
without them cannot even be printed, since repr raises KeyError, yet
rf.drop("Start") built one anyway; it now returns a DataFrame, as
gr.drop("Chromosome") already did.
copy() gets the same treatment as drop(): pandas hands back a frame of our class
already, and a copy keeps every column, so there is nothing to rebuild. That
matters more than it sounds, because head() in pandas is iloc[:n].copy().
At 10^6 rows, flat to 10^8:
RangeFrame.drop 71 -> 52 us
RangeFrame.reindex 61 -> 42 us
RangeFrame.head 20 -> 12 us (and a RangeFrame, not a DataFrame)
RangeFrame[cols] 86 -> 76 us (likewise)
PyRanges.head 110 -> 18 us (42 us before the copy change)
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #171 — this PR targets
fix_new, notmaster. It will be retargeted tomasteronce #171 merges. The diff here is only the two commits on top of that branch.Every pandas operation that returns a new frame rebuilds it through
_constructor, and PyRanges makes that path expensive. A singlegr.drop(columns=["X"])costs four DataFrame constructions and eight manager copies, where one of each will do:PyRanges.__new__builds a second one, purely to look at its columns;PyRanges.__init__builds a third by handing that frame toDataFrame.__init__;RangeFrame.dropbuilds a fourth, wrapping the result pandas had already rebuilt as a PyRanges.None of this scales with the data, so it is the same tax on a thousand rows as on a hundred million — it just stops being visible when there is real work to hide behind.
What changed
_constructor_from_mgron RangeFrame. This is the hook pandas calls when it has a block manager and needs a frame of the caller's class. Its default for a subclass routes throughPyRanges(DataFrame(...)); the manager already carries the column index inaxes[0], so the required-column check needs no frame at all and the result can be built straight from the manager. Both classes share one implementation and declare what they cannot do without:Keep the class while the frame still has those columns, hand back a plain DataFrame when it does not — the rule #171 established for PyRanges, now written once and applied to both.
__new__stops building a frame to read columns. pandas hands it a whole DataFrame on every rebuild, so the columns can be read off that one.drop,drop_and_return,reindexandcopyreturn what pandas built. After the above, pandas already hands back a frame of the right class, or a DataFrame when the frame lost a required column. Rebuilding either costs two more frames and buys nothing.copymatters more than it looks, because pandas implementsheadasiloc[:n].copy().The
lociaccessor is built on first use. Building from a manager never calls__init__, so_locicannot be set there alone.Benchmarks
Microseconds per call at 10^6 rows, median of three interleaved rounds, pandas 3.0.5 on Python 3.13. Verified flat at 10^7 and 10^8 —
PyRanges.headis 17.4 / 17.6 / 17.7 µs across the three sizes, because none of this is proportional to the data.PyRanges.headPyRanges.reset_indexPyRanges.reindexPyRanges.dropPyRanges(df)PyRanges.assignPyRanges.astypeRangeFrame.dropRangeFrame.reindexRangeFrame.headRangeFrame[cols]On the middle column: #171 moves the degrade-to-DataFrame decision into
_constructor_with_fallback, which builds a DataFrame to inspect its columns on every internal rebuild, costing ~22% on these operations. That is what this PR removes, by making the same decision fromaxes[0]without building anything. Merged together, the two land about twice as fast as today's master.Operations whose cost is data movement —
copy, boolean masking — are unchanged. A caveat on measuring those: in a tight loop that discards each result, the allocator penalises whichever variant allocates fewer intermediate objects, and that artifact is larger than the effect being measured; it shows in both directions depending on loop shape. Timed single-shot in fresh processes at 10^7 rows, a boolean mask is 57.3 ms on master and 56.2 ms here; at 10^8 rows every variant lands within 1% of the others (581.8 vs 581.9 ms).Behaviour changes
RangeFrame now survives
head, slicing, masking andcopy. pandas 2 rebuilt a subclassed frame as the subclass; pandas 3 rebuilds it as a plain DataFrame unless_constructorsays otherwise, sorf.head()andrf[cols]had quietly started returning DataFrames whilerf.drop()still returned a RangeFrame. They agree again.Dropping
StartorEndfrom a RangeFrame now yields a DataFrame. A RangeFrame without them cannot even be printed —reprraisesKeyError— yetrf.drop("Start")built one anyway. It now degrades, exactly asgr.drop("Chromosome")does. This is whytest_range_frame_never_degradesfrom #171 is renamed and extended here:rf.reindex(columns=["Start"])is a DataFrame now, since a frame withoutEndis not a RangeFrame.Fixed along the way
.lociraisedAttributeErroron any PyRanges that came back frompickle, because_lociwas only ever set in__init__, which unpickling skips.pr.PyRanges(array, columns=[...])silently returned a column-less DataFrame, because__new__checked the columns of a frame built without them.Risks worth knowing
_constructor_from_mgrand_from_mgrare pandas internals (both present since 2.1). The risk is asymmetric: if_constructor_from_mgrever disappears, pandas simply stops calling our override and_constructorstill handles everything; if_from_mgrdisappears, the hook raises. pandas' own_constructor_from_mgrcarries a shim for GeoDataFrame, so subclasses overriding this hook is an acknowledged pattern. A test asserting both attributes exist would turn a future pandas bump into a CI failure rather than a user-visible one — happy to add it if you want the belt and braces.__init__no longer runs when pandas rebuilds a frame. This is the maintenance trap: anything added toPyRanges.__init__later will silently not happen for rebuilt frames._lociwas exactly that bug;test_loci_survives_a_rebuildguards it.Validation is by column name only — the manager's dtypes are not inspected. Same as before, so nothing is lost, but the fast path adds no safety either.
Subclasses of PyRanges are still not preserved end to end. The hook itself now keeps them (
MyRanges(...).iloc[:1]is aMyRanges), but 36ensure_pyranges(...)call sites and_constructorhardcodepr.PyRanges. Unchanged from master; mentioned because the hook makes it look closer to working than it is.A few pandas paths still call
self._constructor(...)directly rather than through the manager hook. Of ~40 operations I instrumented, onlyastypedoes, and its overhead over plain pandas is already down from 162 µs to about 5 µs. A manager-based_constructorwould take that last bit — measured at 4.7 µs per call — but it means handling every argument shape pandas uses at those call sites. Left out deliberately.Testing
pyright 0 errors, ruff format and check clean, 136 unit tests, 85 module doctests, 10 tutorial/how-to doctests — all with the versions CI pins (pandas 3.0.5, pandas-stubs 3.0.5.260730, pyright 1.1.406, ruranges 0.2.7).
New tests cover rebuilt frames keeping their class for both PyRanges and RangeFrame, degrading when a required column is gone, and
.locisurviving both a rebuild and a pickle round-trip.