mirror of
https://github.com/lancedb/lancedb.git
synced 2026-07-10 22:40:41 +00:00
fix(python): flatten_columns raises when flatten=False (#3629)
### Summary
`flatten_columns` raises `ValueError` when called with `flatten=False`,
even though `False` should mean "do not flatten". This is reachable from
the public API — `Query.to_pandas(flatten=...)` and
`to_batches(flatten=...)` type their `flatten` param as
`Optional[Union[int, bool]]` and pass it straight to `flatten_columns`.
### Cause
`bool` is a subclass of `int`, so `isinstance(False, int)` is `True`.
`flatten=False` skips the `flatten is True` check, falls into the
integer branch, and `False <= 0` evaluates to `True`, raising:
```
ValueError: Please specify a positive integer for flatten or the boolean value `True`
```
### Reproduction
```python
import lancedb
db = lancedb.connect("/tmp/db")
t = db.create_table("t", data=[{"id": 1, "vector": [0.1, 0.2]}])
t.search([0.1, 0.2]).to_pandas(flatten=False) # -> ValueError
```
### Fix
Guard the integer branch with `not isinstance(flatten, bool)` so that
`flatten=False` (and `None`) mean "do not flatten". Behavior is
otherwise unchanged:
- `flatten=True` → flatten all nested levels
- positive `int` → flatten to that depth
- non-positive `int` (e.g. `0`) → still rejected with `ValueError`
Added a regression test in `tests/test_util.py` covering `None`,
`False`, `True`, a positive depth, and `0`.
This commit is contained in:
@@ -177,7 +177,10 @@ def flatten_columns(tbl: pa.Table, flatten: Optional[Union[int, bool]] = None):
|
||||
continue
|
||||
else:
|
||||
break
|
||||
elif isinstance(flatten, int):
|
||||
# `bool` is a subclass of `int`, so guard against it explicitly: `flatten=False`
|
||||
# (and `None`) must mean "do not flatten" rather than falling into the integer
|
||||
# branch and raising on the `flatten <= 0` check.
|
||||
elif isinstance(flatten, int) and not isinstance(flatten, bool):
|
||||
if flatten <= 0:
|
||||
raise ValueError(
|
||||
"Please specify a positive integer for flatten or the boolean "
|
||||
|
||||
@@ -25,10 +25,42 @@ import pandas as pd
|
||||
import polars as pl
|
||||
import pytest
|
||||
import lancedb
|
||||
from lancedb.util import get_uri_scheme, join_uri, value_to_sql
|
||||
from lancedb.util import flatten_columns, get_uri_scheme, join_uri, value_to_sql
|
||||
from utils import exception_output
|
||||
|
||||
|
||||
def _struct_table() -> pa.Table:
|
||||
return pa.table(
|
||||
{
|
||||
"id": [1, 2],
|
||||
"nested": pa.array([{"a": 1, "b": 2}, {"a": 3, "b": 4}]),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_flatten_columns():
|
||||
tbl = _struct_table()
|
||||
|
||||
# None / False mean "do not flatten": the struct column is preserved.
|
||||
# `False` is a regression guard: because bool is a subclass of int it used
|
||||
# to fall into the integer branch and raise ValueError (see issue).
|
||||
for no_flatten in (None, False):
|
||||
result = flatten_columns(tbl, no_flatten)
|
||||
assert result.column_names == ["id", "nested"]
|
||||
|
||||
# True flattens all nested levels.
|
||||
flattened = flatten_columns(tbl, True)
|
||||
assert flattened.column_names == ["id", "nested.a", "nested.b"]
|
||||
|
||||
# A positive integer flattens up to that depth.
|
||||
flattened = flatten_columns(tbl, 1)
|
||||
assert flattened.column_names == ["id", "nested.a", "nested.b"]
|
||||
|
||||
# Non-positive integers are still rejected.
|
||||
with pytest.raises(ValueError):
|
||||
flatten_columns(tbl, 0)
|
||||
|
||||
|
||||
def test_normalize_uri():
|
||||
uris = [
|
||||
"relative/path",
|
||||
|
||||
Reference in New Issue
Block a user