fix(python): align vector parent indices after flattening null lists (#4572)

Arrow's `flatten()` skips child values hidden under null list parents,
but `list_parent_indices()` can retain their indices. This makes
`has_nan_values()` raise `ArrowInvalid` because the parent indices and
NaN mask have different lengths.

Filter out indices belonging to null parents before applying the mask.
The regression test covers `ListArray` and `LargeListArray`,
float16/32/64, and chunked inputs: 12 cases. NaN and null-component
policies stay the same.

The crash also affects #4503. This patch fixes it separately from that
PR's changes to bad-vector policies.

Validation:

- All 12 regression cases failed before the fix. The utility suite now
passes with 87 passed and 1 skipped.
- Ruff lint and formatting pass for both changed files, as does `git
diff --check`.
- The patch applies cleanly to main at `ca275b6d`. The patched helper
from that revision also passes all 12 cases.
- A separate AI reviewer checked commit `343082a6` and found no blocking
issues. They reran the utility suite and compared 4,059 generated cases
plus 12 explicit null-parent cases with a row-by-row reference
implementation. Those checks include slices, chunked inputs, and
fixed-size lists.

SDK tests used base `07e626ec` with the existing compiled extension. The
broader Python run stopped after 474 passed and 23 skipped at
`test_local_module_outside_code_warns_at_registration` with
`ModuleNotFoundError`. That test also fails with the original vector
helper restored.

Repository-wide Ruff reports 49 lint errors and 11 files needing
formatting, all in untouched documentation examples. Full SDK
integration tests against current main require a newer native build and
have not been run.
This commit is contained in:
Srijan Keshri
2026-10-23 03:23:50 +08:00
committed by GitHub
parent 0cb5bd0a5d
commit 118e606e8e
2 changed files with 20 additions and 0 deletions
+5
View File
@@ -5333,6 +5333,11 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray
else:
values_has_nan = pc.is_nan(values)
values_indices = pc.list_parent_indices(arr)
if arr.null_count:
# Match flatten(), which excludes child spans under null parent lists.
values_indices = pc.filter(
values_indices, pc.take(pc.is_valid(arr), values_indices)
)
has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan))
indices = pa.array(range(len(arr)), type=pa.uint32())
return pc.is_in(indices, has_nan_indices)
+15
View File
@@ -21,6 +21,7 @@ from lancedb.table import (
_infer_target_schema,
_merge_metadata,
_sanitize_data,
has_nan_values,
sanitize_create_table,
)
import pyarrow as pa
@@ -489,6 +490,20 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
assert output["vector"].combine_chunks() == expected
@pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()])
@pytest.mark.parametrize("large", [False, True])
@pytest.mark.parametrize("chunked", [False, True])
def test_has_nan_values_skips_hidden_children_of_null_lists(dtype, large, chunked):
array_type = pa.LargeListArray if large else pa.ListArray
values = pa.array([float("nan"), None, float("nan"), 3.0, 4.0, 5.0], type=dtype)
arr = array_type.from_arrays(
[0, 2, 4, 6], values, mask=pa.array([True, False, False])
)
if chunked:
arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2, 1)])
assert has_nan_values(arr).to_pylist() == [False, True, False]
def test_handle_bad_vectors_noop():
# ChunkedArray should be preserved as-is
vector = pa.chunked_array(