From ff81428a9cdf7b06ce6d1d25f48f2ecedc862718 Mon Sep 17 00:00:00 2001
From: Pranav Achar <73564284+PranavAchar01@users.noreply.github.com>
Date: Thu, 9 Jul 2026 11:07:26 -0700
Subject: [PATCH] fix(python): flatten_columns raises when flatten=False
 (#3629)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

### 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`.
---
 python/python/lancedb/util.py    |  5 ++++-
 python/python/tests/test_util.py | 34 +++++++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py
index 932a3e982..6a5be99db 100644
--- a/python/python/lancedb/util.py
+++ b/python/python/lancedb/util.py
@@ -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 "
diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py
index c96407779..da33cc77f 100644
--- a/python/python/tests/test_util.py
+++ b/python/python/tests/test_util.py
@@ -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",