fix: don't project records during broadcasting; push index down #2524
Conversation
Codecov Report
Additional details and impacted files
|
There was a problem hiding this comment.
That does it! This is the entirety of the 1-point plan (discussed on Slack).
I thought it would have to be another top-level rule beside broadcast_any_indexed, but you're right that it doesn't. Some IndexedArrays may be wrapping RecordArrays and others may not be.
Since this passes the buck to a level that might need to deal with "any indexed" again, it might seem like this could lead to infinite recursion, but it can't: you don't have infinitely many nested RecordArrays to swap them with. At some point, you get to an IndexedArray that is not containing a RecordArray, and at that point, need to project().
So this looks like a complete solution to the problem. We'll see, on a case-by-case basis, if there's any option-type data that is still getting touched because of the IndexedArray.simplified call, but it's at least a step forward, a net reduction in projecting.
The
broadcast_and_applylogic currentlyproject()sIndexedArrays during recursion. This is more expensive than necessary; caller-provided callback may exit early, avoiding the need for a projection. Instead, we can just push the indexed node into the contents of theRecordArray.For
dask-awkwardthe existing behaviour has more serious consequence; a much greater number of contents are touched during typetracer time. This PR changes the best-cast behavior of operations that encounter anIndexedArrayofRecordArrayduring broadcasting (e.g.with_fieldafter an advanced index) from touching everything in the record to touching only the offsets above the record. In practice, the presence of option-types within the record array will lead to additional touching of those buffers.