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