From 90b2b264d99dd9474f86f258b1336cc9a15c57cf Mon Sep 17 00:00:00 2001
From: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com>
Date: Tue, 3 Nov 2026 07:15:41 +0530
Subject: [PATCH] feat(python): accept DataFrames and readers in add_columns
Python's `add_columns` can now accept precomputed values from pandas or Polars DataFrames, Arrow Tables, RecordBatches, readers, and iterables of batches on local tables. The sync and async APIs pass an Arrow stream to the existing Rust reader transform without materializing the whole reader.
Rows attach positionally to the current live rows in storage scan order. Input must contain only new columns and exactly one row per live row; pandas indexes are ignored. Readers are consumed once, and their batch boundaries can differ from table fragments. Keep the table unchanged between computing values and adding them.
SQL mappings, Function applications, computed columns, and all-null schema inputs keep their existing behavior. Data inputs cannot be combined with computed columns, and remote tables reject reader inputs with `NotImplementedError`. Tests cover deleted rows, row counts, empty batches and tables, reader errors, collisions, reserved metadata, Arrow types, and the existing input forms. Both guides describe the supported inputs and limits.
Fixes #1959.
---
.../add-computed-column.mdx | 16 +-
docs/web/management/table.mdx | 5 +-
python/python/lancedb/_lancedb.pyi | 3 +
python/python/lancedb/remote/table.py | 9 +-
python/python/lancedb/table.py | 67 +++++-
.../tests/test_first_class_function_slice1.py | 9 +-
python/python/tests/test_remote_db.py | 18 ++
python/python/tests/test_table.py | 198 ++++++++++++++++++
python/python/tests/test_util.py | 8 +
python/src/table.rs | 25 ++-
rust/lancedb/src/table/schema_evolution.rs | 6 +
11 files changed, 343 insertions(+), 21 deletions(-)
diff --git a/docs/web/feature-engineering/add-computed-column.mdx b/docs/web/feature-engineering/add-computed-column.mdx
index 8dd8b3c23..2f5498edc 100644
--- a/docs/web/feature-engineering/add-computed-column.mdx
+++ b/docs/web/feature-engineering/add-computed-column.mdx
@@ -126,12 +126,16 @@ computes every new column from the existing rows in it, rather than growing
the table's version history with many small updates.
-`add_columns` does not accept Python callables, batch UDFs, or PyArrow
-`RecordBatch` iterators for populating new columns, whether through
-`transforms` or `computed`. New column values must be defined with SQL
-expressions, or added as NULL columns from an Arrow field or schema. If your
-transformation cannot be expressed in SQL, compute the values outside
-`add_columns` before writing them back through another workflow.
+On local tables, Python `add_columns` also accepts pandas or Polars
+DataFrames, PyArrow Tables, RecordBatches, RecordBatchReaders, and iterables
+of batches. Pass only the new columns, with exactly one row per current live
+table row in storage scan order. DataFrame indexes are ignored, and reader
+batch boundaries need not match table fragments. Keep the table unchanged
+between computing the values and adding them. Readers are consumed once.
+
+Data inputs cannot be combined with `computed` columns and are not supported
+on remote tables. Python callables and batch UDFs are not accepted; compute
+those values first, then pass the result as a DataFrame or reader.
## The expression {#expressions}
diff --git a/docs/web/management/table.mdx b/docs/web/management/table.mdx
index fc20a3f2e..fb49b3c0b 100644
--- a/docs/web/management/table.mdx
+++ b/docs/web/management/table.mdx
@@ -926,7 +926,10 @@ older version for reads, call `checkout_latest` / `checkoutLatest` before modify
You can add new columns to a table with the [`add_columns`](https://lancedb.github.io/lancedb/python/python/#lancedb.table.Table.add_columns)
method in Python, [`addColumns`](https://lancedb.github.io/lancedb/js/classes/Table/#addcolumns) in TypeScript/JavaScript, or `add_columns` in Rust.
-New columns are populated based on SQL expressions you provide.
+New columns are populated based on SQL expressions you provide. On local tables,
+the Python SDK also accepts precomputed values from DataFrames and Arrow readers;
+see [Create a computed column](/feature-engineering/add-computed-column#computed-vs-transforms)
+for the supported inputs and row-order requirements.
#### Set up the example table {#schema-set-up-the-example-table}
diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi
index d456999f4..b335e7f75 100644
--- a/python/python/lancedb/_lancedb.pyi
+++ b/python/python/lancedb/_lancedb.pyi
@@ -565,6 +565,9 @@ class Table:
) -> list[FtsToken]: ...
async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ...
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
+ async def add_columns_from_reader(
+ self, data: pa.RecordBatchReader
+ ) -> AddColumnsResult: ...
async def add_computed_columns(
self,
columns: list[tuple[str, str]],
diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py
index 3ed31c1c3..b8d4d18c3 100644
--- a/python/python/lancedb/remote/table.py
+++ b/python/python/lancedb/remote/table.py
@@ -1,6 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
+from __future__ import annotations
+
+from collections.abc import Mapping
from datetime import timedelta
import deprecation
import logging
@@ -1093,11 +1096,15 @@ class RemoteTable(Table):
def add_columns(
self,
- transforms: Dict[str, str | FunctionApplication]
+ transforms: Mapping[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
+ | DATA
+ | pa.RecordBatch
+ | pa.RecordBatchReader
+ | pl.DataFrame
| None = None,
*,
computed: Dict[str, str] | None = None,
diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py
index 192de5d0f..49b953076 100644
--- a/python/python/lancedb/table.py
+++ b/python/python/lancedb/table.py
@@ -9,6 +9,7 @@ import json
import deprecation
import warnings
from abc import ABC, abstractmethod
+from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta
from functools import cached_property
@@ -560,7 +561,7 @@ def _into_pyarrow_reader(
and data.__class__.__name__ == "LazyFrame"
):
return data.collect().to_arrow().to_reader()
- elif isinstance(data, Iterable):
+ elif isinstance(data, Iterable) and not isinstance(data, str):
return _iterator_to_reader(data, schema)
else:
raise TypeError(
@@ -2995,11 +2996,15 @@ class Table(ABC):
@abstractmethod
def add_columns(
self,
- transforms: Dict[str, str | FunctionApplication]
+ transforms: Mapping[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
+ | DATA
+ | pa.RecordBatch
+ | pa.RecordBatchReader
+ | pl.DataFrame
| None = None,
*,
computed: Dict[str, str] | None = None,
@@ -3011,8 +3016,9 @@ class Table(ABC):
Parameters
----------
- transforms: Dict[str, str | FunctionApplication], FunctionApplication,
- pa.Field, List[pa.Field], pa.Schema
+ transforms: Mapping[str, str | FunctionApplication], FunctionApplication,
+ pa.Field, List[pa.Field], pa.Schema, pandas.DataFrame,
+ polars.DataFrame, pa.Table, pa.RecordBatch, pa.RecordBatchReader
A map of column name to a SQL expression to use to calculate the
value of the new column. These expressions will be evaluated for
each row in the table, and can reference existing columns.
@@ -3026,6 +3032,16 @@ class Table(ABC):
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only in LanceDB
Enterprise.
+
+ On local tables, a pandas or Polars DataFrame, PyArrow Table,
+ RecordBatch, or RecordBatchReader supplies precomputed new-column
+ values. Supply exactly one row per current live table row, in storage
+ scan order; pandas indexes are ignored. The input contains only new
+ columns. Readers are consumed once, and their batch boundaries need
+ not match table fragments. Keep the table unchanged between computing
+ these values and calling add_columns. Data inputs cannot be combined
+ with computed columns. Remote tables do not support data inputs.
+
computed: Dict[str, str], optional
LanceDB Embedded declaration and refresh require a source build with the
``experimental-functions`` Cargo feature, which published packages
@@ -3081,6 +3097,11 @@ class Table(ABC):
x doubled
0 1 2
1 2 4
+
+ >>> import pyarrow as pa
+ >>> result = table.add_columns(pa.table({"score": [0.25, 0.75]}).to_reader())
+ >>> table.to_arrow()["score"].to_pylist()
+ [0.25, 0.75]
"""
@abstractmethod
@@ -5225,11 +5246,15 @@ class LanceTable(Table):
def add_columns(
self,
- transforms: Dict[str, str | FunctionApplication]
+ transforms: Mapping[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
+ | DATA
+ | pa.RecordBatch
+ | pa.RecordBatchReader
+ | pl.DataFrame
| None = None,
*,
computed: Dict[str, str] | None = None,
@@ -7418,11 +7443,15 @@ class AsyncTable:
async def add_columns(
self,
- transforms: dict[str, str | FunctionApplication]
+ transforms: Mapping[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
+ | DATA
+ | pa.RecordBatch
+ | pa.RecordBatchReader
+ | pl.DataFrame
| None = None,
*,
computed: dict[str, str] | None = None,
@@ -7434,7 +7463,9 @@ class AsyncTable:
Parameters
----------
- transforms: Dict[str, str | FunctionApplication] or FunctionApplication
+ transforms: Mapping[str, str | FunctionApplication], FunctionApplication,
+ pa.Field, List[pa.Field], pa.Schema, pandas.DataFrame,
+ polars.DataFrame, pa.Table, pa.RecordBatch, pa.RecordBatchReader
A map of column name to a SQL expression to use to calculate the
value of the new column. These expressions will be evaluated for
each row in the table, and can reference existing columns.
@@ -7447,6 +7478,16 @@ class AsyncTable:
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only in LanceDB
Enterprise.
+
+ On local tables, a pandas or Polars DataFrame, PyArrow Table,
+ RecordBatch, or RecordBatchReader supplies precomputed new-column
+ values. Supply exactly one row per current live table row, in storage
+ scan order; pandas indexes are ignored. The input contains only new
+ columns. Readers are consumed once, and their batch boundaries need
+ not match table fragments. Keep the table unchanged between computing
+ these values and calling add_columns. Data inputs cannot be combined
+ with computed columns. Remote tables do not support data inputs.
+
computed: Dict[str, str], optional
LanceDB Embedded declaration and refresh require a source build with the
``experimental-functions`` Cargo feature, which published packages
@@ -7515,7 +7556,7 @@ class AsyncTable:
function_output_name = None
if isinstance(transforms, FunctionApplication):
function_application = transforms
- elif isinstance(transforms, dict) and any(
+ elif isinstance(transforms, Mapping) and any(
isinstance(value, FunctionApplication) for value in transforms.values()
):
if len(transforms) != 1 or not all(
@@ -7547,7 +7588,10 @@ class AsyncTable:
):
transforms = pa.schema(transforms)
if computed:
- if transforms:
+ if transforms is not None and (
+ not isinstance(transforms, (Mapping, list, pa.Schema))
+ or bool(transforms)
+ ):
raise ValueError(
"add_columns cannot take both transforms and computed columns"
)
@@ -7564,8 +7608,11 @@ class AsyncTable:
raise ValueError("add_columns requires transforms or computed columns")
if isinstance(transforms, pa.Schema):
return await self._inner.add_columns_with_schema(transforms)
- else:
+ if isinstance(transforms, Mapping):
return await self._inner.add_columns(list(transforms.items()))
+ return await self._inner.add_columns_from_reader(
+ _into_pyarrow_reader(transforms)
+ )
async def refresh_column(
self,
diff --git a/python/python/tests/test_first_class_function_slice1.py b/python/python/tests/test_first_class_function_slice1.py
index 14c35d688..cc4485133 100644
--- a/python/python/tests/test_first_class_function_slice1.py
+++ b/python/python/tests/test_first_class_function_slice1.py
@@ -2,7 +2,9 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import json
+from collections import UserDict
from pathlib import Path
+from types import MappingProxyType
import pytest
@@ -455,13 +457,16 @@ def known_application() -> FunctionApplication:
@pytest.mark.asyncio
-async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomically():
+@pytest.mark.parametrize("mapping_type", [dict, UserDict, MappingProxyType])
+async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomically(
+ mapping_type,
+):
inner = _FunctionDeclarationInner()
table = AsyncTable(inner)
application = known_application()
result = await table.add_columns(
- {"features": application._copy(update={"columns": {}})}
+ mapping_type({"features": application._copy(update={"columns": {}})})
)
assert result == "declared"
assert inner.calls[-1][1] == "features"
diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py
index 9975e8ee1..dbf9cc4d6 100644
--- a/python/python/tests/test_remote_db.py
+++ b/python/python/tests/test_remote_db.py
@@ -5318,6 +5318,24 @@ def test_property_graph_crud_addresses_its_own_routes():
]
+def test_remote_add_columns_reader_not_supported():
+ requests = []
+
+ def handler(request):
+ requests.append(request.path)
+ request.send_response(200)
+ request.send_header("Content-Type", "application/json")
+ request.end_headers()
+ request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
+
+ with mock_lancedb_connection(handler) as db:
+ table = db.open_table("test")
+ requests.clear()
+ with pytest.raises(NotImplementedError, match="Only SQL expressions"):
+ table.add_columns(pa.table({"score": [10]}).to_reader())
+ assert all("/add_columns/" not in path for path in requests)
+
+
@pytest.mark.asyncio
async def test_async_property_graph_handles():
calls = []
diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py
index 6f2e7b5c1..50e4ac12c 100644
--- a/python/python/tests/test_table.py
+++ b/python/python/tests/test_table.py
@@ -10,10 +10,12 @@ import sys
import threading
import warnings
import weakref
+from collections import UserDict
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from decimal import Decimal, localcontext
from time import sleep
+from types import MappingProxyType
from typing import Any, List
from unittest.mock import MagicMock, patch
@@ -4988,6 +4990,69 @@ def test_restore_consistency(tmp_path):
# Schema evolution
+@pytest.mark.parametrize("mapping_type", [dict, UserDict, MappingProxyType])
+def test_add_columns_sql_mapping(mem_db, mapping_type):
+ table = mem_db.create_table("mapping_columns", pa.table({"id": [0, 1]}))
+ result = table.add_columns(mapping_type({"score": "id + 10"}))
+ assert result.version == 2
+ assert table.to_arrow().to_pydict() == {"id": [0, 1], "score": [10, 11]}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mapping_type", [dict, UserDict, MappingProxyType])
+async def test_add_columns_sql_mapping_async(mem_db_async, mapping_type):
+ table = await mem_db_async.create_table("mapping_columns", pa.table({"id": [0, 1]}))
+ result = await table.add_columns(mapping_type({"score": "id + 10"}))
+ assert result.version == 2
+ assert (await table.to_arrow()).to_pydict() == {"id": [0, 1], "score": [10, 11]}
+
+
+def test_add_columns_tuple_batches(mem_db):
+ table = mem_db.create_table("tuple_columns", pa.table({"id": [0, 1]}))
+ result = table.add_columns(
+ (pa.record_batch({"score": [10]}), pa.record_batch({"score": [20]}))
+ )
+ assert result.version == 2
+ assert table.to_arrow().to_pydict() == {"id": [0, 1], "score": [10, 20]}
+
+
+def _add_columns_input(kind):
+ values = pa.table({"score": [10, None, 30]})
+ if kind == "pandas":
+ pd = pytest.importorskip("pandas")
+ return pd.DataFrame(
+ {"score": pd.array([10, None, 30], dtype="Int64")},
+ index=pd.Index([9, 7, 5], name="source_index"),
+ )
+ if kind == "polars":
+ return pl.from_arrow(values)
+ if kind == "batch":
+ return values.to_batches()[0]
+ if kind == "reader":
+ return values.to_reader(max_chunksize=1)
+ return values
+
+
+@pytest.mark.parametrize("kind", ["pandas", "polars", "table", "batch", "reader"])
+def test_add_columns_data(mem_db, kind):
+ table = mem_db.create_table("columns_data", pa.table({"id": [2, 0, 1]}))
+ result = table.add_columns(_add_columns_input(kind))
+ assert result.version == 2
+ assert table.to_arrow().to_pydict() == {"id": [2, 0, 1], "score": [10, None, 30]}
+ assert table.schema.field("score").type == pa.int64()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", ["pandas", "polars", "table", "batch", "reader"])
+async def test_add_columns_data_async(mem_db_async, kind):
+ table = await mem_db_async.create_table("columns_data", pa.table({"id": [2, 0, 1]}))
+ result = await table.add_columns(_add_columns_input(kind))
+ assert result.version == 2
+ data = await table.to_arrow()
+ assert data.to_pydict() == {"id": [2, 0, 1], "score": [10, None, 30]}
+ assert data.schema.field("score").type == pa.int64()
+
+
def test_add_columns(mem_db: DBConnection):
data = pa.table({"id": [0, 1]})
table = LanceTable.create(mem_db, "my_table", data=data)
@@ -5944,6 +6009,139 @@ async def test_refresh_column_async_job_async_table(tmp_path, batch_size):
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
+def test_add_columns_reader_deleted_rows_and_fragments(mem_db):
+ table = mem_db.create_table("column_fragments", pa.table({"id": [9, 3]}))
+ table.add(pa.table({"id": [7, 1, 5]}))
+ table.delete("id = 3")
+ before = table.to_arrow()
+ values = pa.table({"score": [90, 70, 10, 50]})
+ reader = values.to_reader(max_chunksize=3)
+ version = table.version
+ result = table.add_columns(reader)
+ assert result.version == version + 1
+ assert table.to_arrow().to_pydict() == {
+ "id": before["id"].to_pylist(),
+ "score": [90, 70, 10, 50],
+ }
+ assert before["id"].to_pylist() == [9, 7, 1, 5]
+
+
+@pytest.mark.parametrize("count", [0, 1, 3])
+def test_add_columns_reader_wrong_count_is_atomic(mem_db, count):
+ table = mem_db.create_table("column_counts", pa.table({"id": [0, 1]}))
+ before, version = table.to_arrow(), table.version
+ reader = pa.table({"score": pa.array(range(count), type=pa.int64())}).to_reader()
+ with pytest.raises(ValueError, match="Stream"):
+ table.add_columns(reader)
+ assert table.version == version
+ assert table.to_arrow().equals(before)
+
+
+def test_add_columns_reader_empty_table(mem_db):
+ table = mem_db.create_table("empty_columns", schema=pa.schema([("id", pa.int64())]))
+ schema = pa.schema([("score", pa.float32())])
+ result = table.add_columns(pa.RecordBatchReader.from_batches(schema, []))
+ assert result.version == 2
+ assert table.schema.names == ["id", "score"]
+ assert table.schema.field("score").type == pa.float32()
+ assert table.count_rows() == 0
+
+
+def test_add_columns_reader_error_is_atomic(mem_db):
+ table = mem_db.create_table("column_error", pa.table({"id": [0, 1]}))
+ before, version = table.to_arrow(), table.version
+
+ def batches():
+ yield pa.record_batch({"score": [10]})
+ raise RuntimeError("reader failed deliberately")
+
+ reader = pa.RecordBatchReader.from_batches(
+ pa.schema([("score", pa.int64())]), batches()
+ )
+ with pytest.raises((RuntimeError, OSError), match="reader failed deliberately"):
+ table.add_columns(reader)
+ assert table.version == version
+ assert table.to_arrow().equals(before)
+
+
+def test_add_columns_reader_collision_is_atomic(mem_db):
+ table = mem_db.create_table("column_collision", pa.table({"id": [0, 1]}))
+ before, version = table.to_arrow(), table.version
+ with pytest.raises(ValueError, match="id"):
+ table.add_columns(pa.table({"id": [4, 5]}).to_reader())
+ assert table.version == version
+ assert table.to_arrow().equals(before)
+
+
+def test_add_columns_preserves_arrow_types(mem_db):
+ table = mem_db.create_table("column_types", pa.table({"id": [0, 1]}))
+ vectors = pa.array([[1.0, 2.0], [3.0, 4.0]], type=pa.list_(pa.float64(), 2))
+ values = pa.table({"vector": vectors, "score": pa.array([None, 2.5], pa.float32())})
+ table.add_columns(values)
+ assert table.schema.field("vector").type == pa.list_(pa.float64(), 2)
+ assert table.schema.field("score").type == pa.float32()
+ assert table.to_arrow().select(values.column_names).equals(values)
+
+
+def test_add_columns_dataframe_computed_conflict(mem_db):
+ table = mem_db.create_table("column_conflict", pa.table({"id": [0, 1, 2]}))
+ before, version = table.to_arrow(), table.version
+ with pytest.raises(ValueError, match="cannot take both transforms and computed"):
+ table.add_columns(_add_columns_input("pandas"), computed={"x": "id + 1"})
+ assert table.version == version
+ assert table.to_arrow().equals(before)
+
+
+@pytest.mark.experimental_functions
+@pytest.mark.parametrize(
+ "transforms", [{}, pa.schema([]), UserDict(), MappingProxyType({})]
+)
+def test_add_columns_computed_with_empty_transforms(mem_db, transforms):
+ table = mem_db.create_table("empty_transforms", pa.table({"id": [0, 1]}))
+ result = table.add_columns(transforms, computed={"x": "id + 1"})
+ assert result.version == 2
+ table.refresh_column("x")
+ assert table.to_arrow()["x"].to_pylist() == [1, 2]
+
+
+@pytest.mark.parametrize("kind", ["batch", "reader", "leading_and_trailing"])
+def test_add_columns_empty_batches(mem_db, kind):
+ ids = [0] if kind == "leading_and_trailing" else []
+ table = mem_db.create_table(
+ "empty_batches", pa.table({"id": pa.array(ids, pa.int64())})
+ )
+ schema = pa.schema([("score", pa.int64())])
+ empty = pa.RecordBatch.from_arrays([pa.array([], pa.int64())], schema=schema)
+ if kind == "batch":
+ data = empty
+ elif kind == "reader":
+ data = pa.RecordBatchReader.from_batches(schema, [empty])
+ else:
+ values = pa.RecordBatch.from_arrays([pa.array([10])], schema=schema)
+ data = pa.RecordBatchReader.from_batches(schema, [empty, values, empty])
+ result = table.add_columns(data)
+ assert result.version == 2
+ assert table.to_arrow().to_pydict() == {"id": ids, "score": [10] if ids else []}
+
+
+@pytest.mark.parametrize("kind", ["table", "reader"])
+def test_add_columns_data_rejects_computed_metadata(mem_db, kind):
+ table = mem_db.create_table("foreign_metadata", pa.table({"id": [0]}))
+ before, version = table.to_arrow(), table.version
+ schema = pa.schema(
+ [pa.field("score", pa.int64(), metadata={b"computed_column": b"true"})]
+ )
+ data = pa.Table.from_arrays([pa.array([10])], schema=schema)
+ if kind == "reader":
+ data = data.to_reader()
+ with pytest.raises(ValueError, match="computed"):
+ table.add_columns(data)
+ assert table.version == version
+ assert table.to_arrow().equals(before)
+ table.add(pa.table({"id": [1]}))
+ assert table.count_rows() == 2
+
+
@pytest.mark.experimental_functions
@pytest.mark.parametrize("interval", [None, 300])
def test_computed_auto_refresh_options(tmp_path, interval):
diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py
index 1e470e3be..0eb7ece44 100644
--- a/python/python/tests/test_util.py
+++ b/python/python/tests/test_util.py
@@ -780,6 +780,7 @@ class TestModel(lancedb.pydantic.LanceModel):
lambda: pa.table({"a": [1], "b": [2]}),
lambda: pa.table({"a": [1], "b": [2]}).to_reader(),
lambda: iter(pa.table({"a": [1], "b": [2]}).to_batches()),
+ lambda: tuple(pa.table({"a": [1], "b": [2]}).to_batches()),
lambda: lance.write_dataset(
pa.table({"a": [1], "b": [2]}),
"memory://test",
@@ -799,6 +800,7 @@ class TestModel(lancedb.pydantic.LanceModel):
"pa.Table",
"pa.RecordBatchReader",
"batch_iter",
+ "batch_tuple",
"lance.LanceDataset",
"lance.LanceScanner",
"pd.DataFrame",
@@ -813,6 +815,12 @@ def test_into_pyarrow_table(data):
assert output == expected
+@pytest.mark.parametrize("data", ["score", ""])
+def test_into_pyarrow_reader_rejects_text(data):
+ with pytest.raises(TypeError, match="Unknown data type"):
+ _into_pyarrow_reader(data)
+
+
def test_infer_target_schema():
example = pa.schema(
{
diff --git a/python/src/table.rs b/python/src/table.rs
index 1d36a7b79..38b3a7442 100644
--- a/python/src/table.rs
+++ b/python/src/table.rs
@@ -19,7 +19,7 @@ use crate::{
table::scannable::PyScannable,
};
use arrow::{
- array::{Array, LargeBinaryArray},
+ array::{Array, LargeBinaryArray, RecordBatchIterator, RecordBatchReader},
datatypes::{DataType, Schema},
ffi_stream::ArrowArrayStreamReader,
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
@@ -2060,6 +2060,29 @@ impl Table {
})
}
+ pub fn add_columns_from_reader<'a>(
+ self_: PyRef<'a, Self>,
+ data: Bound<'a, PyAny>,
+ ) -> PyResult> {
+ let reader = ArrowArrayStreamReader::from_pyarrow_bound(&data)?;
+ let schema = reader.schema();
+ // Lance's surplus-row check treats trailing empty batches as extra data.
+ let reader = RecordBatchIterator::new(
+ reader.filter(|batch| !batch.as_ref().is_ok_and(|batch| batch.num_rows() == 0)),
+ schema,
+ );
+ let inner = self_.inner_ref()?.clone();
+ future_into_py(self_.py(), async move {
+ let result = inner
+ .add_columns()
+ .transform(NewColumnTransform::Reader(Box::new(reader)))
+ .execute()
+ .await
+ .infer_error()?;
+ Ok(AddColumnsResult::from(result))
+ })
+ }
+
#[pyo3(signature = (columns, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))]
pub fn add_computed_columns(
self_: PyRef<'_, Self>,
diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs
index 8723505ff..db4d6b91f 100644
--- a/rust/lancedb/src/table/schema_evolution.rs
+++ b/rust/lancedb/src/table/schema_evolution.rs
@@ -119,6 +119,12 @@ pub(crate) async fn execute_add_columns(
NewColumnTransform::BatchUDF(udf) => {
computed_columns::ensure_no_foreign_declarations(udf.output_schema.fields())?
}
+ NewColumnTransform::Reader(reader) => {
+ computed_columns::ensure_no_foreign_declarations(reader.schema().fields())?
+ }
+ NewColumnTransform::Stream(stream) => {
+ computed_columns::ensure_no_foreign_declarations(stream.schema().fields())?
+ }
_ => {}
}
commit_add_columns(table, transforms, read_columns).await