refactor: inherit blob metadata for computed projections

This commit is contained in:
Xuanwo
2026-08-28 16:44:17 +08:00
parent 2702cdb219
commit 6c8e5e9d98
11 changed files with 158 additions and 512 deletions
+1 -2
View File
@@ -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]
+1 -2
View File
@@ -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))
+14 -42
View File
@@ -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):
+5 -37
View File
@@ -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
View File
@@ -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))
+16 -84
View File
@@ -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| {
+4 -8
View File
@@ -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)
+3 -34
View File
@@ -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
}
+83 -201
View File
@@ -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
+24 -80
View File
@@ -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);
}
}
+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: &[computed_columns::ComputedColumnDeclaration],
columns: &[(String, String)],
) -> Result<AddColumnsResult> {
use lance::dataset::mem_wal::DatasetMemWalExt;