diff --git a/Cargo.lock b/Cargo.lock index ee2a3d919..f38252fc5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1422,6 +1422,15 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "blake2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b5d4d889834ee8ecfc0f8426ad30faf7cdcb10f741a8e6d7224d95325479f6f" +dependencies = [ + "digest 0.11.3", +] + [[package]] name = "blake3" version = "1.8.5" @@ -2734,7 +2743,7 @@ dependencies = [ "arrow", "arrow-buffer", "base64 0.23.1", - "blake2", + "blake2 0.10.6", "blake3", "chrono", "chrono-tz 0.10.4", @@ -3197,6 +3206,15 @@ dependencies = [ "dirs-sys 0.5.0", ] +[[package]] +name = "dirs" +version = "7.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d57d423b3c82e89b9a24ca3091fee61f456a26edbd28d26c65906f4bc1dcd8f" +dependencies = [ + "dirs-sys 0.5.0", +] + [[package]] name = "dirs-sys" version = "0.4.1" @@ -3442,6 +3460,17 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "erased-serde" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e59b668157f32d3459d4bd634861a06ffa2503a461ca5f0916b9f0252605e66" +dependencies = [ + "serde", + "serde_core", + "typeid", +] + [[package]] name = "errno" version = "0.3.14" @@ -3572,6 +3601,19 @@ dependencies = [ "subtle", ] +[[package]] +name = "file_url" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec49f140842acf886ff3e4fd47600ffdbb5f1b1ba70f6775880e688837a0b174" +dependencies = [ + "itertools 0.15.0", + "percent-encoding", + "thiserror 2.0.18", + "typed-path", + "url", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -3616,6 +3658,15 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "float-cmp" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b09cf3155332e944990140d967ff5eceb70df778b34f77d8075db46e4704e6d8" +dependencies = [ + "num-traits", +] + [[package]] name = "float8" version = "0.7.0" @@ -3747,6 +3798,15 @@ version = "1.20260804.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "82eb03a32a1d50555353c85a7b9d3279a6f1e91af9890b789acdf544ed57c8d7" +[[package]] +name = "fs-err" +version = "3.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5c95b673b8f6f7235229ae11c5642d81b04c2e64c1e2fb417bc0cf73ca45f29" +dependencies = [ + "autocfg", +] + [[package]] name = "fs4" version = "0.13.1" @@ -4323,6 +4383,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "halfbrown" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c7ed2f2edad8a14c8186b847909a41fbb9c3eafa44f88bd891114ed5019da09" +dependencies = [ + "hashbrown 0.16.1", + "serde", +] + [[package]] name = "hash32" version = "0.3.1" @@ -4428,7 +4498,7 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" dependencies = [ - "dirs", + "dirs 6.0.0", "futures", "http", "indicatif", @@ -5624,7 +5694,7 @@ dependencies = [ "datafusion-functions-nested", "datafusion-physical-expr", "datafusion-physical-plan", - "dirs", + "dirs 6.0.0", "fst", "futures", "half", @@ -6132,6 +6202,7 @@ dependencies = [ "pyo3", "pyo3-async-runtimes", "pyo3-build-config", + "rattler_conda_types", "reqwest 0.13.5", "serde", "serde_json", @@ -6140,6 +6211,29 @@ dependencies = [ "uuid", ] +[[package]] +name = "lazy-regex" +version = "3.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0a2cce8a22920a29fdcc8c9170aa5d29edd9fbcf8bea856e24c2cb1c2419f67" +dependencies = [ + "lazy-regex-proc_macros", + "once_cell", + "regex", +] + +[[package]] +name = "lazy-regex-proc_macros" +version = "3.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f1e4b1f1d8c71edb4ff2b2150ec4606445fd2ddf67ddc097655daf59286baa6" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "syn 3.0.3", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -6903,6 +6997,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "nom-language" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2de2bc5b451bfedaef92c90b8939a8fff5770bdcc1fafd6239d086aab8fa6b29" +dependencies = [ + "nom 8.0.0", +] + [[package]] name = "now" version = "0.1.3" @@ -8587,6 +8690,18 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40e24eee682d89fb193496edf918a7f407d30175b2e785fe057e4392dfd182e0" +[[package]] +name = "purl" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60ebe4262ae91ddd28c8721111a0a6e9e58860e211fc92116c4bb85c98fd96ad" +dependencies = [ + "hex", + "percent-encoding", + "serde", + "thiserror 2.0.18", +] + [[package]] name = "pyo3" version = "0.29.3" @@ -8973,6 +9088,99 @@ dependencies = [ "rustversion", ] +[[package]] +name = "rattler_conda_types" +version = "0.58.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fda631350e17df9bc789082d6b6c0fbde38e65cba580593ee3317fb6730644f2" +dependencies = [ + "ahash", + "core-foundation 0.10.1", + "dirs 7.0.0", + "fancy-regex 0.19.2", + "file_url", + "fs-err", + "glob", + "hex", + "indexmap 2.14.0", + "itertools 0.15.0", + "jiff", + "lazy-regex", + "memmap2 0.9.10", + "nom 8.0.0", + "nom-language", + "purl", + "rattler_conda_version", + "rattler_digest", + "rattler_macros", + "rattler_redaction", + "regex", + "serde", + "serde-untagged", + "serde_json", + "serde_repr", + "serde_with", + "serde_yaml", + "simd-json", + "smallvec", + "strum", + "tempfile", + "thiserror 2.0.18", + "tracing", + "typed-path", + "url", +] + +[[package]] +name = "rattler_conda_version" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee922b54cd30c198e93c542096a0d671cefba1a97600768b2ec4af3c11c44f55" +dependencies = [ + "itertools 0.15.0", + "nom 8.0.0", + "serde", + "smallvec", + "thiserror 2.0.18", +] + +[[package]] +name = "rattler_digest" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5bda554663860e074665f838a4d9a89fbf671f576eaf288d29525114b345adc9" +dependencies = [ + "blake2 0.11.0", + "digest 0.11.3", + "hex", + "md-5", + "serde", + "serde-untagged", + "serde_bytes", + "serde_with", + "sha2 0.11.0", + "url", +] + +[[package]] +name = "rattler_macros" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a42690ed14c8b9d610aafc0f2eb715b85c05f3d3a2a307896e4f8d820aaa284f" +dependencies = [ + "quote", + "syn 3.0.3", +] + +[[package]] +name = "rattler_redaction" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b5ee47d0375643b037205a4fe238b90c518ca3fc83c7c240d82c49ba87d61c" +dependencies = [ + "url", +] + [[package]] name = "raw-cpuid" version = "11.6.0" @@ -9926,6 +10134,28 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-untagged" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9faf48a4a2d2693be24c6289dbe26552776eb7737074e6722891fadbe6c5058" +dependencies = [ + "erased-serde", + "serde", + "serde_core", + "typeid", +] + +[[package]] +name = "serde_bytes" +version = "0.11.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" +dependencies = [ + "serde", + "serde_core", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -10056,6 +10286,19 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "serde_yaml" +version = "0.9.34+deprecated" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" +dependencies = [ + "indexmap 2.14.0", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + [[package]] name = "serde_yaml_ng" version = "0.10.0" @@ -10170,7 +10413,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32824fab5e16e6c4d86dc1ba84489390419a39f97699852b66480bb87d297ed8" dependencies = [ "bstr", - "dirs", + "dirs 6.0.0", "os_str_bytes", ] @@ -10235,6 +10478,20 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" +[[package]] +name = "simd-json" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ebe141eb4a23ecc4d5de93fa5b139ac65dc07c38037a2fe9a8fb7c9d36bd832" +dependencies = [ + "halfbrown", + "ref-cast", + "serde", + "serde_json", + "simdutf8", + "value-trait", +] + [[package]] name = "simd_cesu8" version = "1.1.1" @@ -10274,6 +10531,9 @@ name = "smallvec" version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +dependencies = [ + "serde", +] [[package]] name = "smartstring" @@ -11403,6 +11663,12 @@ version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ce898823d8cd0f028662d2f349694b1af0a5cc9be83f1d213a00d3147445e1e" +[[package]] +name = "typeid" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c" + [[package]] name = "typenum" version = "1.20.0" @@ -11583,6 +11849,18 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "value-trait" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "545a3d7089f11f1741d12963c057adf9443fc24325110a6a6daf582e825fd51b" +dependencies = [ + "float-cmp", + "halfbrown", + "itoa", + "ryu", +] + [[package]] name = "vcpkg" version = "0.2.15" @@ -12561,7 +12839,7 @@ dependencies = [ "colored", "const-str", "ctor", - "dirs", + "dirs 6.0.0", "futures", "git-version", "humantime", diff --git a/python/Cargo.toml b/python/Cargo.toml index 07ca7de02..7c11d5fdf 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -37,6 +37,7 @@ pyo3-async-runtimes = { version = "0.29", features = [ "tokio-runtime", ] } reqwest.workspace = true +rattler_conda_types = { version = "0.58.3", default-features = false } pin-project.workspace = true futures.workspace = true serde.workspace = true diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 4029525e9..521cc46bb 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -1310,3 +1310,4 @@ class AuthzObject: def __repr__(self) -> str: ... def validate_function_retry_policy(json: str) -> None: ... +def _validate_conda_requirement(requirement: str) -> None: ... diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 737abc44e..4bf028306 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -44,6 +44,7 @@ from typing import ( ) import pyarrow as pa +from packaging.requirements import Requirement from pydantic import ( AfterValidator, BaseModel, @@ -57,6 +58,7 @@ from pydantic import ( from .schema import blob as _blob_field from .schema import is_blob_v2_field as _is_blob_v2_field from ._blob_reader import BlobReader +from ._lancedb import _validate_conda_requirement from .secrets import EnvVarSecret _Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1) @@ -1701,6 +1703,29 @@ def _packaged_by_source(value: Any) -> bool: return True +def _validate_package_requirements(kind: str, packages: Sequence[str]) -> None: + for index, package in enumerate(packages): + field = f"{kind}[{index}]" + if not isinstance(package, str): + raise TypeError(f"{field} must be a string, got {package!r}") + if not package or package != package.strip(): + raise ValueError( + f"{field} requirement {package!r} must be non-empty and trimmed" + ) + try: + if kind == "conda": + _validate_conda_requirement(package) + else: + candidate = package + if re.match(r"^[A-Za-z][A-Za-z0-9+.-]*://", package): + candidate = f"dependency @ {package}" + Requirement(candidate) + except ValueError as error: + raise ValueError( + f"invalid {field} requirement {package!r}: {error}" + ) from error + + def _requirement_name(requirement: str) -> str: requirement = requirement.rpartition("::")[2] match = _REQUIREMENT_NAME.match(requirement.strip()) @@ -2242,6 +2267,7 @@ class FunctionDefinition: raise ValueError("a Function environment is pip or conda, not both") if conda_channels and not conda: raise ValueError("conda_channels requires conda packages") + _validate_package_requirements("conda" if conda else "pip", conda or pip) gpu_marker = _normalize_gpu_marker(gpu) if gpu_class is not None: _validate_gpu_class_name(gpu_class) @@ -2595,9 +2621,12 @@ def function( non-nullable. Required for batch annotations. When omitted, the output is inferred from the return annotation. pip : sequence of str, optional - Pip requirements for the remote environment. + Pip requirements for the remote environment. Syntax is validated at + decoration time, including extras and environment markers. Named direct + references and bare package or git URLs are supported. conda : sequence of str, optional - Conda packages for the remote environment, instead of ``pip``. + Conda packages for the remote environment, instead of ``pip``. MatchSpec + syntax is validated at decoration time, including channels and builds. conda_channels : sequence of str, optional Conda channels in priority order; requires ``conda``. apt : sequence of str, optional diff --git a/python/python/tests/test_function_requirements.py b/python/python/tests/test_function_requirements.py new file mode 100644 index 000000000..6f94d1313 --- /dev/null +++ b/python/python/tests/test_function_requirements.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pytest + +from lancedb import function, udf +from lancedb.functions import FunctionDefinition + + +def identity(value: int) -> int: + return value + + +@pytest.mark.parametrize( + "package", + [ + "torch>=2.0,,<3", + "torch>=", + "torch=>2.0", + "torch[", + "torch; python_version >=", + "", + " torch", + "torch ", + ], +) +def test_invalid_pip_requirement_fails_before_packaging(package, monkeypatch): + def unexpected_packaging(*args, **kwargs): + pytest.fail("invalid requirements must fail before packaging the function") + + monkeypatch.setattr("lancedb.functions._infer_signature", unexpected_packaging) + with pytest.raises(ValueError) as error: + function(pip=["numpy", package])(identity) + assert "pip[1]" in str(error.value) + assert repr(package) in str(error.value) + + +@pytest.mark.parametrize( + "package", + [ + "numpy>=1,,<3", + "numpy>=", + "numpy[version=]", + "numpy[unknown_field=value]", + "", + " numpy", + "numpy ", + ], +) +def test_invalid_conda_requirement_names_entry(package): + with pytest.raises(ValueError) as error: + function(conda=["scipy", package])(identity) + assert "conda[1]" in str(error.value) + assert repr(package) in str(error.value) + + +@pytest.mark.parametrize("kind", ["pip", "conda"]) +@pytest.mark.parametrize("package", [None, 42, ["numpy"]]) +def test_requirement_must_be_a_string(kind, package): + with pytest.raises(TypeError, match=rf"{kind}\[1\] must be a string"): + function(**{kind: ["numpy", package]})(identity) + + +@pytest.mark.parametrize( + "package", + [ + "numpy", + "torch>=2.0,<3", + 'requests[security,socks]>=2; python_version >= "3.10"', + "numpy (>=1.26, <3)", + "six @ https://example.com/six-1.16.0.tar.gz", + "six @ git+https://github.com/benjaminp/six.git@main", + "https://example.com/six-1.16.0.tar.gz", + "git+https://github.com/benjaminp/six.git@main", + "git+ssh://git@github.com/benjaminp/six.git@main", + ], +) +def test_supported_pip_requirements_are_preserved(package): + definition = function(pip=[package])(identity) + assert definition.registration_request.runtime.environment.packages == (package,) + + +@pytest.mark.parametrize( + "package", + [ + "numpy", + "numpy=1.26", + "numpy>=1.26,<3", + "numpy=1.26=py310_0", + "numpy 1.26 py310_0", + "conda-forge::numpy>=1.26", + "conda-forge/linux-64::numpy>=1.26", + "numpy[version='>=1.26,<3',build='py310*']", + ], +) +def test_supported_conda_requirements_are_preserved(package): + definition = function(conda=[package])(identity) + assert definition.registration_request.runtime.environment.packages == (package,) + + +def test_requirements_still_sort_and_deduplicate(): + definition = function(pip=["torch>=2", "numpy", "torch>=2"])(identity) + assert definition.registration_request.runtime.environment.packages == ( + "numpy", + "torch>=2", + ) + + +def test_deprecated_udf_alias_validates_requirements(): + with pytest.warns(DeprecationWarning, match="use lancedb.function"): + with pytest.raises(ValueError, match=r"pip\[0\]"): + udf(pip=["torch>=2.0,,<3"])(identity) + + +def test_direct_function_definition_validates_requirements(): + with pytest.raises(ValueError, match=r"pip\[0\]"): + FunctionDefinition( + identity, + name=None, + input_schema=None, + output_schema=None, + pip=("torch>=2.0,,<3",), + env={}, + python_version=None, + ) diff --git a/python/src/lib.rs b/python/src/lib.rs index f78c4e7bb..b0698d5e0 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -150,6 +150,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(table::tokenize, m)?)?; m.add_function(wrap_pyfunction!(permutation::async_permutation_builder, m)?)?; m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?; + m.add_function(wrap_pyfunction!(util::validate_conda_requirement, m)?)?; m.add_function(wrap_pyfunction!(util::validate_function_retry_policy, m)?)?; m.add_function(wrap_pyfunction!(util::parse_field_path, m)?)?; m.add_function(wrap_pyfunction!(query::fts_query_to_json, m)?)?; diff --git a/python/src/util.rs b/python/src/util.rs index a956066f1..44efa165a 100644 --- a/python/src/util.rs +++ b/python/src/util.rs @@ -9,6 +9,7 @@ use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, pyfunction, }; +use rattler_conda_types::{MatchSpec, ParseStrictness}; /// A wrapper around a rust builder /// @@ -61,6 +62,13 @@ pub fn validate_table_name(table_name: &str) -> PyResult<()> { .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> {