mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
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:
+24
-1
@@ -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>,
|
||||
|
||||
Reference in New Issue
Block a user