fix(table): rewrap extension arrays instead of casting when projecting - #1599
fix(table): rewrap extension arrays instead of casting when projecting#1599twuebi wants to merge 2 commits into
Conversation
Arrow has no extension-to-extension cast kernel. Projecting a column whose file and requested types are the same named extension with differing parameters fails with "cast_extension has no kernel matching input types", e.g. geoarrow.wkb spelling the same CRS as an SRID vs. an authority code. The values already carry the extension's storage encoding, so rewrapping that storage in the requested type is zero-copy and correct; storage is cast only on a width mismatch, and the requested type stays authoritative.
There was a problem hiding this comment.
This is a clean fix and I think it's right: the extension-to-extension cast gap is real and rewrapping is the correct response. I traced the width-mismatch branch to be sure, and it's not a crash risk: Arrow registers CastFromExtension for every EXTENSION→target cast, so casting the wrapped array unwraps and casts the storage for us. No objection to merging.
A couple of small things I'd tidy while we're here, none blocking:
The width-mismatch branch casts the whole extension array. It works, but I'd cast vals.(array.ExtensionArray).Storage() directly instead — same result, explicit intent, no reliance on the extension-cast registration. And I'd add a one-line comment to rewrapExtension naming the invariant it leans on (shared extension name implies the storage payload is interpretation-compatible), so nobody extends it to a type where that isn't true.
In the negative test, cast_extension is Arrow's internal kernel name rather than a stable surface — I'd loosen that to require.Error with a comment. Details are inline.
One aside, not for this PR: PyIceberg's _cast_if_needed still hits this same failure on the same SRID-vs-authority-code input, so this diverges from (broken) PyIceberg behavior — worth a line in the description as a follow-up.
Nice fix overall. Happy to merge once those are in — or if you'd rather land as-is and follow up, I won't hold it up.
| } | ||
|
|
||
| srcExt, ok := vals.DataType().(arrow.ExtensionType) | ||
| if !ok || srcExt.ExtensionName() != tgtExt.ExtensionName() { |
There was a problem hiding this comment.
The rewrap fires for any two extension arrays that share a name, which is exactly right for geoarrow.wkb — CRS is metadata, the WKB bytes don't change. But the helper reads as general-purpose, and a future extension whose parameters actually affect the storage encoding would get silently rewrapped here instead of erroring.
I'd add a line to the doc comment naming the invariant we're leaning on — something like "safe only because a shared extension name implies the storage payload is interpretation-compatible." That way the next person extending this knows the precondition before they add a type where it doesn't hold.
| vals.(array.ExtensionArray).Storage()), true | ||
| } | ||
|
|
||
| storage := retOrPanic(compute.CastArray(a.ctx, vals, |
There was a problem hiding this comment.
This casts the whole extension array to the target storage type, and it does work — Arrow's addCastFuncs registers CastFromExtension for every EXTENSION→target cast, so the kernel unwraps the extension and casts the storage for us. But it only works because of that registration, and it reads as "cast a WKB array to large_binary" when what we actually want is to widen the storage.
Since we've already confirmed vals is an extension array, I'd cast the storage directly:
storage := retOrPanic(compute.CastArray(a.ctx,
vals.(array.ExtensionArray).Storage(),
compute.SafeCastOptions(tgtExt.StorageType())))Same result, but the intent is explicit and we're not relying on the extension-cast path (it also lines up with the .Storage() call the zero-copy branch already makes). wdyt?
| { | ||
| name: "geometry_metadata_only", | ||
| icebergType: iceberg.GeometryType{}, | ||
| sourceStorage: arrow.BinaryTypes.Binary, |
There was a problem hiding this comment.
These geometry cases leave sourceEdges at its zero value, which geoArrowMetadataToIcebergType happens to treat as planar. It works, but an explicit sourceEdges: geoarrow.EdgePlanar would make the test say what it means rather than lean on the default. Minor — take it or leave it.
| ctx := compute.WithAllocator(context.Background(), mem) | ||
| _, err := table.ToRequestedSchema(ctx, sc, sc, rec, table.SchemaOptions{IncludeFieldIDs: true}) | ||
| require.Error(t, err) | ||
| assert.ErrorContains(t, err, "cast_extension") |
There was a problem hiding this comment.
cast_extension is Arrow's internal kernel name — it shows up in the error text but it isn't a stable API surface, so an Arrow error-message refactor would break this assertion for reasons unrelated to what we're testing. Since the thing we actually care about is "mismatched extension names don't get rewrapped and the cast fails", I'd drop to require.Error(t, err) with a one-line comment saying why. If we want something more specific, better to wrap the failure in an iceberg-go error and match on that. wdyt?
Cast the source storage array directly on a width mismatch instead of relying on Arrow's extension-input cast registration, document the invariant the rewrap depends on (a shared extension name implies the storage payload is interpretation-compatible), make the geometry test cases state planar edges explicitly, and stop matching Arrow's internal kernel name in the negative test.
2 participants
Arrow has no extension-to-extension cast kernel. Projecting a column whose file and requested types are the same named extension with differing parameters fails with "cast_extension has no kernel matching input types", e.g. geoarrow.wkb spelling the same CRS as an SRID vs. an authority code.
The values already carry the extension's storage encoding, so rewrapping that storage in the requested type is zero-copy and correct; storage is cast only on a width mismatch, and the requested type stays authoritative.