mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-03 12:08:52 +00:00
refactor: carry explicit computed output fields
This commit is contained in:
@@ -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 (
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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);
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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")
|
||||
));
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user