tamnd / tamnd/firepanda

P5. The Arrow PyCapsule protocol, the array direction

Open
#202 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Mojo
Stars
1
Forks
0
PR merge metrics
PR metrics pending

Description

Specification: [`07-python-bindings.md`](../blob/main/docs/specs/07-python-bindings.md) section 4, [`12-the-python-front-door-measured.md`](../blob/main/docs/specs/12-the-python-front-door-measured.md) section 3.

`__arrow_c_array__` and `__arrow_c_schema__` on `DataFrame` and `Series`, and construction from anything that has them. This is the rule that keeps the boundary from becoming the bottleneck: plans and scalars cross as objects, data crosses as Arrow, once.

### This is smaller than document 07 assumes

The hard part landed at M2 and knows nothing about Python. `firepanda/io/arrow_c.mojo` declares the structs with the field order and widths pinned by tests, `arrow_export.mojo` fills them in without copying anything and installs release callbacks backed by heap boxes, and `arrow_import.mojo` goes the other way. What is missing between that and the PyCapsule protocol is a wrapper, and the probe that wrote document 12 section 3 measured it at about forty lines.

Mojo exposes what the wrapper needs, with these signatures, read out of the compiler:

```
def PyCapsule_New(pointer: Pointer[NoneType, MutUntrackedOrigin], name: StringSpan[ImmStaticOrigin], destructor: def(PyObjectPtr) abi("C") thin -> None) -> PyObjectPtr
def PyCapsule_GetPointer(capsule: PyObjectPtr, var name: String) -> Pointer[NoneType, MutUntrackedOrigin]
def PyCapsule_IsValid(capsule: PyObjectPtr, var name: String) -> Bool
```

The destructor is required, which is right, because a capsule dropped without the consumer taking it still has to release the struct it holds. `release_schema` and `release_array` are already defined to be no-ops on an already released struct, so the destructor is a `PyCapsule_GetPointer`, a release and a `free`.

The probe produced an int64 array of length five that `pyarrow.array` accepted with the right type, values and null count, and whose values buffer address on the pyarrow side was the address firepanda handed out. So the mechanism is proven and this issue is about coverage and ownership rather than about feasibility.

### Scope

- [x] `firepanda/py/convert.mojo` with the capsule wrapper and the destructors (#215)
- [x] `__arrow_c_array__` and `__arrow_c_schema__` on `DataFrame` (#215)
- [ ] `__arrow_c_array__` and `__arrow_c_schema__` on `Series`. The `PySeries` binding it was waiting on landed in #222 and the export is in #223
- [x] Construction of a `DataFrame` from any object exposing the protocol, through `arrow_import.mojo` (#217). Construction of a `Series` from one is still open and is tracked with the rest of the series construction work, since it needs the same Python to Arrow conversion `pd.Series([1, 2, 3])` needs
- [x] Every dtype firepanda has, not just int64. Strings are the interesting one, because a firepanda string column and an Arrow one are shaped differently and the export path there is doing more work (#215)
- [x] Null bearing columns, since the probe used a column with no nulls and the validity buffer is the one that is allowed to be absent (#215)
- [x] A zero copy test per dtype, asserting buffer address identity and not contents, because contents compare equal for a copy (#215)
- [x] An ownership test: take the capsules, drop the frame, read the values, and assert nothing is freed underneath. Then release and assert the frame is gone (#215)
- [x] `requested_schema` handled or explicitly refused, since the protocol allows a consumer to ask for a cast (#215, refused)

### What the export half turned out to cost

The forty line estimate above was right about the capsule wrapper and wrong about everything around it. A frame is a struct array rather than a list of arrays, so there is a parent to build and release. And a Python held frame cannot be consumed by an export and must not be copied by one, so `PyDataFrame` now holds an `ArcPointer[DataFrame]` and every exported child holds its own share. That is written up in [`15-the-arrow-capsule-boundary.md`](../blob/main/docs/specs/15-the-arrow-capsule-boundary.md).

### What the import half turned out to cost

The same thing happened again, in the other direction. `__arrow_c_array__` is almost never what a container offers: measured against pyarrow 25.0.1, Polars 1.44.1 and pandas 3.0.5, of six container types only `pyarrow.RecordBatch` has it, and a `pyarrow.Table`, a Polars frame and a pandas frame all offer only `__arrow_c_stream__`. So the import could not be built on the array direction with the stream added later, and #217 built both together. That is why this issue and #203 landed in the same two pull requests rather than one after the other.

Pointing the code at Polars also found two live bugs in the array import that shipped with #215, both of them refusals written from a careful reading of the specification that the specification does not actually say. A string view column has a floor of three buffers rather than four, and the buffer of variadic buffer lengths may be absent when there are no variadic buffers, which is every column whose strings all fit inside their views. That is a very common column and every one of them was being refused. Neither was reachable from the export side, which is the argument for testing an interchange format against the libraries that implement it.

Contributor guide

Open the contributing guide

Research direction

Start with sections 4 and 3 of docs/specs/07-python-bindings.md and docs/specs/12-the-python-front-door-measured.md, then read firepanda/py/convert.mojo and the PySeries binding from #223. Trace the existing DataFrame capsule and arrow_import paths before implementing the remaining Series protocol and construction work. Done means Series exposes both Arrow capsule methods and can be constructed from supported protocol objects without violating the documented ownership behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
pandas, python
Domain
data-engineering
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.