mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +00:00
feat(python): add first-class function call authoring
This commit is contained in:
@@ -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
@@ -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
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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>()?;
|
||||
|
||||
Reference in New Issue
Block a user