refactor: carry explicit computed output fields

This commit is contained in:
Xuanwo
2026-08-28 15:43:07 +08:00
parent b40e005c17
commit 2702cdb219
11 changed files with 399 additions and 175 deletions
+1 -1
View File
@@ -38,7 +38,7 @@ from .materialized_view import (
MaterializedView,
MaterializedViewDefinition,
)
from .table import AsyncTable, ComputedColumn as ComputedColumn, Table
from .table import AsyncTable, Table
from .types import BaseTokenizerType
from ._lancedb import Session
from .namespace import (
+1 -1
View File
@@ -353,7 +353,7 @@ class Table:
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
async def add_computed_columns(
self,
columns: list[tuple[str, str, Literal["inferred", "blob_v2"]]],
columns: list[tuple[str | pa.Field, str]],
) -> AddColumnsResult: ...
async def add_function_columns(
self, application_json: str, output_name: Optional[str]
+2 -2
View File
@@ -13,6 +13,7 @@ from typing import (
Iterable,
List,
Optional,
Sequence,
Union,
Literal,
overload,
@@ -71,7 +72,6 @@ from ..table import (
AsyncTable,
BlobMode,
Branches,
ComputedColumn,
IndexStatistics,
Query,
Table,
@@ -984,7 +984,7 @@ class RemoteTable(Table):
| FunctionApplication
| None = None,
*,
computed: Dict[str, str | ComputedColumn] | None = None,
computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms, computed=computed))
+34 -50
View File
@@ -913,44 +913,23 @@ def _normalize_progress(progress):
return progress, False
@dataclass(frozen=True)
class ComputedColumn:
"""A computed-column expression with explicit output semantics.
Plain strings passed to ``Table.add_columns(computed=...)`` infer their
output type. Use :meth:`blob` when a ``LargeBinary`` expression should be
published as Blob v2.
"""
expression: str
output: Literal["inferred", "blob_v2"] = "inferred"
def __post_init__(self):
if not isinstance(self.expression, str):
raise TypeError("ComputedColumn.expression must be a string")
if self.output not in ("inferred", "blob_v2"):
raise ValueError("ComputedColumn.output must be 'inferred' or 'blob_v2'")
@classmethod
def blob(cls, expression: str) -> ComputedColumn:
"""Publish a ``LargeBinary`` expression result as Blob v2."""
return cls(expression=expression, output="blob_v2")
def _normalize_computed_columns(
computed: Dict[str, str | ComputedColumn],
) -> list[tuple[str, str, Literal["inferred", "blob_v2"]]]:
columns: list[tuple[str, str, Literal["inferred", "blob_v2"]]] = []
for name, declaration in computed.items():
if isinstance(declaration, str):
columns.append((name, declaration, "inferred"))
elif isinstance(declaration, ComputedColumn):
columns.append((name, declaration.expression, declaration.output))
else:
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 values must be SQL expression strings or "
"ComputedColumn values"
"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
@@ -2183,7 +2162,7 @@ class Table(ABC):
| pa.Schema
| None = None,
*,
computed: Dict[str, str | ComputedColumn] | None = None,
computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
):
"""
Add new columns with defined values.
@@ -2205,13 +2184,15 @@ class Table(ABC):
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str | ComputedColumn], optional
A map of column name to a SQL expression defining the column. The
column's type and inputs are derived from the expression, so no
data type is supplied. Use ``ComputedColumn.blob(expression)``
when a ``LargeBinary`` result should be stored as Blob v2. All
entries share mapping insertion order, including dependencies
between inferred and Blob outputs.
computed: Dict[str, str] or Sequence[Tuple[str | pa.Field, 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.
Unlike ``transforms``, the expression is stored rather than
evaluated now: the column is committed with no values, and rows get
@@ -4343,7 +4324,7 @@ class LanceTable(Table):
| pa.Schema
| None = None,
*,
computed: Dict[str, str | ComputedColumn] | None = None,
computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms, computed=computed))
@@ -6291,7 +6272,7 @@ class AsyncTable:
| pa.Schema
| None = None,
*,
computed: dict[str, str | ComputedColumn] | None = None,
computed: dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
) -> AddColumnsResult:
"""
Add new columns with defined values.
@@ -6311,11 +6292,14 @@ class AsyncTable:
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str | ComputedColumn], optional
A map of column name to a SQL expression defining the column. The
column's type and inputs are derived from the expression. Use
``ComputedColumn.blob(expression)`` to publish a ``LargeBinary``
result as Blob v2. Mapping insertion order is the declaration and
computed: Dict[str, str] or Sequence[Tuple[str | pa.Field, 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.
Unlike ``transforms``, the expression is stored rather than
+16 -12
View File
@@ -4100,10 +4100,10 @@ def test_computed_column_blob_input_and_explicit_output(tmp_path):
)
table.add_columns(
computed={
"image_copy": lancedb.ComputedColumn.blob("image"),
"payload_copy": "image_copy",
}
computed=[
(lancedb.blob("image_copy"), "image"),
("payload_copy", "image_copy"),
]
)
assert table.refresh_column("image_copy").rows_filled == 2
assert table.refresh_column("payload_copy").rows_filled == 2
@@ -4124,18 +4124,22 @@ def test_blob_output_declaration_rejects_eager_transforms(tmp_path):
with pytest.raises(ValueError):
table.add_columns(
{"a": "x + 1"},
computed={"b": lancedb.ComputedColumn.blob("x")},
computed=[(lancedb.blob("b"), "x")],
)
def test_computed_column_validates_explicit_output():
assert lancedb.ComputedColumn("x + 1").output == "inferred"
assert lancedb.ComputedColumn.blob("image").output == "blob_v2"
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="expression must be a string"):
lancedb.ComputedColumn(42) # type: ignore[arg-type]
with pytest.raises(ValueError, match="output must be"):
lancedb.ComputedColumn("x", output="binary") # type: ignore[arg-type]
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
+15 -10
View File
@@ -13,7 +13,7 @@ use crate::{
};
use arrow::{
array::{Array, LargeBinaryArray},
datatypes::{DataType, Schema},
datatypes::{DataType, Field, Schema},
ffi_stream::ArrowArrayStreamReader,
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
};
@@ -101,6 +101,12 @@ 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)]
@@ -1578,19 +1584,18 @@ impl Table {
pub fn add_computed_columns(
self_: PyRef<'_, Self>,
columns: Vec<(String, String, String)>,
columns: Vec<(ComputedColumnFieldArg, String)>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.add_columns();
for (name, expression, output) in columns {
let declaration = match output.as_str() {
"inferred" => ComputedColumnDeclaration::inferred(name, expression),
"blob_v2" => ComputedColumnDeclaration::blob(name, expression),
output => {
return Err(PyValueError::new_err(format!(
"unsupported computed-column output '{output}'"
)));
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);
+51 -24
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, SchemaRef};
use arrow_schema::{ArrowError, DataType, Schema as ArrowSchema, SchemaRef};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use datafusion_common::DataFusionError;
@@ -3183,21 +3183,31 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
self.schema().await?.as_ref(),
"schema evolution",
)?;
// The server plans the declaration: expression validation, type
// inference and the persisted binding all happen there.
let entries = columns
.iter()
.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<_>>();
// 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<_>>>()?;
let mut body = serde_json::json!({ "new_columns": entries });
self.apply_branch_body(&mut body);
let request = self
@@ -7394,7 +7404,7 @@ mod tests {
assert_eq!(result.version, if old_server { 0 } else { 43 });
}
/// A declaration is sent as `{name, computed}` entries for the server to
/// An inferred 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() {
@@ -7426,7 +7436,7 @@ mod tests {
}
#[tokio::test]
async fn test_add_blob_computed_column_sends_explicit_output_semantics() {
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/" => {
@@ -7435,9 +7445,29 @@ mod tests {
assert_eq!(
value["new_columns"],
serde_json::json!([{
"name": "image_copy",
"computed": "image",
"computed_output": "blob_v2"
"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()
@@ -7450,10 +7480,7 @@ mod tests {
let result = table
.add_columns()
.computed_column(crate::table::ComputedColumnDeclaration::blob(
"image_copy",
"image",
))
.computed_field(crate::blob("image_copy", true), "image")
.execute()
.await
.unwrap();
+4 -4
View File
@@ -95,8 +95,8 @@ pub use cherry_pick::{
};
pub use chrono::Duration;
pub use computed_columns::{
ComputedColumn, ComputedColumnDeclaration, ComputedColumnKind, ComputedColumnOutput,
computed_column_from_field, computed_columns,
ComputedColumn, ComputedColumnDeclaration, ComputedColumnKind, computed_column_from_field,
computed_columns,
};
pub use delete::DeleteResult;
use futures::future::join_all;
@@ -751,8 +751,8 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
/// Declare computed columns, each defined by a SQL expression.
///
/// Where the declaration is planned depends on the backend: a local table
/// validates and types the expression itself, a remote one sends the text
/// for the server to plan.
/// validates and types the expression itself, while a remote one sends the
/// expression and any explicit output field for the server to plan.
async fn add_computed_columns(
&self,
_columns: &[computed_columns::ComputedColumnDeclaration],
+13 -7
View File
@@ -5,6 +5,7 @@
use std::sync::Arc;
use arrow_schema::Field as ArrowField;
use lance::dataset::NewColumnTransform;
use super::BaseTable;
@@ -87,26 +88,31 @@ impl AddColumnsBuilder {
self.computed_column(ComputedColumnDeclaration::inferred(name, expression))
}
/// Add one computed-column declaration with explicit output semantics.
/// Add a computed column whose output schema is the supplied Arrow field.
///
/// Use [`ComputedColumnDeclaration::blob`] when a `LargeBinary` expression
/// should be published as Blob v2. Declarations share one ordered stream,
/// so a later expression may reference a column declared earlier in the
/// same builder.
/// 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;
/// # use lancedb::table::ComputedColumnDeclaration;
/// # async fn declare(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
/// table
/// .add_columns()
/// .computed_column(ComputedColumnDeclaration::blob("image_copy", "image"))
/// .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);
self
+193 -54
View File
@@ -9,12 +9,13 @@
//! refresh fills the rows.
//!
//! The rule is tagged by kind ([`ComputedColumnKind`]) because kinds differ in
//! where the column's type and inputs come from. A SQL expression is
//! self-describing -- both are derived from the expression, so a caller writes
//! neither -- while 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.
//! 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.
//! Registered Functions use an exact remote version plus a schema-level
//! Function binding; unknown newer kinds remain readable and fail closed
//! before mutation.
//!
//! [`computed_columns`] and [`computed_column_from_field`] read declarations
//! back off a schema.
@@ -27,7 +28,9 @@ 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_core::datatypes::{
BLOB_V2_DESC_FIELD, BlobV2Layout, format_field_path_minimal, parse_field_path,
};
use lance_datafusion::planner::Planner;
use lance_namespace::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema};
use serde::{Deserialize, Serialize};
@@ -66,43 +69,59 @@ 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,
#[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 {
/// 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,
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 {
name: name.into(),
target: ComputedColumnTarget::Inferred(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 {
/// 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 {
name: name.into(),
target: ComputedColumnTarget::Explicit(field),
expression: expression.into(),
output: ComputedColumnOutput::BlobV2,
}
}
/// 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),
}
}
}
@@ -1149,15 +1168,20 @@ pub(crate) fn ensure_no_foreign_declarations<'a>(
fields: impl IntoIterator<Item = &'a Arc<ArrowField>>,
) -> Result<()> {
for field in fields {
if field.metadata().keys().any(|k| is_declaration_key(k)) {
return Err(Error::InvalidInput {
message: format!(
"field '{}' carries computed-column metadata; declare computed columns \
with add_columns().computed()",
field.name()
),
});
}
ensure_no_foreign_declaration(field)?;
}
Ok(())
}
fn ensure_no_foreign_declaration(field: &ArrowField) -> Result<()> {
if field.metadata().keys().any(|k| is_declaration_key(k)) {
return Err(Error::InvalidInput {
message: format!(
"field '{}' carries computed-column metadata; declare computed columns \
with add_columns().computed()",
field.name()
),
});
}
Ok(())
}
@@ -1464,35 +1488,65 @@ fn plan_declarations(
let mut fields = Vec::with_capacity(columns.len());
for declaration in columns {
if schema.field_with_name(&declaration.name).is_ok() {
let name = declaration.name();
if schema.field_with_name(name).is_ok() {
return Err(Error::ColumnAlreadyExists {
name: declaration.name.clone(),
name: name.to_string(),
});
}
let bound = bind(schema.clone(), &declaration.name, &declaration.expression)?;
let bound = bind(schema.clone(), name, declaration.expression())?;
// Declared columns start entirely null, so nullability is a property
// of the declaration rather than of what the expression yields.
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(),
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!(
"a Blob v2 computed output requires a LargeBinary expression, got {}",
bound.data_type
"explicit computed output field '{}' must be nullable",
field.name()
),
});
}
let field = crate::blob::blob(&declaration.name, true);
let mut blob_metadata = field.metadata().clone();
blob_metadata.extend(metadata);
field.with_metadata(blob_metadata)
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();
metadata.extend(computed_metadata);
field.clone().with_metadata(metadata)
}
};
schema = Arc::new(ArrowSchema::new_with_metadata(
@@ -1684,6 +1738,91 @@ 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())]));
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);
assert_eq!(
field.metadata().get("semantic").map(String::as_str),
Some("custom")
);
assert_eq!(
field
.metadata()
.get(COMPUTED_COLUMN_META_KEY)
.map(String::as_str),
Some("true")
);
}
#[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();
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")
));
}
/// The binding reaches the schema only if `AllNulls` carries per-field
/// metadata through the commit. The whole representation rests on it.
#[tokio::test]
+69 -10
View File
@@ -33,9 +33,10 @@ use std::collections::HashSet;
use std::sync::Arc;
use arrow_array::{
Array, ArrayRef, BooleanArray, LargeBinaryArray, RecordBatch, RecordBatchOptions,
Array, ArrayRef, BooleanArray, LargeBinaryArray, RecordBatch, RecordBatchOptions, StructArray,
new_null_array,
};
use arrow_schema::Schema as ArrowSchema;
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
use datafusion_expr::ColumnarValue;
use futures::{Stream, StreamExt, TryStreamExt};
use lance::Dataset;
@@ -374,7 +375,10 @@ fn configure_blob_inputs(
Ok(())
}
fn blob_array_from_binary(array: &ArrayRef) -> lance_core::Result<ArrayRef> {
fn blob_array_from_binary(
array: &ArrayRef,
target_field: &ArrowField,
) -> lance_core::Result<ArrayRef> {
let values = array
.as_any()
.downcast_ref::<LargeBinaryArray>()
@@ -392,7 +396,39 @@ fn blob_array_from_binary(array: &ArrayRef) -> lance_core::Result<ArrayRef> {
builder.push_bytes(values.value(index))?;
}
}
builder.finish()
let minimal = builder.finish()?;
let minimal = minimal
.as_any()
.downcast_ref::<StructArray>()
.ok_or_else(|| lance_core::Error::internal("Blob builder returned a non-struct array"))?;
let DataType::Struct(target_fields) = target_field.data_type() else {
return Err(lance_core::Error::invalid_input(format!(
"Blob v2 output field '{}' has non-struct type {}",
target_field.name(),
target_field.data_type()
)));
};
let columns = target_fields
.iter()
.map(|field| match field.name().as_str() {
"data" | "uri" => minimal
.column_by_name(field.name())
.cloned()
.ok_or_else(|| {
lance_core::Error::internal(format!("Blob builder omitted '{}'", field.name()))
}),
"position" | "size" => Ok(new_null_array(field.data_type(), minimal.len())),
name => Err(lance_core::Error::invalid_input(format!(
"Blob v2 output field '{}' has unsupported logical child '{name}'",
target_field.name()
))),
})
.collect::<lance_core::Result<Vec<_>>>()?;
Ok(Arc::new(StructArray::try_new(
target_fields.clone(),
columns,
minimal.nulls().cloned(),
)?))
}
/// How many rows of one fragment would gain a value.
@@ -495,7 +531,7 @@ 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)?
blob_array_from_binary(&merged, projected.field(0))?
} else {
merged
};
@@ -507,13 +543,15 @@ async fn fill_stream(
mod tests {
use std::sync::Arc;
use arrow_array::{Array, Int32Array, LargeBinaryArray, RecordBatch, record_batch};
use arrow_array::{
Array, ArrayRef, Int32Array, LargeBinaryArray, RecordBatch, StructArray, record_batch,
};
use arrow_schema::Field as ArrowField;
use futures::TryStreamExt;
use lance_core::ROW_ID;
use crate::connect;
use crate::query::{ExecutableQuery, QueryBase, Select};
use crate::table::ComputedColumnDeclaration;
use crate::{Error, Result, Table};
async fn table_with(name: &str, values: Vec<i32>) -> Table {
@@ -566,6 +604,25 @@ mod tests {
table.add(batch).execute().await.unwrap();
}
#[test]
fn test_blob_output_matches_complete_logical_field() {
let values: ArrayRef = Arc::new(LargeBinaryArray::from(vec![
Some(b"hello".as_slice()),
None,
]));
let field = ArrowField::new(
"image",
lance_core::datatypes::BLOB_V2_LOGICAL_TYPE.clone(),
true,
);
let output = super::blob_array_from_binary(&values, &field).unwrap();
assert_eq!(output.data_type(), field.data_type());
let output = output.as_any().downcast_ref::<StructArray>().unwrap();
assert_eq!(output.column_by_name("position").unwrap().null_count(), 2);
assert_eq!(output.column_by_name("size").unwrap().null_count(), 2);
}
/// The gate's reproducer: `b = coalesce(a, 0)` refreshed before `a`
/// must not bake zeros from `a`'s placeholder null. It is refused, and
/// names the input, until `a` is filled -- after every append too.
@@ -1345,7 +1402,7 @@ mod tests {
.await;
table
.add_columns()
.computed_column(ComputedColumnDeclaration::blob("image_copy", "image"))
.computed_field(crate::blob("image_copy", true), "image")
.execute()
.await
.unwrap();
@@ -1464,14 +1521,16 @@ mod tests {
let error = table
.add_columns()
.computed_column(ComputedColumnDeclaration::blob("invalid", "id + 1"))
.computed_field(crate::blob("invalid", true), "id + 1")
.execute()
.await
.unwrap_err();
assert!(matches!(
error,
Error::InvalidExpression { column, message }
if column == "invalid" && message.contains("requires a LargeBinary expression")
if column == "invalid"
&& message.contains("explicit output field accepts LargeBinary")
&& message.contains("expression yields Int32")
));
}