feat(python): add first-class function call authoring

This commit is contained in:
Xuanwo
2026-08-12 09:19:49 +08:00
parent f65bf89c98
commit 47213e31f8
5 changed files with 831 additions and 4 deletions
+6
View File
@@ -235,6 +235,12 @@ class Function:
def output_type(self) -> pa.DataType: ...
@property
def output_nullable(self) -> bool: ...
def __call__(self, **kwargs: Any) -> "_FunctionCall": ...
class _FunctionCall:
"""Private unresolved Function call authoring value (FF-028)."""
...
class _FunctionDefinition:
"""Private owner of the Rust FunctionDefinition registration input."""
@@ -0,0 +1,372 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Contract tests for Python exact Function handle call authoring (FF-028)."""
from __future__ import annotations
import contextlib
import http.server
import json
import threading
from collections.abc import Iterator
from typing import Any, Callable
import pyarrow as pa
import pytest
import lancedb
from lancedb import _lancedb as _native
from lancedb.expr import Expr, col, func, lit
_CALL_PATH = "/v1/functions/lookup"
_CALL_CATALOG_NAME = "text.normalize.call-name"
_CALL_FUNCTION_ID = "fn.exact.call-handle"
_LITERAL_PAYLOAD_SENTINEL = "LITERAL_PAYLOAD_SENTINEL_call_xyz_42"
_INT_PAYLOAD_SENTINEL = 2_147_000_123
# Pinned Rust-canonical schema-only type IPC (base64).
_INT32_TYPE_IPC_B64 = (
"QVJST1cxAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP"
"////94AAAAEAAAAAAACgAMAAoACQAEAAoAAAAQAAAAAAEEAAgACAAAAAQACAAAAAQAAAABAAAAFAAAABAAFAAQ"
"AA4ADwAEAAAACAAQAAAAGAAAACAAAAAAAAECHAAAAAgADAAEAAsACAAAACAAAAAAAAABAAAAAAAAAAAAAAAA/"
"////wAAAAAUAAAAAAAAAAwAFAASAAwACAAEAAwAAABsAAAAcAAAABAAAAAAAAQACAAIAAAABAAIAAAABAAAAA"
"EAAAAUAAAAEAAUABAADgAPAAQAAAAIABAAAAAYAAAAIAAAAAAAAQIcAAAACAAMAAQACwAIAAAAIAAAAAAAAAE"
"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACQAAAAQVJST1cx"
)
_UTF8_TYPE_IPC_B64 = (
"QVJST1cxAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP"
"////94AAAAEAAAAAAACgAMAAoACQAEAAoAAAAQAAAAAAEEAAgACAAAAAQACAAAAAQAAAABAAAAFAAAABAAFAAQ"
"AA4ADwAEAAAACAAQAAAAGAAAAAwAAAAAAAEFEAAAAAAAAAAEAAQABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/"
"////wAAAAAQAAAADAAUABIADAAIAAQADAAAAGAAAABkAAAAEAAAAAAABAAIAAgAAAAEAAgAAAAEAAAAAQAAAB"
"QAAAAQABQAEAAOAA8ABAAAAAgAEAAAABgAAAAMAAAAAAABBRAAAAAAAAAABAAEAAQAAAAAAAAAAAAAAAAAAAA"
"AAAAAAAAAAIAAAABBUlJPVzE="
)
_LIST_INT32_TYPE_IPC_B64 = (
"QVJST1cxAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP"
"////+4AAAAEAAAAAAACgAMAAoACQAEAAoAAAAQAAAAAAEEAAgACAAAAAQACAAAAAQAAAABAAAABAAAANz///8c"
"AAAADAAAAAAAAQxcAAAAAQAAABwAAAAEAAQABAAAABAAFAAQAA4ADwAEAAAACAAQAAAAGAAAACAAAAAAAAECH"
"AAAAAgADAAEAAsACAAAACAAAAAAAAABAAAAAAQAAABpdGVtAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP"
"////8AAAAAFAAAAAAAAAAMABQAEgAMAAgABAAMAAAAnAAAAKAAAAAQAAAAAAAEAAgACAAAAAQACAAAAAQAAAA"
"BAAAABAAAANz///8cAAAADAAAAAAAAQxcAAAAAQAAABwAAAAEAAQABAAAABAAFAAQAA4ADwAEAAAACAAQAAAA"
"GAAAACAAAAAAAAECHAAAAAgADAAEAAsACAAAACAAAAAAAAABAAAAAAQAAABpdGVtAAAAAAAAAAAAAAAAAAAA"
"AAAAAAAAAAAAwAAAAEFSUk9XMQ=="
)
_OVERDESIGN_ATTRS = (
"id",
"function_id",
"name",
"connection",
"table",
"snapshot",
"field_id",
"field_ids",
"job",
"job_id",
"artifact",
"digest",
"retry_key",
"idempotency_key",
"user_version",
"execute",
"status",
"wait",
"cancel",
"to_json",
"_to_json",
"serialize",
"geneva",
)
def _sample_function_wire(
*,
function_id: str = _CALL_FUNCTION_ID,
parameters: list[dict[str, str]] | None = None,
output_type_ipc: str = _UTF8_TYPE_IPC_B64,
) -> dict[str, Any]:
return {
"format_version": 1,
"id": function_id,
"signature": {
"parameters": parameters
or [
{"name": "text", "data_type_ipc": _UTF8_TYPE_IPC_B64},
{"name": "limit", "data_type_ipc": _INT32_TYPE_IPC_B64},
],
"output": {
"data_type_ipc": output_type_ipc,
"nullable": True,
},
},
}
def _lookup_success_body(function: dict[str, Any] | None = None) -> bytes:
return json.dumps({"function": function or _sample_function_wire()}).encode("utf-8")
def _read_body(request: http.server.BaseHTTPRequestHandler) -> bytes:
content_len = int(request.headers.get("Content-Length", 0))
if content_len <= 0:
return b""
return request.rfile.read(content_len)
def _make_handler(handler: Callable[[http.server.BaseHTTPRequestHandler], None]):
class _Handler(http.server.BaseHTTPRequestHandler):
def do_GET(self):
handler(self)
def do_POST(self):
handler(self)
def log_message(self, format, *args): # noqa: A003
return
return _Handler
@contextlib.contextmanager
def _mock_remote_db(handler) -> Iterator[Any]:
server = http.server.HTTPServer(("localhost", 0), _make_handler(handler))
port = server.server_address[1]
thread = threading.Thread(target=server.serve_forever)
thread.start()
try:
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=f"http://localhost:{port}",
client_config={
"retry_config": {
"retries": 2,
"backoff_factor": 0.0,
"backoff_jitter": 0.0,
},
"timeout_config": {"connect_timeout": 1},
},
)
yield db
finally:
server.shutdown()
thread.join()
def _lookup_function(function: dict[str, Any] | None = None):
body = _lookup_success_body(function)
def handler(request: http.server.BaseHTTPRequestHandler) -> None:
assert request.command == "POST"
assert request.path == _CALL_PATH
_read_body(request)
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(body)
with _mock_remote_db(handler) as db:
return db.functions.get(_CALL_CATALOG_NAME)
def _authored_call_type():
cls = getattr(_native, "_FunctionCall", None)
if cls is None:
pytest.fail("lancedb._lancedb._FunctionCall is missing")
return cls
def _exception_text(exc: BaseException) -> str:
parts = [str(exc), repr(exc)]
current: BaseException | None = exc
seen: set[int] = set()
while current is not None and id(current) not in seen:
seen.add(id(current))
parts.append(f"{type(current).__name__}: {current}")
current = current.__cause__ or current.__context__
return "\n".join(parts)
def test_function_keyword_call_returns_private_frozen_authored_value():
function = _lookup_function()
assert callable(function)
authored = function(text=col("text"), limit=8)
authored_type = _authored_call_type()
assert type(authored) is authored_type
assert authored_type.__module__ == "lancedb._lancedb"
assert authored_type.__name__ == "_FunctionCall"
# Keyword order must not matter; bindings store/render in signature order.
authored_reversed = function(limit=8, text=col("text"))
assert type(authored_reversed) is authored_type
rendered = repr(authored_reversed)
assert rendered.index("text=") < rendered.index("limit=")
assert 'text=field("text")' in rendered
assert "limit=literal(Int32, null=false)" in rendered
def test_function_call_rejects_positional_missing_and_unknown_args():
function = _lookup_function()
with pytest.raises(TypeError, match="keyword"):
function(col("text"), 8)
with pytest.raises((TypeError, ValueError), match="limit"):
function(text=col("text"))
with pytest.raises((TypeError, ValueError), match="text"):
function(limit=8)
with pytest.raises((TypeError, ValueError), match="unknown|extra"):
function(text=col("text"), limit=8, extra=1)
def test_function_call_accepts_direct_case_sensitive_column_and_rejects_complex_exprs():
function = _lookup_function()
authored = function(text=col("firstName"), limit=1)
assert type(authored) is _authored_call_type()
rendered = repr(authored)
assert 'text=field("firstName")' in rendered
assert "limit=literal(Int32, null=false)" in rendered
complex_exprs = (
col("text") + lit("x"),
col("text").cast(pa.string()),
func("lower", col("text")),
col("text") == lit("x"),
col("text").lower(),
)
for expr in complex_exprs:
with pytest.raises((TypeError, ValueError)):
function(text=expr, limit=1)
# Raw native PyExpr is not the public col() wrapper.
with pytest.raises((TypeError, ValueError)):
function(text=col("text")._inner, limit=1)
# Non-expression / non-literal objects are rejected for field-shaped misuse
# when a column binding is required; plain strings are literals for utf8.
with pytest.raises((TypeError, ValueError)):
function(text=object(), limit=1)
def test_function_call_plain_literal_declared_type_null_and_nested():
function = _lookup_function()
authored = function(text="hello", limit=7)
assert type(authored) is _authored_call_type()
rendered = repr(authored)
assert "text=literal(Utf8, null=false)" in rendered
assert "limit=literal(Int32, null=false)" in rendered
# Plain Python int normalizes to declared Int32 and non-null.
authored_int32 = function(text="hello", limit=2_147_483_647)
assert type(authored_int32) is _authored_call_type()
rendered_int32 = repr(authored_int32)
assert "limit=literal(Int32, null=false)" in rendered_int32
assert "Int64" not in rendered_int32
assert "2147483647" not in rendered_int32
# Plain None keeps each declared parameter type with null=true.
authored_null = function(text=None, limit=None)
assert type(authored_null) is _authored_call_type()
rendered_null = repr(authored_null)
assert "text=literal(Utf8, null=true)" in rendered_null
assert "limit=literal(Int32, null=true)" in rendered_null
list_function = _lookup_function(
_sample_function_wire(
parameters=[
{"name": "values", "data_type_ipc": _LIST_INT32_TYPE_IPC_B64},
]
)
)
authored_list = list_function(values=[1, 2, 3])
assert type(authored_list) is _authored_call_type()
rendered_list = repr(authored_list)
assert "values=literal(List(Int32), null=false)" in rendered_list
assert "[1, 2, 3]" not in rendered_list
authored_list_null = list_function(values=None)
assert type(authored_list_null) is _authored_call_type()
rendered_list_null = repr(authored_list_null)
assert "values=literal(List(Int32), null=true)" in rendered_list_null
def test_function_call_direct_literal_expr_exact_type_only():
function = _lookup_function()
# lit(int) is Int64 in the expression builder; int32 parameter must reject it.
with pytest.raises((TypeError, ValueError), match="limit|int32|type") as raised:
function(text="hello", limit=lit(8))
reject_text = _exception_text(raised.value)
assert "Int64" in reject_text or "int64" in reject_text.lower()
assert "Int32" in reject_text or "int32" in reject_text.lower()
# Exact utf8 literal expression is accepted and stored as Utf8/non-null.
authored = function(text=lit("hello"), limit=8)
assert type(authored) is _authored_call_type()
rendered = repr(authored)
assert "text=literal(Utf8, null=false)" in rendered
assert "limit=literal(Int32, null=false)" in rendered
assert "hello" not in rendered
# Cast / arithmetic around a literal is not a direct Literal node.
with pytest.raises((TypeError, ValueError)):
function(text=lit("hello").cast(pa.string()), limit=8)
def test_function_call_conversion_error_and_repr_are_payload_free():
function = _lookup_function()
with pytest.raises((TypeError, ValueError)) as raised:
function(text="ok", limit=_LITERAL_PAYLOAD_SENTINEL)
text = _exception_text(raised.value)
assert _LITERAL_PAYLOAD_SENTINEL not in text
assert "limit" in text
assert "int32" in text.lower() or "Int32" in text
authored = function(text=_LITERAL_PAYLOAD_SENTINEL, limit=_INT_PAYLOAD_SENTINEL)
rendered = f"{authored!r}\n{authored!s}"
assert _LITERAL_PAYLOAD_SENTINEL not in rendered
assert str(_INT_PAYLOAD_SENTINEL) not in rendered
assert "text=literal(Utf8, null=false)" in rendered
assert "limit=literal(Int32, null=false)" in rendered
assert type(authored).__name__ == "_FunctionCall"
assert "_FunctionCall" in rendered
def test_function_call_private_type_nonconstructible_immutable_and_not_exported():
function = _lookup_function()
authored = function(text=col("text"), limit=1)
authored_type = _authored_call_type()
assert "_FunctionCall" not in getattr(lancedb, "__all__", [])
assert not hasattr(lancedb, "_FunctionCall")
assert getattr(_native, "_FunctionCall", None) is authored_type
with pytest.raises(TypeError):
authored_type()
for attr in _OVERDESIGN_ATTRS:
assert not hasattr(authored, attr)
for attr in ("function", "bindings", "arguments", "parameters", "text", "limit"):
with pytest.raises(AttributeError):
setattr(authored, attr, None)
# Existing Function handle stays frozen / connection-free / name-free.
assert not hasattr(function, "name")
assert not hasattr(function, "connection")
with pytest.raises(AttributeError):
function.id = "mutated"
def test_function_call_does_not_change_col_query_expression_behavior():
# Regression guard: authoring must not alter public col()/Expr query behavior.
expr = col("firstName") > lit(1)
assert isinstance(expr, Expr)
assert expr.to_sql() == "(`firstName` > 1)"
+28 -1
View File
@@ -10,7 +10,7 @@
use std::ops::{Add, Div, Mul, Not, Sub};
use arrow::{datatypes::DataType, pyarrow::PyArrowType};
use datafusion_common::ScalarValue;
use datafusion_common::{Column, ScalarValue};
use lancedb::expr::{
DfExpr, col as ldb_col, contains, expr_cast, is_in, lit as df_lit, lower, upper,
};
@@ -27,6 +27,33 @@ use pyo3::{Bound, PyAny, PyResult, exceptions::PyValueError, prelude::*, pyfunct
#[derive(Clone)]
pub struct PyExpr(pub DfExpr);
/// Crate-private inspection result for Function call authoring (FF-028).
#[derive(Debug, Clone)]
pub(crate) enum DirectExprView<'a> {
/// Direct unqualified DataFusion Column; name is case-sensitive.
UnqualifiedColumn(&'a str),
/// Direct Literal scalar; Arrow type is owned by the scalar value.
Literal(&'a ScalarValue),
}
impl PyExpr {
/// Inspect a direct Column/Literal node for Function call authoring.
///
/// Returns `None` for every other expression shape (arithmetic, cast,
/// scalar function, predicate, alias, qualified column, etc.).
pub(crate) fn as_direct_column_or_literal(&self) -> Option<DirectExprView<'_>> {
match &self.0 {
DfExpr::Column(Column {
relation: None,
name,
..
}) => Some(DirectExprView::UnqualifiedColumn(name.as_str())),
DfExpr::Literal(value, _) => Some(DirectExprView::Literal(value)),
_ => None,
}
}
}
#[pymethods]
impl PyExpr {
// ── comparisons ──────────────────────────────────────────────────────────
+424 -3
View File
@@ -1,20 +1,27 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::collections::HashSet;
use std::fmt;
use arrow::array::{ArrayData, ArrayRef, make_array};
use arrow::datatypes::DataType;
use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use lancedb::function::{
FunctionCapability, FunctionDefinition, FunctionOutput, FunctionParameter, FunctionSignature,
PythonFunctionDefinition,
FunctionArgument, FunctionCapability, FunctionDefinition, FunctionOutput, FunctionParameter,
FunctionSignature, PythonFunctionDefinition,
};
use pyo3::{
Bound, Py, PyAny, PyResult, Python,
exceptions::{PyRuntimeError, PyTypeError, PyValueError},
pyclass, pyfunction, pymethods,
types::{PyAnyMethods, PyBool, PyList, PyListMethods, PyTuple, PyTupleMethods},
types::{
PyAnyMethods, PyBool, PyDict, PyDictMethods, PyList, PyListMethods, PyTuple, PyTupleMethods,
},
};
use crate::error::PythonErrorExt;
use crate::expr::{DirectExprView, PyExpr};
/// Immutable first-class Function handle backed by the exact Rust value.
#[pyclass(frozen, skip_from_py_object)]
@@ -71,6 +78,291 @@ impl Function {
fn __repr__(&self) -> String {
format!("Function(id={:?})", self.inner.id().as_str())
}
/// Author an unresolved function call expression (FF-028).
///
/// Keyword-only. Does not execute. Returns a private frozen authoring value
/// that owns this exact Function and signature-ordered unresolved bindings.
#[pyo3(signature = (*args, **kwargs))]
fn __call__(
&self,
py: Python<'_>,
args: &Bound<'_, PyTuple>,
kwargs: Option<&Bound<'_, PyDict>>,
) -> PyResult<AuthoredFunctionCall> {
if !args.is_empty() {
return Err(PyTypeError::new_err(
"Function.__call__ accepts keyword arguments only",
));
}
let kwargs = match kwargs {
Some(dict) => dict.clone(),
None => PyDict::new(py),
};
AuthoredFunctionCall::try_bind(py, &self.inner, &kwargs)
}
}
/// Signature-ordered unresolved binding for Function call authoring.
///
/// Field bindings keep a case-sensitive column name until a later table API
/// resolves stable field identity in a pinned snapshot. Literal bindings are
/// already canonical [`FunctionArgument`] literals (never field args).
#[derive(Clone)]
pub(crate) enum UnresolvedArgument {
Field {
column_name: String,
},
/// Canonical typed literal; read by the later table-binding slice.
#[allow(dead_code)]
Literal(FunctionArgument),
}
impl UnresolvedArgument {
/// Structural binding text for repr/Debug only.
///
/// Literal bindings expose exact Arrow [`DataType`] and one-row nullness.
/// Never formats literal values, array Debug, IPC bytes, or payload text.
fn format_binding(&self, name: &str) -> String {
match self {
Self::Field { column_name } => {
format!("{name}=field({column_name:?})")
}
Self::Literal(argument) => {
format!(
"{name}=literal({}, null={})",
argument.data_type(),
argument.is_typed_null()
)
}
}
}
}
/// Private, frozen owner of an exact Function plus unresolved call bindings.
///
/// Exposed to Python as `lancedb._lancedb._FunctionCall`. Not constructible
/// from Python, not a catalog/Job/wire/resource, and not serializable.
#[pyclass(
name = "_FunctionCall",
module = "lancedb._lancedb",
frozen,
skip_from_py_object
)]
#[derive(Clone)]
pub struct AuthoredFunctionCall {
function: lancedb::function::Function,
bindings: Vec<(String, UnresolvedArgument)>,
}
impl AuthoredFunctionCall {
pub(crate) fn try_bind(
py: Python<'_>,
function: &lancedb::function::Function,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<Self> {
let parameters = function.signature().parameters();
let mut seen = HashSet::with_capacity(kwargs.len());
let mut by_name = std::collections::HashMap::with_capacity(kwargs.len());
for (key, value) in kwargs.iter() {
let name: String = key.extract().map_err(|_| {
PyTypeError::new_err("Function.__call__ keyword names must be strings")
})?;
if !seen.insert(name.clone()) {
return Err(PyTypeError::new_err(format!(
"duplicate Function argument for parameter `{name}`"
)));
}
by_name.insert(name, value);
}
if by_name.len() != parameters.len() {
// Prefer precise missing/unknown diagnostics over a bare arity error.
for parameter in parameters {
if !by_name.contains_key(parameter.name()) {
return Err(PyTypeError::new_err(format!(
"missing Function argument for parameter `{}`",
parameter.name()
)));
}
}
if let Some(unknown) = by_name.keys().find(|name| {
!parameters
.iter()
.any(|parameter| parameter.name() == name.as_str())
}) {
return Err(PyTypeError::new_err(format!(
"unknown Function argument `{unknown}`"
)));
}
return Err(PyTypeError::new_err(format!(
"Function.__call__ requires exactly {} arguments, got {}",
parameters.len(),
by_name.len()
)));
}
let mut bindings = Vec::with_capacity(parameters.len());
for parameter in parameters {
let Some(value) = by_name.remove(parameter.name()) else {
return Err(PyTypeError::new_err(format!(
"missing Function argument for parameter `{}`",
parameter.name()
)));
};
let argument = bind_argument(py, parameter, &value)?;
bindings.push((parameter.name().to_string(), argument));
}
if let Some(unknown) = by_name.keys().next() {
return Err(PyTypeError::new_err(format!(
"unknown Function argument `{unknown}`"
)));
}
Ok(Self {
function: function.clone(),
bindings,
})
}
/// Crate-private accessor for the later table-binding slice.
#[allow(dead_code)]
pub(crate) fn function(&self) -> &lancedb::function::Function {
&self.function
}
/// Crate-private accessor for signature-ordered unresolved bindings.
#[allow(dead_code)]
pub(crate) fn bindings(&self) -> &[(String, UnresolvedArgument)] {
&self.bindings
}
}
impl fmt::Debug for AuthoredFunctionCall {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
// Never format literal payloads; literals are type + nullness only.
f.debug_struct("_FunctionCall")
.field("function_id", &self.function.id().as_str())
.field(
"bindings",
&self
.bindings
.iter()
.map(|(name, binding)| binding.format_binding(name))
.collect::<Vec<_>>(),
)
.finish()
}
}
#[pymethods]
impl AuthoredFunctionCall {
fn __repr__(&self) -> String {
let bindings = self
.bindings
.iter()
.map(|(name, binding)| binding.format_binding(name))
.collect::<Vec<_>>()
.join(", ");
format!(
"_FunctionCall(function_id={:?}, bindings=[{bindings}])",
self.function.id().as_str()
)
}
}
fn bind_argument(
py: Python<'_>,
parameter: &FunctionParameter,
value: &Bound<'_, PyAny>,
) -> PyResult<UnresolvedArgument> {
if let Some(py_expr) = extract_public_expr_inner(py, value)? {
return match py_expr.as_direct_column_or_literal() {
Some(DirectExprView::UnqualifiedColumn(name)) => Ok(UnresolvedArgument::Field {
column_name: name.to_string(),
}),
Some(DirectExprView::Literal(scalar)) => {
let expected = parameter.data_type();
let actual = scalar.data_type();
if &actual != expected {
return Err(PyTypeError::new_err(format!(
"literal expression type mismatch for parameter `{}`: expected {expected}, got {actual}",
parameter.name()
)));
}
let array = scalar_to_one_row_array(scalar, parameter.name(), expected)?;
let argument = FunctionArgument::try_literal(array)
.map_err(|_| conversion_error(parameter.name(), expected))?;
Ok(UnresolvedArgument::Literal(argument))
}
None => Err(PyTypeError::new_err(format!(
"parameter `{}` requires a direct column reference or literal",
parameter.name()
))),
};
}
let argument =
python_value_to_literal_argument(py, value, parameter.name(), parameter.data_type())?;
Ok(UnresolvedArgument::Literal(argument))
}
fn extract_public_expr_inner<'py>(
py: Python<'py>,
value: &Bound<'py, PyAny>,
) -> PyResult<Option<PyExpr>> {
let expr_cls = py.import("lancedb.expr")?.getattr("Expr")?;
if !value.is_instance(&expr_cls)? {
return Ok(None);
}
let inner = value.getattr("_inner")?;
let py_expr: PyExpr = inner
.extract()
.map_err(|_| PyTypeError::new_err("lancedb.expr.Expr must wrap a native PyExpr"))?;
Ok(Some(py_expr))
}
fn python_value_to_literal_argument(
py: Python<'_>,
value: &Bound<'_, PyAny>,
parameter_name: &str,
data_type: &DataType,
) -> PyResult<FunctionArgument> {
let pa = py.import("pyarrow")?;
let type_obj = data_type
.to_pyarrow(py)
.map_err(|_| conversion_error(parameter_name, data_type))?;
let values = PyList::new(py, std::slice::from_ref(value))
.map_err(|_| conversion_error(parameter_name, data_type))?;
let kwargs = PyDict::new(py);
kwargs
.set_item("type", type_obj)
.map_err(|_| conversion_error(parameter_name, data_type))?;
let array_obj = pa
.call_method("array", (values,), Some(&kwargs))
.map_err(|_| conversion_error(parameter_name, data_type))?;
let array_data = ArrayData::from_pyarrow_bound(&array_obj)
.map_err(|_| conversion_error(parameter_name, data_type))?;
let array: ArrayRef = make_array(array_data);
FunctionArgument::try_literal(array).map_err(|_| conversion_error(parameter_name, data_type))
}
fn scalar_to_one_row_array(
scalar: &datafusion_common::ScalarValue,
parameter_name: &str,
data_type: &DataType,
) -> PyResult<ArrayRef> {
scalar
.to_array_of_size(1)
.map_err(|_| conversion_error(parameter_name, data_type))
}
fn conversion_error(parameter_name: &str, data_type: &DataType) -> pyo3::PyErr {
PyValueError::new_err(format!(
"cannot convert argument for parameter `{parameter_name}` to type {data_type}"
))
}
/// Private, frozen owner of the exact Rust [`FunctionDefinition`].
@@ -274,3 +566,132 @@ fn parse_capability_triple(item: &Bound<'_, PyAny>) -> PyResult<FunctionCapabili
_ => Err(PyValueError::new_err("unsupported capability kind")),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use arrow::array::Int32Array;
use lancedb::expr::col as ldb_col;
use lancedb::expr::lit as ldb_lit;
use lancedb::function::{FunctionId, FunctionOutput, FunctionParameter, FunctionSignature};
fn sample_function() -> lancedb::function::Function {
let signature = FunctionSignature::try_new(
vec![
FunctionParameter::new("text", DataType::Utf8),
FunctionParameter::new("limit", DataType::Int32),
],
FunctionOutput::new(DataType::Utf8, true),
)
.expect("signature");
lancedb::function::Function::new(
FunctionId::try_new("fn.exact.call-handle").expect("id"),
signature,
)
}
#[test]
fn authored_call_normalizes_binding_order_and_preserves_column_case() {
let function = sample_function();
let int_lit =
FunctionArgument::try_literal(Arc::new(Int32Array::from(vec![8])) as ArrayRef)
.expect("literal");
let authored = AuthoredFunctionCall {
function: function.clone(),
bindings: vec![
(
"text".to_string(),
UnresolvedArgument::Field {
column_name: "firstName".to_string(),
},
),
("limit".to_string(), UnresolvedArgument::Literal(int_lit)),
],
};
assert_eq!(authored.function().id().as_str(), "fn.exact.call-handle");
assert_eq!(authored.bindings().len(), 2);
assert_eq!(authored.bindings()[0].0, "text");
match &authored.bindings()[0].1 {
UnresolvedArgument::Field { column_name } => assert_eq!(column_name, "firstName"),
UnresolvedArgument::Literal(_) => panic!("expected field binding, got literal"),
}
match &authored.bindings()[1].1 {
UnresolvedArgument::Literal(argument) => {
assert_eq!(argument.data_type(), &DataType::Int32);
assert!(!argument.is_typed_null());
}
UnresolvedArgument::Field { .. } => panic!("expected literal binding, got field"),
}
let rendered = authored.__repr__();
assert!(rendered.starts_with("_FunctionCall(function_id="));
assert!(
rendered.find("text=field(\"firstName\")").unwrap()
< rendered.find("limit=literal(Int32, null=false)").unwrap()
);
assert!(!rendered.contains('8'));
}
#[test]
fn authored_call_repr_and_debug_omit_literal_payload() {
let function = sample_function();
let sentinel = FunctionArgument::try_literal(Arc::new(arrow::array::StringArray::from(
vec![Some("LITERAL_PAYLOAD_SENTINEL_call_xyz_42")],
)) as ArrayRef)
.expect("literal");
let authored = AuthoredFunctionCall {
function,
bindings: vec![
("text".to_string(), UnresolvedArgument::Literal(sentinel)),
(
"limit".to_string(),
UnresolvedArgument::Literal(
FunctionArgument::try_literal(Arc::new(Int32Array::from(vec![Some(
2_147_000_123,
)])) as ArrayRef)
.expect("int literal"),
),
),
],
};
let rendered = format!("{authored:?}\n{}", authored.__repr__());
assert!(!rendered.contains("LITERAL_PAYLOAD_SENTINEL_call_xyz_42"));
assert!(!rendered.contains("2147000123"));
assert!(rendered.contains("text=literal(Utf8, null=false)"));
assert!(rendered.contains("limit=literal(Int32, null=false)"));
}
#[test]
fn typed_null_literal_argument_round_trips_type() {
let null =
FunctionArgument::try_literal(
Arc::new(Int32Array::from(vec![None as Option<i32>])) as ArrayRef
)
.expect("typed null");
assert!(null.is_typed_null());
assert_eq!(null.data_type(), &DataType::Int32);
}
#[test]
fn direct_expr_view_accepts_column_and_literal_only() {
let column = PyExpr(ldb_col("firstName"));
match column.as_direct_column_or_literal() {
Some(DirectExprView::UnqualifiedColumn(name)) => assert_eq!(name, "firstName"),
_ => panic!("expected unqualified column"),
}
let literal = PyExpr(ldb_lit(8i64));
match literal.as_direct_column_or_literal() {
Some(DirectExprView::Literal(value)) => {
assert_eq!(value.data_type(), DataType::Int64);
}
_ => panic!("expected literal"),
}
let complex = PyExpr(ldb_col("text").eq(ldb_lit("x")));
assert!(complex.as_direct_column_or_literal().is_none());
}
}
+1
View File
@@ -48,6 +48,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<Table>()?;
m.add_class::<crate::function::Function>()?;
m.add_class::<crate::function::PyFunctionDefinition>()?;
m.add_class::<crate::function::AuthoredFunctionCall>()?;
m.add_class::<crate::job::Job>()?;
m.add_class::<crate::job::JobInfo>()?;
m.add_class::<crate::job::JobDescription>()?;