Skip to content

Commit 707c1b9

Browse files
kurodo3[bot]claudeeywalker
authored
fix: extension-type metadata dropped for list-backed logical types (ITL-627) (#257)
* docs(specs): add ITL-627 design spec for list-backed extension type fixes Root causes identified: - ListLogicalType.get_polars_extension_type() omits metadata= arg, causing Polars to export b'' on to_arrow() which _deserialize rejects - SemanticHashingVisitor.visit_extension short-circuits on isinstance check for generic aliases like list[File] Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs(specs): expand ITL-627 spec scope and add implementation plan - Add Defect 3: MergeJoin drops extension type when aggregating logical-type columns - Expand Defect 2 tests: set[File], list[list[T]], Dataclass with list[T] field - Add list[File] x list[File] -> list[list[File]] MergeJoin case (was wrongly out of scope) - Correct Fix 3: binary_output_schema needs no change (Schema stores Python types already correct) - Add implementation plan with complete code for all 4 tasks * fix(logical_types): pass metadata to make_polars_extension_type in ListLogicalType ListLogicalType.get_polars_extension_type() was not passing the JSON metadata bytes to make_polars_extension_type, so ext_metadata() returned None. Polars exported b'' on to_arrow(), which _deserialize rejected with ValueError during the Join/MergeJoin Polars round-trip. One-line fix covers both list[T] and set[T]. Fixes ITL-627 (Defect 1). * test(operators): add regression tests for list extension column round-trip Exercises the full Join and MergeJoin Polars round-trip with list[Path] extension columns, confirming Fix 1 (ITL-627 Defect 1) at the operator integration level. * test(operators): add data integrity assertions and clean up registration boilerplate Keep registration boilerplate in both join tests and add a comment explaining why it is required: ArrowTableStream does not trigger LogicalType registration, so without explicit pa/pl registration Polars degrades the extension type to its storage type during the operator's round-trip. Add row-count and value spot-check assertions to verify data integrity after the join. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * test(operators): fix any->all in path value assertions * fix(hashing): hash list-backed extension columns element-by-element SemanticHashingVisitor.visit_extension previously short-circuited on `not isinstance(python_type, type)` for list[File] (a GenericAlias), returning the extension type unchanged — raw JSON path strings were hashed instead of file contents. Fix: detect list-backed extension types before the isinstance guard and delegate to _visit_list_elements with a virtual large_list(elem_ext_type) so each element is hashed identically to the scalar visit_extension path. Handles set[T], list[list[T]], and arbitrary nesting depth via recursion. Adds 9 regression tests covering: list[File], set[File], list[list[File]], struct with list[File] field, passthrough when no handler, passthrough for list[Path], content-change sensitivity, same-content determinism, and element-wise hash symmetry with scalar File hashing. * refactor(hashing): move method-body imports to module level; improve docstring - Remove redundant from-body imports of File, pa, uuid, PythonTypeHandlerRegistry, SemanticAwarePythonHasher in test_extension_type_hashing.py — all were already available at module level. - Add LogicalPath, ListLogicalType, PythonTypeHandlerRegistry, SemanticAwarePythonHasher to module-level imports. - Add idempotency note to _make_list_file_ext_type / _make_scalar_file_ext_type helpers. - Add Args:/Returns: sections and passthrough-case documentation to SemanticHashingVisitor.visit_extension docstring (Google style). - Add defensive-guard comment on the `if args:` fallthrough. * fix(operators): MergeJoin produces extension-typed list when merging logical-type columns binary_static_process used pa.array(merged_vals) which inferred array type from raw storage values, producing plain large_list(storage_type) and losing the extension wrapper. Fix snapshots the element Arrow type before the Polars round-trip, then builds the merged array as pa.ExtensionArray.from_storage(ListLogicalType, ...) when the element is an extension type. Handles nested list[list[T]] naturally. Fixes ITL-627 (Defect 3). * refactor(operators): rename list_lt -> list_logical_type; add comments; improve test messages - Rename list_lt to list_logical_type in binary_static_process for clarity. - Add late-import comment explaining the circular dependency guard. - Add failure messages to shape assertions in TestMergeJoinLogicalTypeColumns. * fix(tests): use type-converter cache for list[Path] extension type in Join/MergeJoin regression tests Fresh ListLogicalType instances create different underlying Arrow extension classes. When two tests both call pa.register_extension_type() for the same extension name, the second call is a no-op — so the second test holds a class object that is never globally registered. Join's table.cast() then tries to cast between two different class objects for the same extension name, raising ArrowTypeError. Fix: use ctx.type_converter.register_python_class(list[Path]) which goes through the type converter's shared cache, guaranteeing the same class object across tests. * fix(hashing): address four review issues from brian-arnold 1. **Test contamination** (test_list_logical_type.py): Replace the fresh `ListLogicalType(LogicalPath())` + manual `pa.register_extension_type` / `pl.register_extension_type` with `ctx.type_converter.register_python_class(list[Path])`. The manual approach registered a different class object than the orcapod registry held, causing `ArrowTypeError: Casting from extension<list[orcapod.path]> to different extension type` when tests ran in `test_logical_types` before `test_core/operators` order. 2. **list[T] / set[T] hash collision** (visitors.py): The outer extension name was not encoded in the result, so `extension<list[orcapod.file]>` and `extension<set[orcapod.file]>` with identical contents produced identical hashes. Fixed by folding `extension_type.extension_name` into the combined result, mirroring the scalar path encoding. 3. **Two crash paths** (visitors.py): (a) `list[list[Path]]` — `is_nested_list_or_set` committed to recursing before checking whether the innermost type has a handler, causing `_visit_list_elements` to return `large_list(extension<...>)` which Arrow rejects in array creation. (b) Empty/null inner list — `_visit_list_elements` fell back to the element extension type when no non-null element was present, latching an extension type in `_process_table_columns`. Fixed by making the hash-vs-passthrough decision type-driven (unwrap nesting to innermost, check handler) before visiting any rows, and discarding the returned list type from `_visit_list_elements` (using `_`) since we return `(large_binary(), combined)` regardless. 4. **Wrong context in MergeJoin** (merge_join.py): `get_default_context().type_converter` bypassed the operator's actual context. Replaced with `left_stream.data_context.type_converter`, hoisted above the colliding-keys loop, and removed the late import. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: agent-kurodo[bot] <268466204+agent-kurodo[bot]@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Edgar Y. Walker <eywalker@users.noreply.github.com>
1 parent a19b814 commit 707c1b9

9 files changed

Lines changed: 1979 additions & 9 deletions

File tree

src/orcapod/core/operators/merge_join.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,15 @@ def binary_static_process(
182182
# Find colliding data columns
183183
colliding_keys = set(left_data_keys) & set(right_data_keys)
184184

185+
# Snapshot Arrow types of colliding columns BEFORE the Polars round-trip.
186+
# The round-trip may strip or alter extension metadata; we need the original
187+
# element type to reconstruct the correct list extension type after merging.
188+
colliding_col_types: dict[str, "pa.DataType"] = {
189+
col: left_table.schema.field(col).type
190+
for col in colliding_keys
191+
if col in left_table.schema.names
192+
}
193+
185194
# Capture nullable flags from input schemas BEFORE Polars conversion.
186195
# Polars' join discards nullable info (defaults all to True); we derive
187196
# the output schema from the inputs instead of from data null counts.
@@ -232,6 +241,11 @@ def binary_static_process(
232241
)
233242
joined = joined.drop(COMMON_JOIN_KEY)
234243

244+
# Use the left stream's type converter — not the default context — so that a
245+
# MergeJoin over streams built with a non-default DataContext reconstructs the
246+
# merged column's extension type from the correct registry.
247+
tc = left_stream.data_context.type_converter
248+
235249
# Process colliding data columns: merge into sorted lists
236250
for col in colliding_keys:
237251
left_col_name = col
@@ -273,7 +287,18 @@ def binary_static_process(
273287
joined = joined.drop(left_col_name)
274288
joined = joined.drop(right_col_name)
275289

276-
merged_array = pa.array(merged_vals)
290+
elem_arrow_type = colliding_col_types.get(col)
291+
if elem_arrow_type is not None and isinstance(elem_arrow_type, pa.ExtensionType):
292+
elem_python_type = tc.arrow_type_to_python_type(elem_arrow_type)
293+
list_logical_type = tc.get_logical_type_for_python_type(list[elem_python_type])
294+
if list_logical_type is not None:
295+
list_ext_type = list_logical_type.get_arrow_extension_type()
296+
storage_array = pa.array(merged_vals, type=list_ext_type.storage_type)
297+
merged_array = pa.ExtensionArray.from_storage(list_ext_type, storage_array)
298+
else:
299+
merged_array = pa.array(merged_vals)
300+
else:
301+
merged_array = pa.array(merged_vals)
277302
joined = joined.add_column(col_idx, left_col_name, merged_array)
278303

279304
if has_source:

src/orcapod/hashing/visitors.py

Lines changed: 84 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -198,21 +198,101 @@ def visit_extension(
198198
extension_type: "pa.ExtensionType",
199199
storage_value: Any,
200200
) -> tuple["pa.DataType", Any]:
201-
"""Hash an extension type value to pa.large_binary(), or passthrough."""
201+
"""Hash an extension type value to ``pa.large_binary()``, or passthrough.
202+
203+
For list-backed extension types (e.g. ``extension<list[orcapod.file]>``),
204+
delegates to ``_visit_list_elements`` with a virtual
205+
``large_list(elem_ext_type)`` so that each element is hashed identically
206+
to the scalar ``visit_extension`` path. This covers ``list[T]``,
207+
``set[T]``, and arbitrary nesting depth via recursion.
208+
209+
Three passthrough cases (extension type and storage value returned unchanged):
210+
- ``storage_value`` is ``None``.
211+
- The Python type could not be resolved (``typing.Any`` or not a plain ``type``).
212+
- The element type has no registered semantic handler and is not a nested list/set.
213+
214+
Args:
215+
extension_type: The Arrow extension type to process.
216+
storage_value: The storage-level value (result of ``to_pylist()`` on the column).
217+
218+
Returns:
219+
Tuple of ``(new_arrow_type, new_data)``. For hashable scalar types returns
220+
``(pa.large_binary(), hash_bytes)``. For list/set-backed types returns
221+
``(pa.large_list(...), [hash_bytes, ...])``. Passthroughs return the
222+
original ``(extension_type, storage_value)``.
223+
"""
202224
if storage_value is None:
203225
return extension_type, None
204226

205227
# Resolve extension type → Python type.
206228
python_type = self._type_converter.arrow_type_to_python_type(extension_type)
207229

230+
# Detect list-backed extension types: extension<list[orcapod.file]>,
231+
# extension<set[orcapod.file]>, etc. list[File] is a types.GenericAlias
232+
# (not isinstance(..., type)), so the guard below would incorrectly skip it.
233+
# We intercept here and hash each element, folding the outer extension name
234+
# into the result (mirrors the scalar path) to prevent list[T]/set[T] collisions.
235+
if (
236+
typing.get_origin(python_type) in (list, set)
237+
and pa.types.is_large_list(extension_type.storage_type)
238+
):
239+
args = typing.get_args(python_type)
240+
# Defensive guard: a well-formed list[T]/set[T] always has args, but if
241+
# not, fall through to the isinstance(python_type, type) passthrough below.
242+
if args:
243+
elem_python_type = args[0]
244+
245+
# Type-driven hashability: unwrap list/set nesting to the innermost
246+
# non-container type and check whether it has a semantic handler.
247+
# Decision is made once per column (not per row) so empty and null
248+
# inner lists are handled correctly without crashing.
249+
inner = elem_python_type
250+
while typing.get_origin(inner) in (list, set):
251+
inner_args = typing.get_args(inner)
252+
if not inner_args:
253+
break
254+
inner = inner_args[0]
255+
hashable = (
256+
isinstance(inner, type)
257+
and self._python_hasher.type_handler_registry.has_handler(inner)
258+
)
259+
if not hashable:
260+
# Innermost element type has no semantic handler — whole-column
261+
# passthrough, identical to main-branch behaviour.
262+
return extension_type, storage_value
263+
264+
# Hashable: delegate element-level hashing to _visit_list_elements.
265+
# Using the converter's element arrow type (which may itself be an
266+
# extension<list[...]>) ensures each element recurses back into
267+
# visit_extension, producing large_binary() per element.
268+
# We discard the returned list type (it may hold an extension type
269+
# when data is empty) and derive the output type from the outer name.
270+
elem_arrow_type = self._type_converter.python_type_to_arrow_type(
271+
elem_python_type
272+
)
273+
virtual_list_type = pa.large_list(elem_arrow_type)
274+
_, list_data = self._visit_list_elements(virtual_list_type, storage_value)
275+
276+
# Fold the outer extension name into the result, the same way the
277+
# scalar path does. This ensures list[T] and set[T] with identical
278+
# contents produce distinct hashes.
279+
type_name = extension_type.extension_name.replace(".", ":")
280+
combined = (
281+
type_name.encode("utf-8")
282+
+ b"::"
283+
+ b"\x00".join(
284+
elem if isinstance(elem, bytes) else b""
285+
for elem in (list_data or [])
286+
)
287+
)
288+
return pa.large_binary(), combined
289+
208290
# If the converter couldn't resolve to a concrete class, passthrough.
209291
if python_type is typing.Any or not isinstance(python_type, type):
210292
return extension_type, storage_value
211293

212294
# Only hash if a semantic hasher is registered for this Python type.
213-
if not self._python_hasher.type_handler_registry.has_handler(
214-
python_type
215-
):
295+
if not self._python_hasher.type_handler_registry.has_handler(python_type):
216296
return extension_type, storage_value
217297

218298
# Convert storage value → Python object and hash it.

src/orcapod/logical_types/list_logical_type_factory.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ def get_polars_extension_type(self) -> pl.BaseExtension:
150150
polars_ext_class = make_polars_extension_type(
151151
self._logical_type_name,
152152
self._storage_type,
153+
metadata=self._metadata_bytes.decode("utf-8"),
153154
)
154155
self._polars_ext = polars_ext_class()
155156
return self._polars_ext

0 commit comments

Comments
 (0)