diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 7798bdf73..2556810c4 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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.""" diff --git a/python/python/tests/test_first_class_function_call.py b/python/python/tests/test_first_class_function_call.py new file mode 100644 index 000000000..65cbea325 --- /dev/null +++ b/python/python/tests/test_first_class_function_call.py @@ -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)" diff --git a/python/src/expr.rs b/python/src/expr.rs index 242e88b05..5dde47d8a 100644 --- a/python/src/expr.rs +++ b/python/src/expr.rs @@ -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> { + 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 ────────────────────────────────────────────────────────── diff --git a/python/src/function.rs b/python/src/function.rs index c83710000..5f49e901f 100644 --- a/python/src/function.rs +++ b/python/src/function.rs @@ -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 { + 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 { + 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::>(), + ) + .finish() + } +} + +#[pymethods] +impl AuthoredFunctionCall { + fn __repr__(&self) -> String { + let bindings = self + .bindings + .iter() + .map(|(name, binding)| binding.format_binding(name)) + .collect::>() + .join(", "); + format!( + "_FunctionCall(function_id={:?}, bindings=[{bindings}])", + self.function.id().as_str() + ) + } +} + +fn bind_argument( + py: Python<'_>, + parameter: &FunctionParameter, + value: &Bound<'_, PyAny>, +) -> PyResult { + 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> { + 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 { + 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 { + 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 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])) 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()); + } +} diff --git a/python/src/lib.rs b/python/src/lib.rs index c6355985a..dd49ce6f5 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -48,6 +48,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?;