mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-04 12:38:38 +00:00
refactor: inherit blob metadata for computed projections
This commit is contained in:
@@ -352,8 +352,7 @@ 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 | pa.Field, str]],
|
||||
self, columns: list[tuple[str, str]]
|
||||
) -> AddColumnsResult: ...
|
||||
async def add_function_columns(
|
||||
self, application_json: str, output_name: Optional[str]
|
||||
|
||||
@@ -13,7 +13,6 @@ from typing import (
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
Union,
|
||||
Literal,
|
||||
overload,
|
||||
@@ -984,7 +983,7 @@ class RemoteTable(Table):
|
||||
| FunctionApplication
|
||||
| None = None,
|
||||
*,
|
||||
computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
|
||||
computed: Dict[str, str] | None = None,
|
||||
) -> AddColumnsResult:
|
||||
return LOOP.run(self._table.add_columns(transforms, computed=computed))
|
||||
|
||||
|
||||
@@ -913,26 +913,6 @@ def _normalize_progress(progress):
|
||||
return progress, False
|
||||
|
||||
|
||||
def _normalize_computed_columns(
|
||||
computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]],
|
||||
) -> list[tuple[str | pa.Field, str]]:
|
||||
columns: list[tuple[str | pa.Field, str]] = []
|
||||
declarations = computed.items() if isinstance(computed, dict) else computed
|
||||
for declaration in declarations:
|
||||
if not isinstance(declaration, (tuple, list)) or len(declaration) != 2:
|
||||
raise TypeError(
|
||||
"computed sequences must contain (column name or pyarrow Field, "
|
||||
"SQL expression) pairs"
|
||||
)
|
||||
field, expression = declaration
|
||||
if not isinstance(field, (str, pa.Field)):
|
||||
raise TypeError("computed targets must be column names or pyarrow Fields")
|
||||
if not isinstance(expression, str):
|
||||
raise TypeError("computed values must be SQL expression strings")
|
||||
columns.append((field, expression))
|
||||
return columns
|
||||
|
||||
|
||||
class Table(ABC):
|
||||
"""
|
||||
A Table is a collection of Records in a LanceDB Database.
|
||||
@@ -2162,7 +2142,7 @@ class Table(ABC):
|
||||
| pa.Schema
|
||||
| None = None,
|
||||
*,
|
||||
computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
|
||||
computed: Dict[str, str] | None = None,
|
||||
):
|
||||
"""
|
||||
Add new columns with defined values.
|
||||
@@ -2184,15 +2164,12 @@ class Table(ABC):
|
||||
atomic binding; aliases come from ``rename(columns=...)``.
|
||||
Function columns are supported only on LanceDB Cloud and
|
||||
Enterprise.
|
||||
computed: Dict[str, str] or Sequence[Tuple[str | pa.Field, str]], optional
|
||||
computed: Dict[str, str], optional
|
||||
A mapping from output column names to SQL expressions derives each
|
||||
output field from its expression. An ordered sequence may instead
|
||||
use a pyarrow Field as a target, supplying its name, type,
|
||||
nullability, and extension metadata; use
|
||||
``(lancedb.blob("name"), expression)`` for a Blob v2 output.
|
||||
Explicit fields must be nullable and their expression result type
|
||||
must be compatible. Mapping or sequence order is declaration and
|
||||
dependency order.
|
||||
output field from its expression. A direct projection of a Blob v2
|
||||
field inherits Blob v2 semantics; other expressions derive their
|
||||
ordinary Arrow type. Mapping order is declaration and dependency
|
||||
order.
|
||||
|
||||
Unlike ``transforms``, the expression is stored rather than
|
||||
evaluated now: the column is committed with no values, and rows get
|
||||
@@ -4324,7 +4301,7 @@ class LanceTable(Table):
|
||||
| pa.Schema
|
||||
| None = None,
|
||||
*,
|
||||
computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
|
||||
computed: Dict[str, str] | None = None,
|
||||
) -> AddColumnsResult:
|
||||
return LOOP.run(self._table.add_columns(transforms, computed=computed))
|
||||
|
||||
@@ -6272,7 +6249,7 @@ class AsyncTable:
|
||||
| pa.Schema
|
||||
| None = None,
|
||||
*,
|
||||
computed: dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
|
||||
computed: dict[str, str] | None = None,
|
||||
) -> AddColumnsResult:
|
||||
"""
|
||||
Add new columns with defined values.
|
||||
@@ -6292,15 +6269,12 @@ class AsyncTable:
|
||||
atomic binding; aliases come from ``rename(columns=...)``.
|
||||
Function columns are supported only on LanceDB Cloud and
|
||||
Enterprise.
|
||||
computed: Dict[str, str] or Sequence[Tuple[str | pa.Field, str]], optional
|
||||
computed: Dict[str, str], optional
|
||||
A mapping from output column names to SQL expressions derives each
|
||||
output field from its expression. An ordered sequence may instead
|
||||
use a pyarrow Field as a target, supplying its name, type,
|
||||
nullability, and extension metadata; use
|
||||
``(lancedb.blob("name"), expression)`` for a Blob v2 output.
|
||||
Explicit fields must be nullable and their expression result type
|
||||
must be compatible. Mapping or sequence order is declaration and
|
||||
dependency order.
|
||||
output field from its expression. A direct projection of a Blob v2
|
||||
field inherits Blob v2 semantics; other expressions derive their
|
||||
ordinary Arrow type. Mapping order is declaration and dependency
|
||||
order.
|
||||
|
||||
Unlike ``transforms``, the expression is stored rather than
|
||||
evaluated now: the column is committed with no values, and rows get
|
||||
@@ -6358,9 +6332,7 @@ class AsyncTable:
|
||||
raise ValueError(
|
||||
"add_columns cannot take both transforms and computed columns"
|
||||
)
|
||||
return await self._inner.add_computed_columns(
|
||||
_normalize_computed_columns(computed)
|
||||
)
|
||||
return await self._inner.add_computed_columns(list(computed.items()))
|
||||
if transforms is None:
|
||||
raise ValueError("add_columns requires transforms or computed columns")
|
||||
if isinstance(transforms, pa.Schema):
|
||||
|
||||
@@ -4087,7 +4087,7 @@ def test_computed_column_rejects_transforms_and_computed_together(tmp_path):
|
||||
table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"})
|
||||
|
||||
|
||||
def test_computed_column_blob_input_and_explicit_output(tmp_path):
|
||||
def test_computed_column_blob_projection_inherits_semantics(tmp_path):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("computed_column_blob", schema=schema)
|
||||
@@ -4099,49 +4099,17 @@ def test_computed_column_blob_input_and_explicit_output(tmp_path):
|
||||
]
|
||||
)
|
||||
|
||||
table.add_columns(
|
||||
computed=[
|
||||
(lancedb.blob("image_copy"), "image"),
|
||||
("payload_copy", "image_copy"),
|
||||
]
|
||||
)
|
||||
table.add_columns(computed={"image_copy": "image", "second_copy": "image_copy"})
|
||||
assert table.refresh_column("image_copy").rows_filled == 2
|
||||
assert table.refresh_column("payload_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"]
|
||||
assert table.refresh_column("second_copy").rows_filled == 2
|
||||
assert table.blob_columns() == ["image", "image_copy", "second_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])
|
||||
copied = table.fetch_blobs("second_copy", [row_id for _, row_id in rows])
|
||||
assert copied.to_pylist() == [b"hello", b"", None]
|
||||
|
||||
|
||||
def test_blob_output_declaration_rejects_eager_transforms(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("computed_column_blob_mixed", [{"x": 1}])
|
||||
with pytest.raises(ValueError):
|
||||
table.add_columns(
|
||||
{"a": "x + 1"},
|
||||
computed=[(lancedb.blob("b"), "x")],
|
||||
)
|
||||
|
||||
|
||||
def test_computed_column_validates_declaration_mapping(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("computed_column_mapping", [{"x": 1}])
|
||||
|
||||
with pytest.raises(
|
||||
TypeError, match="targets must be column names or pyarrow Fields"
|
||||
):
|
||||
table.add_columns(computed={42: "x"}) # type: ignore[dict-item]
|
||||
with pytest.raises(TypeError, match="values must be SQL expression strings"):
|
||||
table.add_columns(computed={"copy": 42}) # type: ignore[dict-item]
|
||||
with pytest.raises(TypeError, match="sequences must contain"):
|
||||
table.add_columns(computed=[("copy", "x", "extra")]) # type: ignore[list-item]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_computed_column_async(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
|
||||
+6
-21
@@ -13,16 +13,15 @@ use crate::{
|
||||
};
|
||||
use arrow::{
|
||||
array::{Array, LargeBinaryArray},
|
||||
datatypes::{DataType, Field, Schema},
|
||||
datatypes::{DataType, Schema},
|
||||
ffi_stream::ArrowArrayStreamReader,
|
||||
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
|
||||
};
|
||||
use lancedb::blob::{BlobFile, BlobRangeRequest};
|
||||
use lancedb::index::scalar::FtsIndexBuilder;
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration, ComputedColumnDeclaration, Duration, FieldMetadataUpdate,
|
||||
FtsToken as LanceDbFtsToken, NewColumnTransform, OptimizeAction, OptimizeOptions, Ref,
|
||||
Table as LanceDbTable,
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use pyo3::{
|
||||
@@ -101,12 +100,6 @@ enum PredicateArg {
|
||||
Sql(String),
|
||||
}
|
||||
|
||||
#[derive(FromPyObject)]
|
||||
pub enum ComputedColumnFieldArg {
|
||||
Name(String),
|
||||
Field(PyArrowType<Field>),
|
||||
}
|
||||
|
||||
/// Statistics about a compaction operation.
|
||||
#[pyclass(get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -1584,21 +1577,13 @@ impl Table {
|
||||
|
||||
pub fn add_computed_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
columns: Vec<(ComputedColumnFieldArg, String)>,
|
||||
columns: Vec<(String, String)>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let mut builder = inner.add_columns();
|
||||
for (field, expression) in columns {
|
||||
let declaration = match field {
|
||||
ComputedColumnFieldArg::Name(name) => {
|
||||
ComputedColumnDeclaration::inferred(name, expression)
|
||||
}
|
||||
ComputedColumnFieldArg::Field(PyArrowType(field)) => {
|
||||
ComputedColumnDeclaration::with_field(field, expression)
|
||||
}
|
||||
};
|
||||
builder = builder.computed_column(declaration);
|
||||
for (name, expression) in columns {
|
||||
builder = builder.computed(name, expression);
|
||||
}
|
||||
let result = builder.execute().await.infer_error()?;
|
||||
Ok(AddColumnsResult::from(result))
|
||||
|
||||
@@ -55,7 +55,7 @@ use crate::{
|
||||
};
|
||||
use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
|
||||
use arrow_ipc::reader::{FileReader, StreamReader};
|
||||
use arrow_schema::{ArrowError, DataType, Schema as ArrowSchema, SchemaRef};
|
||||
use arrow_schema::{ArrowError, DataType, SchemaRef};
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use datafusion_common::DataFusionError;
|
||||
@@ -3174,40 +3174,24 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn add_computed_columns(
|
||||
&self,
|
||||
columns: &[crate::table::computed_columns::ComputedColumnDeclaration],
|
||||
) -> Result<AddColumnsResult> {
|
||||
async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result<AddColumnsResult> {
|
||||
self.check_mutable().await?;
|
||||
crate::table::computed_columns::ensure_no_function_bindings_for_mutation(
|
||||
self.schema().await?.as_ref(),
|
||||
"schema evolution",
|
||||
)?;
|
||||
// The server plans the declaration: expression validation and the
|
||||
// persisted binding happen there. An explicit field carries its own
|
||||
// output schema and extension metadata.
|
||||
let entries =
|
||||
columns
|
||||
.iter()
|
||||
.map(|column| -> Result<serde_json::Value> {
|
||||
let mut entry = serde_json::json!({ "computed": column.expression() });
|
||||
if let Some(field) = column.field() {
|
||||
let mut schema = lance_namespace::schema::arrow_schema_to_json(
|
||||
&ArrowSchema::new(vec![field.clone()]),
|
||||
)?;
|
||||
entry["field"] = serde_json::to_value(schema.fields.remove(0)).map_err(
|
||||
|error| Error::Runtime {
|
||||
message: format!(
|
||||
"failed to serialize explicit computed output field: {error}"
|
||||
),
|
||||
},
|
||||
)?;
|
||||
} else {
|
||||
entry["name"] = serde_json::json!(column.name());
|
||||
}
|
||||
Ok(entry)
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
// The server plans the declaration against its table schema, including
|
||||
// Blob v2 semantics inherited by a direct field projection.
|
||||
let entries = columns
|
||||
.iter()
|
||||
.map(
|
||||
|(name, expression)| lance_namespace::models::AddColumnsEntry {
|
||||
name: name.clone(),
|
||||
computed: Some(Some(expression.clone())),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.collect::<Vec<_>>();
|
||||
let mut body = serde_json::json!({ "new_columns": entries });
|
||||
self.apply_branch_body(&mut body);
|
||||
let request = self
|
||||
@@ -7404,8 +7388,8 @@ mod tests {
|
||||
assert_eq!(result.version, if old_server { 0 } else { 43 });
|
||||
}
|
||||
|
||||
/// An inferred declaration is sent as `{name, computed}` for the server to
|
||||
/// plan; the client never types the expression itself.
|
||||
/// A declaration is sent as `{name, computed}` for the server to plan; the
|
||||
/// client never types the expression itself.
|
||||
#[tokio::test]
|
||||
async fn test_add_computed_columns_sends_the_expression() {
|
||||
let table = Table::new_with_handler("my_table", |request| match request.url().path() {
|
||||
@@ -7435,58 +7419,6 @@ mod tests {
|
||||
assert_eq!(result.version, 7);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_add_computed_column_sends_explicit_output_field() {
|
||||
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!([{
|
||||
"computed": "image",
|
||||
"field": {
|
||||
"metadata": {
|
||||
"ARROW:extension:name": "lance.blob.v2"
|
||||
},
|
||||
"name": "image_copy",
|
||||
"nullable": true,
|
||||
"type": {
|
||||
"type": "struct",
|
||||
"fields": [
|
||||
{
|
||||
"name": "data",
|
||||
"nullable": true,
|
||||
"type": { "type": "large_binary" }
|
||||
},
|
||||
{
|
||||
"name": "uri",
|
||||
"nullable": true,
|
||||
"type": { "type": "utf8" }
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}])
|
||||
);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 8}"#.to_string())
|
||||
.unwrap()
|
||||
}
|
||||
path => panic!("Unexpected path: {path}"),
|
||||
});
|
||||
|
||||
let result = table
|
||||
.add_columns()
|
||||
.computed_field(crate::blob("image_copy", true), "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| {
|
||||
|
||||
@@ -95,8 +95,7 @@ pub use cherry_pick::{
|
||||
};
|
||||
pub use chrono::Duration;
|
||||
pub use computed_columns::{
|
||||
ComputedColumn, ComputedColumnDeclaration, ComputedColumnKind, computed_column_from_field,
|
||||
computed_columns,
|
||||
ComputedColumn, ComputedColumnKind, computed_column_from_field, computed_columns,
|
||||
};
|
||||
pub use delete::DeleteResult;
|
||||
use futures::future::join_all;
|
||||
@@ -752,10 +751,10 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
///
|
||||
/// Where the declaration is planned depends on the backend: a local table
|
||||
/// validates and types the expression itself, while a remote one sends the
|
||||
/// expression and any explicit output field for the server to plan.
|
||||
/// expression for the server to plan.
|
||||
async fn add_computed_columns(
|
||||
&self,
|
||||
_columns: &[computed_columns::ComputedColumnDeclaration],
|
||||
_columns: &[(String, String)],
|
||||
) -> Result<AddColumnsResult> {
|
||||
Err(Error::NotSupported {
|
||||
message: "computed columns are not supported on this table type".into(),
|
||||
@@ -3524,10 +3523,7 @@ impl BaseTable for NativeTable {
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn add_computed_columns(
|
||||
&self,
|
||||
columns: &[computed_columns::ComputedColumnDeclaration],
|
||||
) -> Result<AddColumnsResult> {
|
||||
async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result<AddColumnsResult> {
|
||||
let result = schema_evolution::execute_declare(self, columns).await?;
|
||||
self.bump_freshness();
|
||||
Ok(result)
|
||||
|
||||
@@ -5,11 +5,9 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::Field as ArrowField;
|
||||
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};
|
||||
@@ -18,7 +16,7 @@ use crate::{Error, Result};
|
||||
pub struct AddColumnsBuilder {
|
||||
parent: Arc<dyn BaseTable>,
|
||||
transform: Option<NewColumnTransform>,
|
||||
computed: Vec<ComputedColumnDeclaration>,
|
||||
computed: Vec<(String, String)>,
|
||||
function: Option<(FunctionApplication, Option<String>)>,
|
||||
read_columns: Option<Vec<String>>,
|
||||
}
|
||||
@@ -84,37 +82,8 @@ impl AddColumnsBuilder {
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn computed(self, name: impl Into<String>, expression: impl Into<String>) -> Self {
|
||||
self.computed_column(ComputedColumnDeclaration::inferred(name, expression))
|
||||
}
|
||||
|
||||
/// Add a computed column whose output schema is the supplied Arrow field.
|
||||
///
|
||||
/// Extension semantics such as Blob v2 come from the field metadata. The
|
||||
/// expression's inferred type must be compatible with the field semantics.
|
||||
///
|
||||
/// ```
|
||||
/// # use lancedb::Table;
|
||||
/// # async fn declare(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// table
|
||||
/// .add_columns()
|
||||
/// .computed_field(lancedb::blob("image_copy", true), "image")
|
||||
/// .execute()
|
||||
/// .await?;
|
||||
/// table.refresh_column("image_copy").await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn computed_field(self, field: ArrowField, expression: impl Into<String>) -> Self {
|
||||
self.computed_column(ComputedColumnDeclaration::with_field(field, expression))
|
||||
}
|
||||
|
||||
/// Add one pre-built computed-column declaration.
|
||||
///
|
||||
/// Declarations share one ordered stream, so a later expression may
|
||||
/// reference a column declared earlier in the same builder.
|
||||
pub fn computed_column(mut self, declaration: ComputedColumnDeclaration) -> Self {
|
||||
self.computed.push(declaration);
|
||||
pub fn computed(mut self, name: impl Into<String>, expression: impl Into<String>) -> Self {
|
||||
self.computed.push((name.into(), expression.into()));
|
||||
self
|
||||
}
|
||||
|
||||
|
||||
@@ -10,9 +10,10 @@
|
||||
//!
|
||||
//! The rule is tagged by kind ([`ComputedColumnKind`]) because kinds differ in
|
||||
//! where the column's type and inputs come from. A SQL expression determines
|
||||
//! its inputs and physical result type; a caller may also supply an Arrow field
|
||||
//! when the output needs extension metadata that expressions do not carry.
|
||||
//! A kind resolved through a registry cannot be typed without consulting it.
|
||||
//! its inputs and physical result type. A direct projection of a Blob v2 field
|
||||
//! also inherits that field's semantic type while execution continues to use
|
||||
//! `LargeBinary`. A kind resolved through a registry cannot be typed without
|
||||
//! consulting it.
|
||||
//! Registered Functions use an exact remote version plus a schema-level
|
||||
//! Function binding; unknown newer kinds remain readable and fail closed
|
||||
//! before mutation.
|
||||
@@ -24,19 +25,19 @@ 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_common::{ScalarValue, tree_node::TreeNode};
|
||||
use datafusion_expr::Expr;
|
||||
use datafusion_physical_plan::PhysicalExpr;
|
||||
use lance::dataset::NewColumnTransform;
|
||||
use lance_arrow::FieldExt;
|
||||
use lance_core::datatypes::{
|
||||
BLOB_V2_DESC_FIELD, BlobV2Layout, format_field_path_minimal, parse_field_path,
|
||||
};
|
||||
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};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::function::{FunctionApplication, FunctionBinding};
|
||||
use crate::utils::resolve_arrow_field_path;
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Field metadata key marking a column as computed. The value is `"true"`.
|
||||
@@ -69,63 +70,6 @@ pub const SQL_KIND: &str = "sql";
|
||||
/// Value of [`KIND_META_KEY`] for a registered Function binding.
|
||||
pub const FUNCTION_KIND: &str = "function";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum ComputedColumnTarget {
|
||||
Inferred(String),
|
||||
Explicit(ArrowField),
|
||||
}
|
||||
|
||||
/// One SQL computed-column declaration before it is planned.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ComputedColumnDeclaration {
|
||||
target: ComputedColumnTarget,
|
||||
expression: String,
|
||||
}
|
||||
|
||||
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 {
|
||||
target: ComputedColumnTarget::Inferred(name.into()),
|
||||
expression: expression.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Declare a computed column with an explicit Arrow field.
|
||||
///
|
||||
/// The field is the source of truth for the output's name, type,
|
||||
/// nullability, and extension metadata. Its semantics must accept the
|
||||
/// expression's inferred type, and it must be nullable because declaration
|
||||
/// commits the column before refresh fills its values.
|
||||
pub fn with_field(field: ArrowField, expression: impl Into<String>) -> Self {
|
||||
Self {
|
||||
target: ComputedColumnTarget::Explicit(field),
|
||||
expression: expression.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Name of the column to declare.
|
||||
pub fn name(&self) -> &str {
|
||||
match &self.target {
|
||||
ComputedColumnTarget::Inferred(name) => name,
|
||||
ComputedColumnTarget::Explicit(field) => field.name(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Immutable SQL expression evaluated by refresh.
|
||||
pub fn expression(&self) -> &str {
|
||||
&self.expression
|
||||
}
|
||||
|
||||
/// Explicit output field, or `None` when the expression determines it.
|
||||
pub fn field(&self) -> Option<&ArrowField> {
|
||||
match &self.target {
|
||||
ComputedColumnTarget::Inferred(_) => None,
|
||||
ComputedColumnTarget::Explicit(field) => Some(field),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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";
|
||||
@@ -1238,6 +1182,36 @@ pub(crate) struct BoundExpression {
|
||||
pub data_type: DataType,
|
||||
/// Blob v2 leaves the scan must materialize as `LargeBinary`.
|
||||
pub blob_paths: Vec<String>,
|
||||
/// A directly projected Blob v2 field whose semantics the output inherits.
|
||||
projected_blob_field: Option<ArrowField>,
|
||||
}
|
||||
|
||||
fn is_direct_field_projection(expr: &Expr) -> bool {
|
||||
match expr {
|
||||
Expr::Column(_) => true,
|
||||
Expr::ScalarFunction(function)
|
||||
if function.name() == "get_field" && function.args.len() == 2 =>
|
||||
{
|
||||
is_direct_field_projection(&function.args[0])
|
||||
&& matches!(
|
||||
&function.args[1],
|
||||
Expr::Literal(ScalarValue::Utf8(Some(_)), _)
|
||||
)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn projected_blob_field(schema: &ArrowSchema, expr: &Expr) -> Result<Option<ArrowField>> {
|
||||
if !is_direct_field_projection(expr) {
|
||||
return Ok(None);
|
||||
}
|
||||
let paths = Planner::column_names_in_expr(expr);
|
||||
let [path] = paths.as_slice() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let (_, field) = resolve_arrow_field_path(schema, path)?;
|
||||
Ok(field.is_blob_v2().then_some(field))
|
||||
}
|
||||
|
||||
fn collect_blob_paths(field: &ArrowField, parent: &[String], paths: &mut Vec<Vec<String>>) {
|
||||
@@ -1372,6 +1346,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
|
||||
let parsed = planner
|
||||
.parse_expr(expression)
|
||||
.map_err(|e| invalid(e.to_string()))?;
|
||||
let projected_blob_field = projected_blob_field(schema.as_ref(), &parsed)?;
|
||||
|
||||
// A declaration is evaluated more than once -- staging and writing are
|
||||
// separate passes, and a refresh years later replays the same text -- so
|
||||
@@ -1459,6 +1434,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
|
||||
format_field_path_minimal(&segments)
|
||||
})
|
||||
.collect(),
|
||||
projected_blob_field,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1474,10 +1450,7 @@ 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.
|
||||
fn plan_declarations(
|
||||
schema: SchemaRef,
|
||||
columns: &[ComputedColumnDeclaration],
|
||||
) -> Result<Vec<ArrowField>> {
|
||||
fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<Vec<ArrowField>> {
|
||||
if columns.is_empty() {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "at least one computed column is required".into(),
|
||||
@@ -1487,67 +1460,29 @@ fn plan_declarations(
|
||||
let mut schema = schema;
|
||||
let mut fields = Vec::with_capacity(columns.len());
|
||||
|
||||
for declaration in columns {
|
||||
let name = declaration.name();
|
||||
for (name, expression) in columns {
|
||||
if schema.field_with_name(name).is_ok() {
|
||||
return Err(Error::ColumnAlreadyExists {
|
||||
name: name.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let bound = bind(schema.clone(), name, declaration.expression())?;
|
||||
let bound = bind(schema.clone(), name, expression)?;
|
||||
|
||||
// Declared columns start entirely null, so nullability is a property
|
||||
// of the declaration rather than of what the expression yields.
|
||||
let computed_metadata = computed_column_metadata(declaration.expression(), &bound.inputs);
|
||||
let field = match declaration.field() {
|
||||
None => ArrowField::new(name, bound.data_type, true).with_metadata(computed_metadata),
|
||||
Some(field) => {
|
||||
ensure_no_foreign_declaration(field)?;
|
||||
if !field.is_nullable() {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"explicit computed output field '{}' must be nullable",
|
||||
field.name()
|
||||
),
|
||||
});
|
||||
}
|
||||
let expression_type = if field.is_blob_v2() {
|
||||
match field.data_type() {
|
||||
DataType::Struct(fields)
|
||||
if BlobV2Layout::classify(fields) == Some(BlobV2Layout::Logical) =>
|
||||
{
|
||||
// A Blob field stores descriptors, but refresh
|
||||
// publishes materialized bytes through Lance's
|
||||
// Blob conversion path.
|
||||
DataType::LargeBinary
|
||||
}
|
||||
data_type => {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"explicit Blob v2 output field '{}' has unsupported logical type {}",
|
||||
field.name(),
|
||||
data_type
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
} else {
|
||||
field.data_type().clone()
|
||||
};
|
||||
if expression_type != bound.data_type {
|
||||
return Err(Error::InvalidExpression {
|
||||
column: name.to_string(),
|
||||
message: format!(
|
||||
"explicit output field accepts {}, but the expression yields {}",
|
||||
expression_type, bound.data_type,
|
||||
),
|
||||
});
|
||||
}
|
||||
let mut metadata = field.metadata().clone();
|
||||
let computed_metadata = computed_column_metadata(expression, &bound.inputs);
|
||||
let field = match bound.projected_blob_field {
|
||||
Some(source) => {
|
||||
let mut metadata = source.metadata().clone();
|
||||
metadata.retain(|key, _| !is_declaration_key(key));
|
||||
metadata.extend(computed_metadata);
|
||||
field.clone().with_metadata(metadata)
|
||||
source
|
||||
.with_name(name)
|
||||
.with_nullable(true)
|
||||
.with_metadata(metadata)
|
||||
}
|
||||
None => ArrowField::new(name, bound.data_type, true).with_metadata(computed_metadata),
|
||||
};
|
||||
schema = Arc::new(ArrowSchema::new_with_metadata(
|
||||
schema
|
||||
@@ -1565,13 +1500,7 @@ fn plan_declarations(
|
||||
}
|
||||
|
||||
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)
|
||||
plan_declarations(schema, columns)
|
||||
}
|
||||
|
||||
/// Run the schema-level checks of
|
||||
@@ -1610,7 +1539,7 @@ pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) ->
|
||||
/// public way in.
|
||||
pub(crate) fn declare(
|
||||
schema: SchemaRef,
|
||||
columns: &[ComputedColumnDeclaration],
|
||||
columns: &[(String, String)],
|
||||
) -> Result<NewColumnTransform> {
|
||||
let fields = plan_declarations(schema, columns)?;
|
||||
Ok(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(
|
||||
@@ -1738,89 +1667,42 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_explicit_field_is_the_output_schema() {
|
||||
let table = table_with_ints("explicit_field").await;
|
||||
let field = ArrowField::new("copy", DataType::Int32, true)
|
||||
.with_metadata(HashMap::from([("semantic".into(), "custom".into())]));
|
||||
#[test]
|
||||
fn test_direct_blob_projection_inherits_semantics() {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![crate::blob("image", false)]));
|
||||
let fields = plan(
|
||||
schema,
|
||||
&[
|
||||
("first".to_string(), "image".to_string()),
|
||||
("second".to_string(), "first".to_string()),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
table
|
||||
.add_columns()
|
||||
.computed_field(field, "x")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let schema = table.schema().await.unwrap();
|
||||
let field = schema.field_with_name("copy").unwrap();
|
||||
assert_eq!(field.data_type(), &DataType::Int32);
|
||||
for field in &fields {
|
||||
assert!(field.is_blob_v2());
|
||||
assert!(field.is_nullable());
|
||||
}
|
||||
assert_eq!(
|
||||
field.metadata().get("semantic").map(String::as_str),
|
||||
Some("custom")
|
||||
);
|
||||
assert_eq!(
|
||||
field
|
||||
fields[1]
|
||||
.metadata()
|
||||
.get(COMPUTED_COLUMN_META_KEY)
|
||||
.get(EXPRESSION_META_KEY)
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
Some("first")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_explicit_field_must_be_nullable() {
|
||||
let table = table_with_ints("explicit_field_nullable").await;
|
||||
let error = table
|
||||
.add_columns()
|
||||
.computed_field(ArrowField::new("copy", DataType::Int32, false), "x")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
#[test]
|
||||
fn test_blob_expression_transformation_does_not_inherit_semantics() {
|
||||
let schema = Arc::new(ArrowSchema::new(vec![crate::blob("image", true)]));
|
||||
let fields = plan(
|
||||
schema,
|
||||
&[("payload".to_string(), "coalesce(image, image)".to_string())],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(error, Error::InvalidInput { message } if message.contains("must be nullable"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_explicit_field_cannot_supply_binding_metadata() {
|
||||
let table = table_with_ints("explicit_field_binding").await;
|
||||
let field =
|
||||
ArrowField::new("copy", DataType::Int32, true).with_metadata(HashMap::from([(
|
||||
EXPRESSION_META_KEY.to_string(),
|
||||
"other".to_string(),
|
||||
)]));
|
||||
let error = table
|
||||
.add_columns()
|
||||
.computed_field(field, "x")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::InvalidInput { message }
|
||||
if message.contains("carries computed-column metadata")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_explicit_blob_field_requires_a_logical_blob_layout() {
|
||||
let table = table_with_ints("explicit_blob_layout").await;
|
||||
let metadata = crate::blob("copy", true).metadata().clone();
|
||||
let field = ArrowField::new("copy", DataType::LargeBinary, true).with_metadata(metadata);
|
||||
let error = table
|
||||
.add_columns()
|
||||
.computed_field(field, "x")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::InvalidInput { message }
|
||||
if message.contains("unsupported logical type LargeBinary")
|
||||
));
|
||||
assert!(!fields[0].is_blob_v2());
|
||||
assert_eq!(fields[0].data_type(), &DataType::LargeBinary);
|
||||
}
|
||||
|
||||
/// The binding reaches the schema only if `AllNulls` carries per-field
|
||||
|
||||
@@ -1338,51 +1338,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[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() {
|
||||
async fn test_refresh_inherits_and_publishes_blob_output() {
|
||||
use arrow_array::UInt64Array;
|
||||
use lance_arrow::{
|
||||
BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, BLOB_INLINE_SIZE_THRESHOLD_META_KEY,
|
||||
@@ -1402,7 +1358,7 @@ mod tests {
|
||||
.await;
|
||||
table
|
||||
.add_columns()
|
||||
.computed_field(crate::blob("image_copy", true), "image")
|
||||
.computed("image_copy", "image")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -1515,28 +1471,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[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_field(crate::blob("invalid", true), "id + 1")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::InvalidExpression { column, message }
|
||||
if column == "invalid"
|
||||
&& message.contains("explicit output field accepts LargeBinary")
|
||||
&& message.contains("expression yields Int32")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_materializes_nested_struct_blob_input() {
|
||||
use arrow_array::{Int32Array, StructArray};
|
||||
async fn test_refresh_inherits_nested_struct_blob_input() {
|
||||
use arrow_array::{Int32Array, StructArray, UInt64Array};
|
||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
@@ -1574,21 +1510,27 @@ mod tests {
|
||||
.rows_filled,
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
table.blob_columns().await.unwrap(),
|
||||
vec!["metadata.image".to_string(), "payload_copy".to_string()]
|
||||
);
|
||||
let batches = table
|
||||
.query()
|
||||
.select(Select::columns(&["payload_copy"]))
|
||||
.with_row_id()
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let payloads = batches[0]
|
||||
.column_by_name("payload_copy")
|
||||
let row_ids = batches[0]
|
||||
.column_by_name(ROW_ID)
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<LargeBinaryArray>()
|
||||
.unwrap();
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap()
|
||||
.values();
|
||||
let payloads = table.fetch_blobs("payload_copy", row_ids).await.unwrap();
|
||||
assert_eq!(payloads.value(0), b"nested");
|
||||
assert!(payloads.is_null(1));
|
||||
}
|
||||
@@ -1655,8 +1597,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_materializes_external_blob_input() {
|
||||
use arrow_array::{Int32Array, StringArray};
|
||||
async fn test_refresh_inherits_external_blob_input() {
|
||||
use arrow_array::{Int32Array, StringArray, UInt64Array};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
@@ -1713,19 +1655,21 @@ mod tests {
|
||||
);
|
||||
let batches = table
|
||||
.query()
|
||||
.select(Select::columns(&["payload_copy"]))
|
||||
.with_row_id()
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let payloads = batches[0]
|
||||
.column_by_name("payload_copy")
|
||||
let row_ids = batches[0]
|
||||
.column_by_name(ROW_ID)
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<LargeBinaryArray>()
|
||||
.unwrap();
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap()
|
||||
.values();
|
||||
let payloads = table.fetch_blobs("payload_copy", row_ids).await.unwrap();
|
||||
assert_eq!(payloads.value(0), payload);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -124,7 +124,7 @@ pub(crate) async fn execute_add_columns(
|
||||
/// declaration metadata.
|
||||
pub(crate) async fn execute_declare(
|
||||
table: &NativeTable,
|
||||
columns: &[computed_columns::ComputedColumnDeclaration],
|
||||
columns: &[(String, String)],
|
||||
) -> Result<AddColumnsResult> {
|
||||
use lance::dataset::mem_wal::DatasetMemWalExt;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user