rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-09-04 12:38:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support blob computed column refresh

Browse Source
This commit is contained in:
Xuanwo
2026-08-28 14:32:29 +08:00
parent 6ab3b9eb30
commit c93911033f
11 changed files with 892 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/_lancedb.pyi
+3 -1
View File
@@ -352,7 +352,9 @@ class Table:
async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ...
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
async def add_computed_columns(
self, columns: list[tuple[str, str]]
self,
columns: list[tuple[str, str]],
blob_columns: Optional[list[tuple[str, str]]] = None,
) -> AddColumnsResult: ...
async def add_function_columns(
self, application_json: str, output_name: Optional[str]
python/python/lancedb/remote/table.py
+6 -1
View File
@@ -976,8 +976,13 @@ class RemoteTable(Table):
| None = None,
*,
computed: Dict[str, str] | None = None,
computed_blobs: Dict[str, str] | None = None,
) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms, computed=computed))
return LOOP.run(
self._table.add_columns(
transforms, computed=computed, computed_blobs=computed_blobs
)
)
def refresh_column(self, column: str):
return LOOP.run(self._table.refresh_column(column))
python/python/lancedb/table.py
+25 -4
View File
@@ -2143,6 +2143,7 @@ class Table(ABC):
| None = None,
*,
computed: Dict[str, str] | None = None,
computed_blobs: Dict[str, str] | None = None,
):
"""
Add new columns with defined values.
@@ -2184,6 +2185,12 @@ class Table(ABC):
server, and the refresh runs as a server job -- see
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
Cannot be combined with ``transforms``.
computed_blobs: Dict[str, str], optional
A map of Blob v2 output column names to SQL expressions returning
``LargeBinary`` payload bytes. Blob inputs named by an expression
are materialized as bytes, and refresh stores the result as Blob
v2 so ``blob_columns()`` and Blob read APIs continue to recognize
it. Cannot be combined with ``transforms``.
Returns
-------
@@ -4300,8 +4307,13 @@ class LanceTable(Table):
| None = None,
*,
computed: Dict[str, str] | None = None,
computed_blobs: Dict[str, str] | None = None,
) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms, computed=computed))
return LOOP.run(
self._table.add_columns(
transforms, computed=computed, computed_blobs=computed_blobs
)
)
def refresh_column(self, column: str) -> "RefreshColumnResult":
"""Fill a computed column's unfilled rows. See
@@ -6248,6 +6260,7 @@ class AsyncTable:
| None = None,
*,
computed: dict[str, str] | None = None,
computed_blobs: dict[str, str] | None = None,
) -> AddColumnsResult:
"""
Add new columns with defined values.
@@ -6283,6 +6296,11 @@ class AsyncTable:
On LanceDB Cloud and Enterprise the expression is planned by
the server. Cannot be combined with ``transforms``.
computed_blobs: Dict[str, str], optional
A map of Blob v2 output column names to SQL expressions returning
``LargeBinary`` payload bytes. Blob inputs are materialized as
payload bytes during refresh. Cannot be combined with
``transforms``.
Returns
-------
@@ -6306,7 +6324,7 @@ class AsyncTable:
function_output_name, function_application = next(iter(transforms.items()))
if function_application is not None:
if computed:
if computed or computed_blobs:
raise ValueError(
"add_columns cannot mix a Function application with SQL "
"computed columns"
@@ -6322,12 +6340,15 @@ class AsyncTable:
{isinstance(f, pa.Field) for f in transforms}
):
transforms = pa.schema(transforms)
if computed:
if computed or computed_blobs:
if transforms:
raise ValueError(
"add_columns cannot take both transforms and computed columns"
)
return await self._inner.add_computed_columns(list(computed.items()))
return await self._inner.add_computed_columns(
list((computed or {}).items()),
list((computed_blobs or {}).items()),
)
if transforms is None:
raise ValueError("add_columns requires transforms or computed columns")
if isinstance(transforms, pa.Schema):
python/python/tests/test_table.py
+36
View File
@@ -4087,6 +4087,42 @@ def test_computed_column_rejects_transforms_and_computed_together(tmp_path):
table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"})
def test_computed_blob_input_and_explicit_output(tmp_path):
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
db = lancedb.connect(tmp_path)
table = db.create_table("computed_blob", schema=schema)
table.add(
[
{"id": 1, "image": b"hello"},
{"id": 2, "image": b""},
{"id": 3, "image": None},
]
)
table.add_columns(
computed={"payload_copy": "image"},
computed_blobs={"image_copy": "image"},
)
assert table.refresh_column("payload_copy").rows_filled == 2
assert table.refresh_column("image_copy").rows_filled == 2
values = table.to_arrow()["payload_copy"].combine_chunks().to_pylist()
assert values == [b"hello", b"", None]
assert table.blob_columns() == ["image", "image_copy"]
hits = table.search().with_row_id(True).limit(10).to_arrow()
rows = sorted(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
copied = table.fetch_blobs("image_copy", [row_id for _, row_id in rows])
assert copied.to_pylist() == [b"hello", b"", None]
def test_computed_blob_rejects_eager_transforms(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("computed_blob_mixed", [{"x": 1}])
with pytest.raises(ValueError):
table.add_columns({"a": "x + 1"}, computed_blobs={"b": "x"})
@pytest.mark.asyncio
async def test_computed_column_async(tmp_path):
db = await lancedb.connect_async(tmp_path)
python/src/table.rs
+5
View File
@@ -1575,9 +1575,11 @@ impl Table {
})
}
#[pyo3(signature = (columns, blob_columns=None))]
pub fn add_computed_columns(
self_: PyRef<'_, Self>,
columns: Vec<(String, String)>,
blob_columns: Option<Vec<(String, String)>>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
@@ -1585,6 +1587,9 @@ impl Table {
for (name, expression) in columns {
builder = builder.computed(name, expression);
}
for (name, expression) in blob_columns.unwrap_or_default() {
builder = builder.computed_blob(name, expression);
}
let result = builder.execute().await.infer_error()?;
Ok(AddColumnsResult::from(result))
})
rust/lancedb/src/remote/table.rs
+46 -8
View File
@@ -3174,7 +3174,10 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
}
}
async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result<AddColumnsResult> {
async fn add_computed_columns(
&self,
columns: &[crate::table::computed_columns::ComputedColumnDeclaration],
) -> Result<AddColumnsResult> {
self.check_mutable().await?;
crate::table::computed_columns::ensure_no_function_bindings_for_mutation(
self.schema().await?.as_ref(),
@@ -3184,13 +3187,16 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
// inference and the persisted binding all happen there.
let entries = columns
.iter()
.map(
|(name, expression)| lance_namespace::models::AddColumnsEntry {
name: name.clone(),
computed: Some(Some(expression.clone())),
..Default::default()
},
)
.map(|column| {
let mut entry = serde_json::json!({
"name": column.name,
"computed": column.expression,
});
if column.output == crate::table::computed_columns::ComputedColumnOutput::BlobV2 {
entry["computed_output"] = serde_json::json!("blob_v2");
}
entry
})
.collect::<Vec<_>>();
let mut body = serde_json::json!({ "new_columns": entries });
self.apply_branch_body(&mut body);
@@ -7419,6 +7425,38 @@ mod tests {
assert_eq!(result.version, 7);
}
#[tokio::test]
async fn test_add_blob_computed_column_sends_explicit_output_semantics() {
let table = Table::new_with_handler("my_table", |request| match request.url().path() {
"/v1/table/my_table/describe/" => simple_describe_response(),
"/v1/table/my_table/add_columns/" => {
let body = request.body().unwrap().as_bytes().unwrap();
let value: serde_json::Value = serde_json::from_slice(body).unwrap();
assert_eq!(
value["new_columns"],
serde_json::json!([{
"name": "image_copy",
"computed": "image",
"computed_output": "blob_v2"
}])
);
http::Response::builder()
.status(200)
.body(r#"{"version": 8}"#.to_string())
.unwrap()
}
path => panic!("Unexpected path: {path}"),
});
let result = table
.add_columns()
.computed_blob("image_copy", "image")
.execute()
.await
.unwrap();
assert_eq!(result.version, 8);
}
#[tokio::test]
async fn test_add_scalar_function_column_sends_atomic_null_declaration() {
let table = Table::new_with_handler("my_table", |request| {
rust/lancedb/src/table.rs
+5 -2
View File
@@ -754,7 +754,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
/// for the server to plan.
async fn add_computed_columns(
&self,
_columns: &[(String, String)],
_columns: &[computed_columns::ComputedColumnDeclaration],
) -> Result<AddColumnsResult> {
Err(Error::NotSupported {
message: "computed columns are not supported on this table type".into(),
@@ -3523,7 +3523,10 @@ impl BaseTable for NativeTable {
Ok(result)
}
async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result<AddColumnsResult> {
async fn add_computed_columns(
&self,
columns: &[computed_columns::ComputedColumnDeclaration],
) -> Result<AddColumnsResult> {
let result = schema_evolution::execute_declare(self, columns).await?;
self.bump_freshness();
Ok(result)
rust/lancedb/src/table/add_columns.rs
+30 -2
View File
@@ -8,6 +8,7 @@ use std::sync::Arc;
use lance::dataset::NewColumnTransform;
use super::BaseTable;
use super::computed_columns::ComputedColumnDeclaration;
use super::schema_evolution::AddColumnsResult;
use crate::function::FunctionApplication;
use crate::{Error, Result};
@@ -16,7 +17,7 @@ use crate::{Error, Result};
pub struct AddColumnsBuilder {
parent: Arc<dyn BaseTable>,
transform: Option<NewColumnTransform>,
computed: Vec<(String, String)>,
computed: Vec<ComputedColumnDeclaration>,
function: Option<(FunctionApplication, Option<String>)>,
read_columns: Option<Vec<String>>,
}
@@ -83,7 +84,34 @@ impl AddColumnsBuilder {
/// # }
/// ```
pub fn computed(mut self, name: impl Into<String>, expression: impl Into<String>) -> Self {
self.computed.push((name.into(), expression.into()));
self.computed
.push(ComputedColumnDeclaration::inferred(name, expression));
self
}
/// Add a Blob v2 column defined by a `LargeBinary` expression and filled
/// by a later refresh.
///
/// Blob inputs in the expression are materialized as their payload bytes.
/// The expression result is wrapped back into the Blob v2 logical type
/// before publication, so queries and [`Table::blob_columns`](super::Table::blob_columns)
/// continue to recognize the output as a Blob column.
///
/// ```
/// # use lancedb::Table;
/// # async fn declare(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
/// table
/// .add_columns()
/// .computed_blob("image_copy", "image")
/// .execute()
/// .await?;
/// table.refresh_column("image_copy").await?;
/// # Ok(())
/// # }
/// ```
pub fn computed_blob(mut self, name: impl Into<String>, expression: impl Into<String>) -> Self {
self.computed
.push(ComputedColumnDeclaration::blob(name, expression));
self
}
rust/lancedb/src/table/computed_columns.rs
+225 -19
View File
@@ -19,13 +19,15 @@
//! [`computed_columns`] and [`computed_column_from_field`] read declarations
//! back off a schema.
use std::collections::{BTreeSet, HashMap};
use std::collections::{BTreeSet, HashMap, HashSet};
use std::sync::Arc;
use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef};
use datafusion_common::tree_node::TreeNode;
use datafusion_physical_plan::PhysicalExpr;
use lance::dataset::NewColumnTransform;
use lance_arrow::FieldExt;
use lance_core::datatypes::{BLOB_V2_DESC_FIELD, format_field_path_minimal, parse_field_path};
use lance_datafusion::planner::Planner;
use lance_namespace::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema};
use serde::{Deserialize, Serialize};
@@ -64,6 +66,47 @@ pub const SQL_KIND: &str = "sql";
/// Value of [`KIND_META_KEY`] for a registered Function binding.
pub const FUNCTION_KIND: &str = "function";
/// How a SQL computed declaration chooses its stored output type.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum ComputedColumnOutput {
/// Use the expression's inferred Arrow type.
Inferred,
/// Store a materialized `LargeBinary` result as a Blob v2 column.
BlobV2,
}
/// One SQL computed-column declaration before it is planned.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ComputedColumnDeclaration {
/// Name of the column to declare.
pub name: String,
/// Immutable SQL expression evaluated by refresh.
pub expression: String,
/// Stored semantic type of the result.
pub output: ComputedColumnOutput,
}
impl ComputedColumnDeclaration {
/// Declare a computed column whose type is inferred from its expression.
pub fn inferred(name: impl Into<String>, expression: impl Into<String>) -> Self {
Self {
name: name.into(),
expression: expression.into(),
output: ComputedColumnOutput::Inferred,
}
}
/// Declare a Blob v2 output backed by a `LargeBinary` expression.
pub fn blob(name: impl Into<String>, expression: impl Into<String>) -> Self {
Self {
name: name.into(),
expression: expression.into(),
output: ComputedColumnOutput::BlobV2,
}
}
}
/// Synthetic result identity used when the entire Function result maps to one
/// table column (scalar or struct-as-one-column).
pub const WHOLE_RESULT_FIELD: &str = "$value";
@@ -1162,15 +1205,124 @@ pub(crate) struct BoundExpression {
/// The columns the expression names, as written; nested inputs keep
/// their dotted path.
pub inputs: Vec<String>,
/// The top-level columns evaluation reads, in [`Self::read_schema`]
/// order. A nested input appears through its root.
/// The top-level columns evaluation reads, in physical-expression order.
/// A nested input appears through its root.
pub roots: Vec<String>,
/// The projected schema evaluation runs against.
pub read_schema: SchemaRef,
/// The compiled expression.
pub physical: Arc<dyn PhysicalExpr>,
/// The type the expression yields.
pub data_type: DataType,
/// Blob v2 leaves the scan must materialize as `LargeBinary`.
pub blob_paths: Vec<String>,
}
fn collect_blob_paths(field: &ArrowField, parent: &[String], paths: &mut Vec<Vec<String>>) {
let mut path = parent.to_vec();
path.push(field.name().clone());
if field.is_blob_v2() {
paths.push(path);
return;
}
match field.data_type() {
DataType::Struct(children) => {
for child in children {
collect_blob_paths(child, &path, paths);
}
}
DataType::List(child)
| DataType::LargeList(child)
| DataType::FixedSizeList(child, _)
| DataType::Map(child, _) => collect_blob_paths(child, &path, paths),
_ => {}
}
}
fn schema_blob_paths(schema: &ArrowSchema) -> Vec<Vec<String>> {
let mut paths = Vec::new();
for field in schema.fields() {
collect_blob_paths(field, &[], &mut paths);
}
paths
}
fn transform_blob_field(
field: &ArrowField,
parent: &[String],
materialized: &HashSet<Vec<String>>,
) -> ArrowField {
let mut path = parent.to_vec();
path.push(field.name().clone());
if field.is_blob_v2() {
if materialized.contains(&path) {
return ArrowField::new(field.name(), DataType::LargeBinary, field.is_nullable());
}
return ArrowField::new(
field.name(),
BLOB_V2_DESC_FIELD.data_type().clone(),
field.is_nullable(),
)
.with_metadata(BLOB_V2_DESC_FIELD.metadata().clone());
}
let data_type = match field.data_type() {
DataType::Struct(children) => DataType::Struct(
children
.iter()
.map(|child| Arc::new(transform_blob_field(child, &path, materialized)))
.collect(),
),
DataType::List(child) => {
DataType::List(Arc::new(transform_blob_field(child, &path, materialized)))
}
DataType::LargeList(child) => {
DataType::LargeList(Arc::new(transform_blob_field(child, &path, materialized)))
}
DataType::FixedSizeList(child, size) => DataType::FixedSizeList(
Arc::new(transform_blob_field(child, &path, materialized)),
*size,
),
DataType::Map(child, sorted) => DataType::Map(
Arc::new(transform_blob_field(child, &path, materialized)),
*sorted,
),
_ => return field.clone(),
};
ArrowField::new(field.name(), data_type, field.is_nullable())
.with_metadata(field.metadata().clone())
}
fn blob_runtime_schema(schema: &ArrowSchema, materialized: &HashSet<Vec<String>>) -> SchemaRef {
Arc::new(ArrowSchema::new_with_metadata(
schema
.fields()
.iter()
.map(|field| Arc::new(transform_blob_field(field, &[], materialized)))
.collect::<Fields>(),
schema.metadata().clone(),
))
}
fn referenced_blob_paths(schema: &ArrowSchema, inputs: &[String]) -> Result<Vec<Vec<String>>> {
let input_paths = inputs
.iter()
.map(|input| {
parse_field_path(input).map_err(|error| Error::InvalidInput {
message: format!("invalid computed-column input path '{input}': {error}"),
})
})
.collect::<Result<Vec<_>>>()?;
Ok(schema_blob_paths(schema)
.into_iter()
.filter(|blob_path| {
input_paths.iter().any(|input_path| {
input_path.len() <= blob_path.len()
&& input_path
.iter()
.zip(blob_path)
.all(|(input, blob)| input == blob)
})
})
.collect())
}
/// Parse, resolve and compile `expression` against `schema`.
@@ -1185,7 +1337,14 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
message,
};
let planner = Planner::new(schema.clone());
// Blob v2 is a semantic type whose runtime expression ABI is
// `LargeBinary`. Parse against that ABI first so a direct Blob reference
// is not mistaken for its storage descriptor struct.
let all_blob_paths = schema_blob_paths(schema.as_ref())
.into_iter()
.collect::<HashSet<_>>();
let parsing_schema = blob_runtime_schema(schema.as_ref(), &all_blob_paths);
let planner = Planner::new(parsing_schema);
let parsed = planner
.parse_expr(expression)
.map_err(|e| invalid(e.to_string()))?;
@@ -1218,13 +1377,19 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
inputs.sort();
inputs.dedup();
let blob_paths = referenced_blob_paths(schema.as_ref(), &inputs)?;
let runtime_schema = blob_runtime_schema(
schema.as_ref(),
&blob_paths.iter().cloned().collect::<HashSet<_>>(),
);
// A nested input is recorded by its path but read through its root
// column; Schema::index_of resolves top-level names only. Resolved here
// rather than left to the planner so an unknown column names itself in
// the error instead of surfacing as a plan failure.
let mut indices = Vec::with_capacity(inputs.len());
for input in &inputs {
let index = schema
let index = runtime_schema
.index_of(root(input))
.map_err(|_| invalid(format!("unknown column '{input}'")))?;
if !indices.contains(&index) {
@@ -1237,7 +1402,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
// compiles the expression has to be built on the projected schema
// evaluation will actually read.
let read_schema = Arc::new(
schema
runtime_schema
.project(&indices)
.map_err(|e| invalid(e.to_string()))?,
);
@@ -1247,7 +1412,8 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
.map(|field| field.name().clone())
.collect();
let optimized = planner
let runtime_planner = Planner::new(runtime_schema);
let optimized = runtime_planner
.optimize_expr(parsed)
.map_err(|e| invalid(e.to_string()))?;
let physical = Planner::new(read_schema.clone())
@@ -1260,9 +1426,15 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
Ok(BoundExpression {
inputs,
roots,
read_schema,
physical,
data_type,
blob_paths: blob_paths
.iter()
.map(|path| {
let segments = path.iter().map(String::as_str).collect::<Vec<_>>();
format_field_path_minimal(&segments)
})
.collect(),
})
}
@@ -1278,7 +1450,10 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
/// batch may declare `a` and then `b = a + 1` in one commit. Refresh order
/// then matters, and refresh enforces it: `b` is refused while `a` still has
/// unfilled rows.
pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Vec<ArrowField>> {
fn plan_declarations(
schema: SchemaRef,
columns: &[ComputedColumnDeclaration],
) -> Result<Vec<ArrowField>> {
if columns.is_empty() {
return Err(Error::InvalidInput {
message: "at least one computed column is required".into(),
@@ -1288,17 +1463,38 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
let mut schema = schema;
let mut fields = Vec::with_capacity(columns.len());
for (name, expression) in columns {
if schema.field_with_name(name).is_ok() {
return Err(Error::ColumnAlreadyExists { name: name.clone() });
for declaration in columns {
if schema.field_with_name(&declaration.name).is_ok() {
return Err(Error::ColumnAlreadyExists {
name: declaration.name.clone(),
});
}
let bound = bind(schema.clone(), name, expression)?;
let bound = bind(schema.clone(), &declaration.name, &declaration.expression)?;
// Declared columns start entirely null, so nullability is a property
// of the declaration rather than of what the expression yields.
let field = ArrowField::new(name, bound.data_type, true)
.with_metadata(computed_column_metadata(expression, &bound.inputs));
let metadata = computed_column_metadata(&declaration.expression, &bound.inputs);
let field = match declaration.output {
ComputedColumnOutput::Inferred => {
ArrowField::new(&declaration.name, bound.data_type, true).with_metadata(metadata)
}
ComputedColumnOutput::BlobV2 => {
if bound.data_type != DataType::LargeBinary {
return Err(Error::InvalidExpression {
column: declaration.name.clone(),
message: format!(
"a Blob v2 computed output requires a LargeBinary expression, got {}",
bound.data_type
),
});
}
let field = crate::blob::blob(&declaration.name, true);
let mut blob_metadata = field.metadata().clone();
blob_metadata.extend(metadata);
field.with_metadata(blob_metadata)
}
};
schema = Arc::new(ArrowSchema::new_with_metadata(
schema
.fields()
@@ -1314,6 +1510,16 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
Ok(fields)
}
pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Vec<ArrowField>> {
let declarations = columns
.iter()
.map(|(name, expression)| {
ComputedColumnDeclaration::inferred(name.clone(), expression.clone())
})
.collect::<Vec<_>>();
plan_declarations(schema, &declarations)
}
/// Run the schema-level checks of
/// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) against
/// `schema` without committing: the Function-binding guard and the planning of
@@ -1350,9 +1556,9 @@ pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) ->
/// public way in.
pub(crate) fn declare(
schema: SchemaRef,
columns: &[(String, String)],
columns: &[ComputedColumnDeclaration],
) -> Result<NewColumnTransform> {
let fields = plan(schema, columns)?;
let fields = plan_declarations(schema, columns)?;
Ok(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(
fields,
))))
rust/lancedb/src/table/refresh.rs
+510 -6
View File
@@ -29,9 +29,12 @@
//! inputs masked to null first, so a poison value in a row nobody is filling
//! cannot fail the refresh.
use std::collections::HashSet;
use std::sync::Arc;
use arrow_array::{ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions};
use arrow_array::{
Array, ArrayRef, BooleanArray, LargeBinaryArray, RecordBatch, RecordBatchOptions,
};
use arrow_schema::Schema as ArrowSchema;
use datafusion_expr::ColumnarValue;
use futures::{Stream, StreamExt, TryStreamExt};
@@ -40,7 +43,7 @@ use lance::dataset::WriteDestination;
use lance::dataset::fragment::FileFragment;
use lance::dataset::transaction::Operation;
use lance_core::ROW_ID;
use lance_core::datatypes::Schema as LanceSchema;
use lance_core::datatypes::{BlobHandling, Schema as LanceSchema};
use serde::{Deserialize, Serialize};
use super::computed_columns::{BoundExpression, ComputedColumnKind, computed_column_from_field};
@@ -104,6 +107,7 @@ async fn execute_refresh_column_with_source(
fields: vec![field.clone()],
metadata: Default::default(),
};
let output_is_blob = field.is_blob_v2();
let mut rows_filled = 0u64;
let mut replacements = Vec::new();
@@ -113,7 +117,8 @@ async fn execute_refresh_column_with_source(
continue;
}
rows_filled += gained;
let values = fill_stream(&dataset, &fragment, bound.clone(), column).await?;
let values =
fill_stream(&dataset, &fragment, bound.clone(), column, output_is_blob).await?;
replacements.push(fragment.write_columns(values, &column_schema).await?);
}
@@ -294,12 +299,15 @@ fn evaluation_batch(
mask_out: Option<&BooleanArray>,
) -> lance_core::Result<RecordBatch> {
let mut columns = Vec::with_capacity(bound.roots.len());
let mut fields = Vec::with_capacity(bound.roots.len());
for name in &bound.roots {
let column = batch.column_by_name(name).ok_or_else(|| {
let index = batch.schema_ref().index_of(name).map_err(|_| {
lance_core::Error::invalid_input(format!(
"refreshing a computed column read no {name} column"
))
})?;
let column = batch.column(index);
fields.push(batch.schema_ref().field(index).clone());
// Rows outside the mask must not reach the expression: a value in a
// deleted or already-filled row can be one it would choke on.
columns.push(match mask_out {
@@ -308,7 +316,7 @@ fn evaluation_batch(
});
}
Ok(RecordBatch::try_new_with_options(
bound.read_schema.clone(),
Arc::new(ArrowSchema::new(fields)),
columns,
&RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
)?)
@@ -329,6 +337,64 @@ fn evaluate(bound: &BoundExpression, batch: &RecordBatch) -> lance_core::Result<
}
}
fn materialized_blob_ids(schema: &LanceSchema, paths: &[String]) -> Result<HashSet<u32>> {
paths
.iter()
.map(|path| {
let field = schema
.resolve(path)
.and_then(|fields| fields.last().copied())
.ok_or_else(|| Error::InvalidInput {
message: format!("computed Blob input '{path}' no longer exists"),
})?;
if !field.is_blob_v2() {
return Err(Error::InvalidInput {
message: format!("computed Blob input '{path}' is no longer Blob v2"),
});
}
u32::try_from(field.id).map_err(|_| Error::InvalidInput {
message: format!(
"computed Blob input '{path}' has invalid field id {}",
field.id
),
})
})
.collect()
}
fn configure_blob_inputs(
scanner: &mut lance::dataset::scanner::Scanner,
schema: &LanceSchema,
bound: &BoundExpression,
extra_blob_id: Option<u32>,
) -> Result<()> {
let mut ids = materialized_blob_ids(schema, &bound.blob_paths)?;
ids.extend(extra_blob_id);
scanner.blob_handling(BlobHandling::SomeBlobsBinary(ids));
Ok(())
}
fn blob_array_from_binary(array: &ArrayRef) -> lance_core::Result<ArrayRef> {
let values = array
.as_any()
.downcast_ref::<LargeBinaryArray>()
.ok_or_else(|| {
lance_core::Error::invalid_input(format!(
"a Blob v2 computed output produced {}, expected LargeBinary",
array.data_type()
))
})?;
let mut builder = lance::blob::BlobArrayBuilder::new(values.len());
for index in 0..values.len() {
if values.is_null(index) {
builder.push_null()?;
} else {
builder.push_bytes(values.value(index))?;
}
}
builder.finish()
}
/// How many rows of one fragment would gain a value.
///
/// Scans only the unfilled live rows -- deleted rows never reach the
@@ -347,6 +413,7 @@ async fn count_fragment_gains(
.with_row_id()
.filter(&format!("{} IS NULL", quote_identifier(column)))?
.project(&bound.roots)?;
configure_blob_inputs(&mut scanner, dataset.schema(), bound, None)?;
let mut gained = 0u64;
let mut batches = scanner.try_into_stream().await?;
@@ -368,6 +435,7 @@ async fn fill_stream(
fragment: &FileFragment,
bound: Arc<BoundExpression>,
column: &str,
output_is_blob: bool,
) -> Result<impl Stream<Item = lance_core::Result<RecordBatch>> + Send + use<>> {
let mut projection: Vec<String> = bound.roots.clone();
projection.push(column.to_string());
@@ -377,6 +445,20 @@ async fn fill_stream(
.with_row_id()
.include_deleted_rows()
.project(&projection)?;
let output_blob_id = output_is_blob
.then(|| {
dataset
.schema()
.field(column)
.and_then(|field| u32::try_from(field.id).ok())
})
.flatten();
configure_blob_inputs(
&mut scanner,
dataset.schema(),
bound.as_ref(),
output_blob_id,
)?;
let projected = Arc::new(ArrowSchema::new(vec![
ArrowSchema::from(dataset.schema())
@@ -412,6 +494,11 @@ async fn fill_stream(
let computed = evaluate(&bound, &evaluation_batch(&batch, &bound, Some(&keep))?)?;
let merged = arrow_select::zip::zip(&fill, &computed, existing)?;
let merged = if output_is_blob {
blob_array_from_binary(&merged)?
} else {
merged
};
Ok(RecordBatch::try_new(projected.clone(), vec![merged])?)
}))
}
@@ -420,8 +507,9 @@ async fn fill_stream(
mod tests {
use std::sync::Arc;
use arrow_array::{Int32Array, record_batch};
use arrow_array::{Array, Int32Array, LargeBinaryArray, RecordBatch, record_batch};
use futures::TryStreamExt;
use lance_core::ROW_ID;
use crate::connect;
use crate::query::{ExecutableQuery, QueryBase, Select};
@@ -1164,4 +1252,420 @@ mod tests {
let err = table.refresh_column("embedding").await.unwrap_err();
assert!(matches!(err, Error::NotSupported { message } if message.contains("udf")));
}
fn blob_batch(ids: Vec<i32>, payloads: Vec<Option<&[u8]>>) -> RecordBatch {
use arrow_array::Int32Array;
use arrow_schema::{Field, Schema};
let mut builder = lance::blob::BlobArrayBuilder::new(payloads.len());
for payload in payloads {
match payload {
Some(payload) => builder.push_bytes(payload).unwrap(),
None => builder.push_null().unwrap(),
}
}
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", arrow_schema::DataType::Int32, false),
crate::blob("image", true),
])),
vec![Arc::new(Int32Array::from(ids)), builder.finish().unwrap()],
)
.unwrap()
}
async fn create_blob_table(path: &std::path::Path, batch: RecordBatch) -> Table {
let conn = connect(path.to_str().unwrap()).execute().await.unwrap();
conn.create_table("blobs", batch).execute().await.unwrap()
}
#[tokio::test]
async fn test_refresh_materializes_top_level_blob_input() {
let tmp = tempfile::tempdir().unwrap();
let table = create_blob_table(
tmp.path(),
blob_batch(vec![1, 2, 3], vec![Some(b"hello"), Some(b""), None]),
)
.await;
table
.add_columns()
.computed("payload_copy", "image")
.execute()
.await
.unwrap();
let result = table.refresh_column("payload_copy").await.unwrap();
assert_eq!(result.rows_filled, 2);
let batches = table
.query()
.select(Select::columns(&["payload_copy"]))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let payloads = batches[0]
.column_by_name("payload_copy")
.unwrap()
.as_any()
.downcast_ref::<LargeBinaryArray>()
.unwrap();
assert_eq!(payloads.value(0), b"hello");
assert_eq!(payloads.value(1), b"");
assert!(payloads.is_null(2));
assert_eq!(
table
.count_rows(Some("payload_copy IS NULL".to_string()))
.await
.unwrap(),
1
);
}
#[tokio::test]
async fn test_refresh_publishes_explicit_blob_output() {
use arrow_array::UInt64Array;
use lance_arrow::{
BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, BLOB_INLINE_SIZE_THRESHOLD_META_KEY,
};
use lance_core::datatypes::BlobKind;
use crate::table::schema_evolution::FieldMetadataUpdate;
let tmp = tempfile::tempdir().unwrap();
let table = create_blob_table(
tmp.path(),
blob_batch(
vec![1, 2, 3, 4],
vec![Some(b"hello"), Some(b"ab"), Some(b""), None],
),
)
.await;
table
.add_columns()
.computed_blob("image_copy", "image")
.execute()
.await
.unwrap();
table
.update_field_metadata(&[FieldMetadataUpdate::new("image_copy")
.set(BLOB_INLINE_SIZE_THRESHOLD_META_KEY, "1")
.set(BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, "4")])
.await
.unwrap();
let first_refresh = table.refresh_column("image_copy").await.unwrap();
assert_eq!(first_refresh.rows_filled, 3);
assert_eq!(
table.blob_columns().await.unwrap(),
vec!["image".to_string(), "image_copy".to_string()]
);
let batches = table
.query()
.with_row_id()
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let batch = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap();
assert!(
batch
.column_by_name("image_copy")
.unwrap()
.as_any()
.is::<arrow_array::StructArray>()
);
let row_ids = batch
.column_by_name(ROW_ID)
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap()
.values()
.to_vec();
let original = table.fetch_blobs("image", &row_ids).await.unwrap();
let copied = table.fetch_blobs("image_copy", &row_ids).await.unwrap();
assert_eq!(original, copied);
let ids = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
let files = table
.fetch_blob_files("image_copy", &row_ids)
.await
.unwrap();
let mut layouts = ids
.values()
.iter()
.copied()
.zip(files)
.map(|(id, file)| (id, file.and_then(|file| file.kind())))
.collect::<Vec<_>>();
layouts.sort_by_key(|(id, _)| *id);
assert_eq!(
layouts,
vec![
(1, Some(BlobKind::Dedicated)),
(2, Some(BlobKind::Packed)),
(3, Some(BlobKind::Inline)),
(4, None),
]
);
table
.add(blob_batch(vec![5], vec![Some(b"appended")]))
.execute()
.await
.unwrap();
table
.optimize(crate::table::OptimizeAction::Compact {
options: crate::table::CompactionOptions::default(),
remap_options: None,
})
.await
.unwrap();
assert_eq!(
table
.refresh_column("image_copy")
.await
.unwrap()
.rows_filled,
1
);
assert_eq!(
table
.refresh_column("image_copy")
.await
.unwrap()
.rows_filled,
0
);
table.checkout(first_refresh.version).await.unwrap();
assert_eq!(table.count_rows(None).await.unwrap(), 4);
assert_eq!(
table.blob_columns().await.unwrap(),
vec!["image".to_string(), "image_copy".to_string()]
);
table.checkout_latest().await.unwrap();
}
#[tokio::test]
async fn test_explicit_blob_output_requires_large_binary_expression() {
let tmp = tempfile::tempdir().unwrap();
let table = create_blob_table(tmp.path(), blob_batch(vec![1], vec![Some(b"hello")])).await;
let error = table
.add_columns()
.computed_blob("invalid", "id + 1")
.execute()
.await
.unwrap_err();
assert!(matches!(
error,
Error::InvalidExpression { column, message }
if column == "invalid" && message.contains("requires a LargeBinary expression")
));
}
#[tokio::test]
async fn test_refresh_materializes_nested_struct_blob_input() {
use arrow_array::{Int32Array, StructArray};
use arrow_schema::{DataType, Field, Fields, Schema};
let tmp = tempfile::tempdir().unwrap();
let mut blob_builder = lance::blob::BlobArrayBuilder::new(2);
blob_builder.push_bytes(b"nested").unwrap();
blob_builder.push_null().unwrap();
let blob_field = crate::blob("image", true);
let metadata_fields = Fields::from(vec![blob_field.clone()]);
let metadata = StructArray::new(
metadata_fields.clone(),
vec![blob_builder.finish().unwrap()],
None,
);
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("metadata", DataType::Struct(metadata_fields), true),
])),
vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(metadata)],
)
.unwrap();
let table = create_blob_table(tmp.path(), batch).await;
table
.add_columns()
.computed("payload_copy", "metadata.image")
.execute()
.await
.unwrap();
assert_eq!(
table
.refresh_column("payload_copy")
.await
.unwrap()
.rows_filled,
1
);
let batches = table
.query()
.select(Select::columns(&["payload_copy"]))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let payloads = batches[0]
.column_by_name("payload_copy")
.unwrap()
.as_any()
.downcast_ref::<LargeBinaryArray>()
.unwrap();
assert_eq!(payloads.value(0), b"nested");
assert!(payloads.is_null(1));
}
#[tokio::test]
async fn test_refresh_preserves_list_shape_when_materializing_blob_input() {
use arrow_array::{Int32Array, ListArray};
use arrow_buffer::{OffsetBuffer, ScalarBuffer};
use arrow_schema::{DataType, Field, Schema};
let tmp = tempfile::tempdir().unwrap();
let mut blob_builder = lance::blob::BlobArrayBuilder::new(3);
blob_builder.push_bytes(b"a").unwrap();
blob_builder.push_bytes(b"bb").unwrap();
blob_builder.push_null().unwrap();
let item = Arc::new(crate::blob("item", true));
let images = ListArray::new(
item.clone(),
OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])),
blob_builder.finish().unwrap(),
None,
);
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("images", DataType::List(item), true),
])),
vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(images)],
)
.unwrap();
let table = create_blob_table(tmp.path(), batch).await;
table
.add_columns()
.computed("image_payloads", "images")
.execute()
.await
.unwrap();
assert_eq!(
table
.refresh_column("image_payloads")
.await
.unwrap()
.rows_filled,
2
);
let batches = table
.query()
.select(Select::columns(&["image_payloads"]))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let output = batches[0]
.column_by_name("image_payloads")
.unwrap()
.as_any()
.downcast_ref::<ListArray>()
.unwrap();
assert_eq!(output.value_offsets(), &[0, 2, 3]);
assert!(output.values().as_any().is::<LargeBinaryArray>());
}
#[tokio::test]
async fn test_refresh_materializes_external_blob_input() {
use arrow_array::{Int32Array, StringArray};
use arrow_schema::{DataType, Field, Schema};
let tmp = tempfile::tempdir().unwrap();
let payload = b"external-payload";
let path = tmp.path().join("payload.bin");
std::fs::write(&path, payload).unwrap();
let uri = url::Url::from_file_path(path).unwrap().to_string();
let conn = connect(tmp.path().join("db").to_str().unwrap())
.execute()
.await
.unwrap();
let table = conn
.create_empty_table(
"external",
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
crate::blob("image", true),
])),
)
.execute()
.await
.unwrap();
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("image", DataType::Utf8, true),
])),
vec![
Arc::new(Int32Array::from(vec![1])),
Arc::new(StringArray::from(vec![Some(uri)])),
],
)
.unwrap();
table
.add(batch)
.allow_external_blob_outside_bases(true)
.execute()
.await
.unwrap();
table
.add_columns()
.computed("payload_copy", "image")
.execute()
.await
.unwrap();
assert_eq!(
table
.refresh_column("payload_copy")
.await
.unwrap()
.rows_filled,
1
);
let batches = table
.query()
.select(Select::columns(&["payload_copy"]))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let payloads = batches[0]
.column_by_name("payload_copy")
.unwrap()
.as_any()
.downcast_ref::<LargeBinaryArray>()
.unwrap();
assert_eq!(payloads.value(0), payload);
}
}
rust/lancedb/src/table/schema_evolution.rs
+1 -1
View File
@@ -124,7 +124,7 @@ pub(crate) async fn execute_add_columns(
/// declaration metadata.
pub(crate) async fn execute_declare(
table: &NativeTable,
columns: &[(String, String)],
columns: &[computed_columns::ComputedColumnDeclaration],
) -> Result<AddColumnsResult> {
use lance::dataset::mem_wal::DatasetMemWalExt;