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.
This commit is contained in:
Srijan Keshri
2026-11-03 09:45:41 +08:00
committed by GitHub
parent 8ee7698545
commit 90b2b264d9
11 changed files with 343 additions and 21 deletions
@@ -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.
<Info>
`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.
</Info>
## The expression {#expressions}
+4 -1
View File
@@ -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}
+3
View File
@@ -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]],
+8 -1
View File
@@ -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,
+57 -10
View File
@@ -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,
@@ -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"
+18
View File
@@ -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 = []
+198
View File
@@ -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):
+8
View File
@@ -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(
{
+24 -1
View File
@@ -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<Bound<'a, PyAny>> {
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>,
@@ -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