Files
lancedb/python/src/util.rs
T
Jack Ye ee231aad06 fix(python): validate function package requirements before registration (#4971)
Function dependency typos currently reach the environment build before
failing. Validate pip requirements with
`packaging.requirements.Requirement` and conda requirements with
Rattler's MatchSpec parser when constructing a function definition,
before source packaging or registration. Errors identify the original
entry and its index.

```python
@lancedb.function(pip=["numpy", "torch>=2.0,,<3"])
def score(value: float) -> float:
    return value
# ValueError: invalid pip[1] requirement 'torch>=2.0,,<3': ...
```

Preserves pip extras, markers, named references, bare archive/git URLs,
and conda channel/build specifications. The existing native extension
exposes the conda parser, so no additional Python package is required.
Adds regression coverage for invalid and supported forms, including the
deprecated `udf` alias.
2026-11-11 22:56:36 -08:00

92 lines
3.2 KiB
Rust

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::Mutex;
use lancedb::DistanceType;
use pyo3::{
PyResult,
exceptions::{PyRuntimeError, PyValueError},
pyfunction,
};
use rattler_conda_types::{MatchSpec, ParseStrictness};
/// A wrapper around a rust builder
///
/// Rust builders are often implemented so that the builder methods
/// consume the builder and return a new one. This is not compatible
/// with the pyo3, which, being garbage collected, cannot easily obtain
/// ownership of an object.
///
/// This wrapper converts the compile-time safety of rust into runtime
/// errors if any attempt to use the builder happens after it is consumed.
pub struct BuilderWrapper<T> {
name: String,
inner: Mutex<Option<T>>,
}
impl<T> BuilderWrapper<T> {
pub fn new(name: impl AsRef<str>, inner: T) -> Self {
Self {
name: name.as_ref().to_string(),
inner: Mutex::new(Some(inner)),
}
}
pub fn consume<O>(&self, mod_fn: impl FnOnce(T) -> O) -> PyResult<O> {
let mut inner = self.inner.lock().unwrap();
let inner_builder = inner.take().ok_or_else(|| {
PyRuntimeError::new_err(format!("{} has already been consumed", self.name))
})?;
let result = mod_fn(inner_builder);
Ok(result)
}
}
pub fn parse_distance_type(distance_type: impl AsRef<str>) -> PyResult<DistanceType> {
match distance_type.as_ref().to_lowercase().as_str() {
"l2" => Ok(DistanceType::L2),
"cosine" => Ok(DistanceType::Cosine),
"dot" => Ok(DistanceType::Dot),
"hamming" => Ok(DistanceType::Hamming),
_ => Err(PyValueError::new_err(format!(
"Invalid distance type '{}'. Must be one of l2, cosine, dot, or hamming",
distance_type.as_ref()
))),
}
}
#[pyfunction]
pub fn validate_table_name(table_name: &str) -> PyResult<()> {
lancedb::utils::validate_table_name(table_name)
.map_err(|e| PyValueError::new_err(e.to_string()))
}
#[pyfunction(name = "_validate_conda_requirement")]
pub fn validate_conda_requirement(requirement: &str) -> PyResult<()> {
MatchSpec::from_str(requirement, ParseStrictness::Lenient)
.map(|_| ())
.map_err(|error| PyValueError::new_err(error.to_string()))
}
/// Parse field references with the same identifier syntax as Lance writes.
#[pyfunction(name = "_parse_field_path")]
pub fn parse_field_path(path: &str) -> PyResult<Vec<String>> {
lance_core::datatypes::parse_field_path(path).map_err(|e| PyValueError::new_err(e.to_string()))
}
/// A wrapper around a LanceDB type to allow it to be used in Python
#[derive(Debug, Clone)]
pub struct PyLanceDB<T>(pub T);
/// Validate the Function retry wire contract using the same rules as Rust APIs.
#[pyo3::pyfunction]
pub fn validate_function_retry_policy(json: &str) -> pyo3::PyResult<()> {
let policy: lancedb::function::FunctionRetryPolicy =
serde_json::from_str(json).map_err(|error| {
pyo3::exceptions::PyValueError::new_err(format!("invalid retry_policy: {error}"))
})?;
policy
.validate()
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))
}