Merge remote-tracking branch 'refs/remotes/origin/main' into gatekeeper/fix-2820-1

# Conflicts:
#	rust/lancedb/src/remote/table.rs

python/python/lancedb/functions.py

@@ -222,6 +222,7 @@ class PythonEnvironmentSpec(_RemoteValue):
 
     kind: str
     packages: tuple[str, ...] = ()
+    channels: tuple[str, ...] = ()
     path: Optional[str] = None
     modules: tuple[str, ...] = ()
     image: Optional[str] = None
@@ -909,13 +910,25 @@ class UdfDefinition:
         pip: tuple[str, ...],
         env: Mapping[str, str],
         python_version: Optional[str],
+        conda: tuple[str, ...] = (),
+        conda_channels: tuple[str, ...] = (),
     ):
         function_name = name or function.__name__
         if not _FUNCTION_NAME.fullmatch(function_name):
             raise ValueError(f"invalid Function name: {function_name!r}")
-        packages = tuple(sorted(set(pip)))
+        if pip and conda:
+            raise ValueError("a Function environment is pip or conda, not both")
+        if conda_channels and not conda:
+            raise ValueError("conda_channels requires conda packages")
+        packages = tuple(sorted(set(conda if conda else pip)))
         if any(not package or package != package.strip() for package in packages):
-            raise ValueError("pip requirements must be non-empty and trimmed")
+            raise ValueError("package requirements must be non-empty and trimmed")
+        if conda:
+            environment_spec = PythonEnvironmentSpec(
+                kind="conda", packages=packages, channels=tuple(conda_channels)
+            )
+        else:
+            environment_spec = PythonEnvironmentSpec(kind="pip", packages=packages)
         environment = dict(env)
         if any(
             not isinstance(key, str) or not isinstance(value, str)
@@ -929,7 +942,7 @@ class UdfDefinition:
             kind="python",
             python_version=python_version
             or f"{sys.version_info.major}.{sys.version_info.minor}",
-            environment=PythonEnvironmentSpec(kind="pip", packages=packages),
+            environment=environment_spec,
             env=environment,
         )
         self._function = function
@@ -976,6 +989,8 @@ def udf(
     pip: tuple[str, ...] | list[str] = (),
     env: Optional[Mapping[str, str]] = None,
     python_version: Optional[str] = None,
+    conda: tuple[str, ...] | list[str] = (),
+    conda_channels: tuple[str, ...] | list[str] = (),
 ) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
 
 
@@ -988,6 +1003,8 @@ def udf(
     pip: tuple[str, ...] | list[str] = (),
     env: Optional[Mapping[str, str]] = None,
     python_version: Optional[str] = None,
+    conda: tuple[str, ...] | list[str] = (),
+    conda_channels: tuple[str, ...] | list[str] = (),
 ):
     """Prepare a scalar Python callable for remote Function registration.
 
@@ -1010,6 +1027,10 @@ def udf(
         provided together with ``input_schema``.
     pip : sequence of str, optional
         Pip requirements for the remote environment.
+    conda : sequence of str, optional
+        Conda packages for the remote environment, instead of ``pip``.
+    conda_channels : sequence of str, optional
+        Conda channels in priority order; requires ``conda``.
     env : mapping of str to str, optional
         Environment variables included in the Function definition.
     python_version : str, optional
@@ -1049,6 +1070,8 @@ def udf(
             pip=tuple(pip),
             env={} if env is None else env,
             python_version=python_version,
+            conda=tuple(conda),
+            conda_channels=tuple(conda_channels),
         )
 
     if function is None:

python/python/lancedb/index.py

@@ -7,6 +7,7 @@ from typing import List, Literal, Optional
 from ._lancedb import (
     IndexConfig,
 )
+from .query import DocumentGranularity
 from .types import BaseTokenizerType
 
 lang_mapping = {
@@ -121,6 +122,11 @@ class FTS:
 
     >>> config = FTS(block_size=256)
 
+    Create an index that treats each deepest-list element as one document:
+
+    >>> from lancedb.query import DocumentGranularity
+    >>> config = FTS(document_granularity=DocumentGranularity.LIST_ELEMENT)
+
     Attributes
     ----------
     with_position : bool, default False
@@ -172,6 +178,11 @@ class FTS:
         roughly half of the available CPU cores. The effective value is
         limited by the available compute capacity. This build-only setting is
         not persisted with the index and does not apply to remote tables.
+    document_granularity : DocumentGranularity, default ROW
+        ``ROW`` treats the selected text in one table row as one document.
+        ``LIST_ELEMENT`` treats each element of the deepest list on the indexed
+        field path as one document and returns its physical coordinates in
+        ``_doc_index`` for matching queries.
 
     Notes
     -----
@@ -196,6 +207,7 @@ class FTS:
     custom_stop_words: Optional[List[str]] = None
     memory_limit: Optional[int] = None
     num_workers: Optional[int] = None
+    document_granularity: DocumentGranularity = DocumentGranularity.ROW
 
 
 @dataclass

python/python/lancedb/query.py

@@ -376,6 +376,13 @@ class FullTextOperator(str, Enum):
     OR = "OR"
 
 
+class DocumentGranularity(str, Enum):
+    """The unit treated as one full-text-search document."""
+
+    ROW = "row"
+    LIST_ELEMENT = "list_element"
+
+
 class Occur(str, Enum):
     SHOULD = "SHOULD"
     MUST = "MUST"
@@ -479,6 +486,10 @@ class MatchQuery(FullTextQuery):
     prefix_length : int, optional
         The number of beginning characters being unchanged for fuzzy matching.
         This is useful to achieve prefix matching.
+    document_granularity : DocumentGranularity, optional
+        Explicitly select row or deepest-list-element documents. If omitted,
+        the indexed granularity is inferred. When both granularities are indexed
+        for the field, this must be specified. With no index, row granularity is used.
     """
 
     query: str
@@ -488,6 +499,9 @@ class MatchQuery(FullTextQuery):
     max_expansions: int = pydantic.Field(50, kw_only=True)
     operator: FullTextOperator = pydantic.Field(FullTextOperator.OR, kw_only=True)
     prefix_length: int = pydantic.Field(0, kw_only=True)
+    document_granularity: Optional[DocumentGranularity] = pydantic.Field(
+        None, kw_only=True
+    )
 
     def query_type(self) -> FullTextQueryType:
         return FullTextQueryType.MATCH
@@ -504,11 +518,20 @@ class PhraseQuery(FullTextQuery):
         The query string to match against.
     column : str
         The name of the column to match against.
+    slop : int, default 0
+        The maximum number of intervening positions permitted in the phrase.
+    document_granularity : DocumentGranularity, optional
+        Explicitly select row or deepest-list-element documents. If omitted,
+        the indexed granularity is inferred. When both granularities are indexed
+        for the field, this must be specified. With no index, row granularity is used.
     """
 
     query: str
     column: str
     slop: int = pydantic.Field(0, kw_only=True)
+    document_granularity: Optional[DocumentGranularity] = pydantic.Field(
+        None, kw_only=True
+    )
 
     def query_type(self) -> FullTextQueryType:
         return FullTextQueryType.MATCH_PHRASE

python/python/lancedb/remote/table.py

@@ -61,6 +61,7 @@ from lancedb.table import _normalize_progress
 
 from ..query import (
     AnalyzePlanDistributedMetrics,
+    DocumentGranularity,
     LanceQueryBuilder,
     LanceTakeQueryBuilder,
     LanceVectorQueryBuilder,
@@ -349,6 +350,7 @@ class RemoteTable(Table):
         ngram_max_length: int = 3,
         prefix_only: bool = False,
         block_size: int = 128,
+        document_granularity: DocumentGranularity = DocumentGranularity.ROW,
         name: Optional[str] = None,
     ):
         """Create a full-text search index on a column.
@@ -371,6 +373,7 @@ class RemoteTable(Table):
             ngram_max_length=ngram_max_length,
             prefix_only=prefix_only,
             block_size=block_size,
+            document_granularity=document_granularity,
         )
         LOOP.run(
             self._table.create_index(

python/python/lancedb/table.py

@@ -85,6 +85,7 @@ from .query import (
     AsyncQuery,
     AsyncTakeQuery,
     AsyncVectorQuery,
+    DocumentGranularity,
     FullTextQuery,
     LanceEmptyQueryBuilder,
     LanceFtsQueryBuilder,
@@ -1168,6 +1169,7 @@ class Table(ABC):
         ngram_max_length: int = 3,
         prefix_only: bool = False,
         block_size: int = 128,
+        document_granularity: DocumentGranularity = DocumentGranularity.ROW,
         wait_timeout: Optional[timedelta] = None,
         name: Optional[str] = None,
     ):
@@ -1246,6 +1248,11 @@ class Table(ABC):
             The number of documents per compressed posting block. Must be 128
             or 256. A value of 256 uses the experimental FTS V3 format and
             may introduce breaking changes.
+        document_granularity: DocumentGranularity, default ROW
+            ``ROW`` treats the selected text in one table row as one document.
+            ``LIST_ELEMENT`` treats each element of the deepest list on the field
+            path as one document and returns its physical coordinates in
+            ``_doc_index`` for matching queries.
         wait_timeout: timedelta, optional
             The timeout to wait if indexing is asynchronous.
         name: str, optional
@@ -3273,6 +3280,7 @@ class LanceTable(Table):
         ngram_max_length: int = 3,
         prefix_only: bool = False,
         block_size: int = 128,
+        document_granularity: DocumentGranularity = DocumentGranularity.ROW,
         name: Optional[str] = None,
     ):
         """Create a full-text search index on a column.
@@ -3324,7 +3332,11 @@ class LanceTable(Table):
             tokenizer_configs = self.infer_tokenizer_configs(tokenizer_name)
             tokenizer_configs["custom_stop_words"] = custom_stop_words
 
-        config = FTS(block_size=block_size, **tokenizer_configs)
+        config = FTS(
+            block_size=block_size,
+            document_granularity=document_granularity,
+            **tokenizer_configs,
+        )
 
         try:
             LOOP.run(

python/python/tests/test_first_class_function_slice2.py

@@ -69,6 +69,26 @@ def _run_packaged(definition, *args):
     return namespace[definition.registration_request.artifact.entrypoint](*args)
 
 
+def test_udf_conda_environment():
+    @udf(conda=["scipy", "numpy"], conda_channels=["conda-forge", "defaults"])
+    def halve(value: float) -> float:
+        return value / 2
+
+    request = json.loads(halve.registration_request.to_canonical_json())
+    assert request["runtime"]["environment"] == {
+        "kind": "conda",
+        "packages": ["numpy", "scipy"],
+        "channels": ["conda-forge", "defaults"],
+    }
+    pip_request = json.loads(normalize_score.registration_request.to_canonical_json())
+    assert "channels" not in pip_request["runtime"]["environment"]
+
+    with pytest.raises(ValueError, match="not both"):
+        udf(name="both", pip=["numpy"], conda=["numpy"])(lambda value: value)
+    with pytest.raises(ValueError, match="requires conda"):
+        udf(name="channels", conda_channels=["conda-forge"])(lambda value: value)
+
+
 def test_udf_packages_attribute_access_and_body_imports():
     @udf
     def word_norm(body: str) -> float:

python/python/tests/test_fts.py

@@ -25,6 +25,7 @@ from lancedb.db import DBConnection
 from lancedb.index import FTS
 from lancedb.query import (
     BoostQuery,
+    DocumentGranularity,
     MatchQuery,
     MultiMatchQuery,
     PhraseQuery,
@@ -245,6 +246,55 @@ def test_create_inverted_index_rejects_invalid_block_size(table):
         table.create_index("text", config=FTS(block_size=129))
 
 
+def test_list_element_document_granularity(tmp_path):
+    docs_type = pa.list_(pa.struct([pa.field("content", pa.string())]))
+    docs = pa.array(
+        [
+            [
+                {"content": "alpha beta"},
+                None,
+                {"content": ""},
+                {"content": "the and"},
+                {"content": "alpha beta"},
+            ]
+        ],
+        type=docs_type,
+    )
+    table = ldb.connect(tmp_path).create_table(
+        "list_element_docs", pa.table({"id": [0], "docs": docs})
+    )
+    row_table = ldb.connect(tmp_path).create_table(
+        "row_docs", pa.table({"id": [0], "docs": docs})
+    )
+    row_table.create_index("docs.content", config=FTS())
+    row_result = row_table.search(MatchQuery("alpha", "docs.content")).to_arrow()
+    assert row_result.num_rows == 1
+    assert "_doc_index" not in row_result.column_names
+
+    granularity = DocumentGranularity.LIST_ELEMENT
+    table.create_index(
+        "docs.content",
+        config=FTS(with_position=True, document_granularity=granularity),
+    )
+    assert table.list_indices()[0].columns == ["docs.content"]
+
+    def coordinates(query):
+        result = table.search(query).limit(10).to_arrow()
+        doc_index_type = result.schema.field("_doc_index").type
+        assert pa.types.is_list(doc_index_type)
+        assert doc_index_type.value_type == pa.uint32()
+        return sorted(result["_doc_index"].to_pylist())
+
+    assert coordinates(
+        MatchQuery("alpha", "docs.content", document_granularity=granularity)
+    ) == [[0], [4]]
+    assert coordinates(
+        PhraseQuery("alpha beta", "docs.content", document_granularity=granularity)
+    ) == [[0], [4]]
+    assert coordinates(MatchQuery("alpha", "docs.content")) == [[0], [4]]
+    assert FTS().document_granularity is DocumentGranularity.ROW
+
+
 def test_create_inverted_index_respects_build_memory_limit(table):
     with pytest.raises(ValueError, match="exceeds worker memory limit"):
         table.create_index(
@@ -1089,6 +1139,20 @@ def test_fts_query_to_json():
     )
     assert json_str == expected
 
+    # Test MatchQuery with list-element document granularity
+    match_query = MatchQuery(
+        "hello world",
+        "text",
+        document_granularity=DocumentGranularity.LIST_ELEMENT,
+    )
+    json_str = match_query.to_json()
+    expected = (
+        '{"match":{"column":"text","terms":"hello world","boost":1.0,'
+        '"fuzziness":0,"max_expansions":50,"operator":"Or","prefix_length":0,'
+        '"document_granularity":"list_element"}}'
+    )
+    assert json_str == expected
+
     # Test MatchQuery with options
     match_query = MatchQuery("puppy", "text", fuzziness=2, boost=1.5, prefix_length=3)
     json_str = match_query.to_json()
@@ -1098,6 +1162,19 @@ def test_fts_query_to_json():
     )
     assert json_str == expected
 
+    # Test PhraseQuery with list-element document granularity
+    phrase_query = PhraseQuery(
+        "quick brown fox",
+        "title",
+        document_granularity=DocumentGranularity.LIST_ELEMENT,
+    )
+    json_str = phrase_query.to_json()
+    expected = (
+        '{"phrase":{"column":"title","terms":"quick brown fox","slop":0,'
+        '"document_granularity":"list_element"}}'
+    )
+    assert json_str == expected
+
     # Test PhraseQuery
     phrase_query = PhraseQuery("quick brown fox", "title")
     json_str = phrase_query.to_json()

python/python/tests/test_remote_db.py

@@ -1643,6 +1643,49 @@ def test_query_sync_fts():
         )
 
 
+def test_query_sync_fts_document_granularity():
+    from lancedb.query import DocumentGranularity, MatchQuery
+
+    def handler(body):
+        assert body == {
+            "full_text_query": {
+                "query": {
+                    "match": {
+                        "column": "docs.content",
+                        "terms": "alpha",
+                        "boost": 1.0,
+                        "fuzziness": 0,
+                        "max_expansions": 50,
+                        "operator": "Or",
+                        "prefix_length": 0,
+                        "document_granularity": "list_element",
+                    }
+                }
+            },
+            "k": 10,
+            "prefilter": True,
+            "vector": [],
+            "version": None,
+        }
+        return pa.table(
+            {
+                "id": [1, 1],
+                "_doc_index": pa.array([[0], [4]], type=pa.list_(pa.uint32())),
+            }
+        )
+
+    with query_test_table(handler, server_version=Version("0.6.0")) as table:
+        result = table.search(
+            MatchQuery(
+                "alpha",
+                "docs.content",
+                document_granularity=DocumentGranularity.LIST_ELEMENT,
+            )
+        ).to_arrow()
+
+    assert result["_doc_index"].to_pylist() == [[0], [4]]
+
+
 def test_query_sync_hybrid():
     def handler(body):
         if "full_text_query" in body:

python/python/tests/test_s3.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import copy
+from concurrent.futures import ThreadPoolExecutor
 from datetime import timedelta
 import threading
 
@@ -86,6 +87,25 @@ def test_s3_lifecycle(s3_bucket: str):
     asyncio.run(test())
 
 
+@pytest.mark.s3_test
+def test_concurrent_open_table(s3_bucket: str):
+    uri = f"s3://{s3_bucket}/test_concurrent_open_table"
+    db = lancedb.connect(uri, storage_options=copy.copy(CONFIG))
+    db.create_table("test", pa.table({"x": [1, 2, 3]}))
+
+    num_workers = 32
+    barrier = threading.Barrier(num_workers)
+
+    def open_and_count(_):
+        barrier.wait()
+        return db.open_table("test").count_rows()
+
+    with ThreadPoolExecutor(max_workers=num_workers) as pool:
+        row_counts = list(pool.map(open_and_count, range(num_workers)))
+
+    assert row_counts == [3] * num_workers
+
+
 @pytest.fixture()
 def kms_key():
     kms = get_boto3_client("kms", endpoint_url=CONFIG["aws_endpoint"])

python/src/index.rs

@@ -8,7 +8,7 @@ use lancedb::index::vector::{
 };
 use lancedb::index::{
     Index as LanceDbIndex,
-    scalar::{BTreeIndexBuilder, FmIndexBuilder, FtsIndexBuilder},
+    scalar::{BTreeIndexBuilder, DocumentGranularity, FmIndexBuilder, FtsIndexBuilder},
 };
 use pyo3::IntoPyObject;
 use pyo3::types::PyStringMethods;
@@ -60,7 +60,11 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
                     .ngram_min_length(params.ngram_min_length)
                     .ngram_max_length(params.ngram_max_length)
                     .ngram_prefix_only(params.prefix_only)
-                    .custom_stop_words(params.custom_stop_words);
+                    .custom_stop_words(params.custom_stop_words)
+                    .document_granularity(
+                        DocumentGranularity::try_from(params.document_granularity.as_str())
+                            .map_err(|err| PyValueError::new_err(err.to_string()))?,
+                    );
                 if let Some(memory_limit) = params.memory_limit {
                     inner_opts = inner_opts.memory_limit_mb(memory_limit);
                 }
@@ -221,6 +225,7 @@ struct FtsParams {
     block_size: usize,
     memory_limit: Option<u64>,
     num_workers: Option<usize>,
+    document_granularity: String,
 }
 
 #[derive(FromPyObject)]
@@ -481,6 +486,7 @@ mod tests {
     block_size = 128
     memory_limit = 2048
     num_workers = 7
+    document_granularity = 'row'
 
 config = FTS()",
                 None,

python/src/query.rs

@@ -16,8 +16,8 @@ use arrow::pyarrow::FromPyArrow;
 use arrow::pyarrow::IntoPyArrow;
 use arrow::pyarrow::ToPyArrow;
 use lancedb::index::scalar::{
-    BooleanQuery, BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur,
-    Operator, PhraseQuery,
+    BooleanQuery, BoostQuery, DocumentGranularity, FtsQuery, FullTextSearchQuery, MatchQuery,
+    MultiMatchQuery, Occur, Operator, PhraseQuery,
 };
 use lancedb::query::AnalyzePlanDistributedMetrics;
 use lancedb::query::QueryBase;
@@ -76,8 +76,16 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
                 let max_expansions = ob.getattr("max_expansions")?.extract()?;
                 let operator = ob.getattr("operator")?.extract::<String>()?;
                 let prefix_length = ob.getattr("prefix_length")?.extract()?;
+                let document_granularity = ob
+                    .getattr("document_granularity")?
+                    .extract::<Option<String>>()?
+                    .map(|value| {
+                        DocumentGranularity::try_from(value.as_str())
+                            .map_err(|err| PyValueError::new_err(err.to_string()))
+                    })
+                    .transpose()?;
 
-                Ok(Self(
+                let mut query =
                     MatchQuery::new(query)
                         .with_column(Some(column))
                         .with_boost(boost)
@@ -86,21 +94,32 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
                         .with_operator(Operator::try_from(operator.as_str()).map_err(|e| {
                             PyValueError::new_err(format!("Invalid operator: {}", e))
                         })?)
-                        .with_prefix_length(prefix_length)
-                        .into(),
-                ))
+                        .with_prefix_length(prefix_length);
+                if let Some(document_granularity) = document_granularity {
+                    query = query.with_document_granularity(document_granularity);
+                }
+                Ok(Self(query.into()))
             }
             "PhraseQuery" => {
                 let query = ob.getattr("query")?.extract()?;
                 let column = ob.getattr("column")?.extract()?;
                 let slop = ob.getattr("slop")?.extract()?;
+                let document_granularity = ob
+                    .getattr("document_granularity")?
+                    .extract::<Option<String>>()?
+                    .map(|value| {
+                        DocumentGranularity::try_from(value.as_str())
+                            .map_err(|err| PyValueError::new_err(err.to_string()))
+                    })
+                    .transpose()?;
 
-                Ok(Self(
-                    PhraseQuery::new(query)
-                        .with_column(Some(column))
-                        .with_slop(slop)
-                        .into(),
-                ))
+                let mut query = PhraseQuery::new(query)
+                    .with_column(Some(column))
+                    .with_slop(slop);
+                if let Some(document_granularity) = document_granularity {
+                    query = query.with_document_granularity(document_granularity);
+                }
+                Ok(Self(query.into()))
             }
             "BoostQuery" => {
                 let positive: Self = ob.getattr("positive")?.extract()?;
@@ -167,6 +186,13 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
                 kwargs.set_item("max_expansions", query.max_expansions)?;
                 kwargs.set_item::<_, &str>("operator", query.operator.into())?;
                 kwargs.set_item("prefix_length", query.prefix_length)?;
+                if let Some(document_granularity) = query.document_granularity {
+                    let value = match document_granularity {
+                        DocumentGranularity::Row => "row",
+                        DocumentGranularity::ListElement => "list_element",
+                    };
+                    kwargs.set_item("document_granularity", value)?;
+                }
                 namespace
                     .getattr(intern!(py, "MatchQuery"))?
                     .call((query.terms, query.column.unwrap()), Some(&kwargs))
@@ -174,6 +200,13 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
             FtsQuery::Phrase(query) => {
                 let kwargs = PyDict::new(py);
                 kwargs.set_item("slop", query.slop)?;
+                if let Some(document_granularity) = query.document_granularity {
+                    let value = match document_granularity {
+                        DocumentGranularity::Row => "row",
+                        DocumentGranularity::ListElement => "list_element",
+                    };
+                    kwargs.set_item("document_granularity", value)?;
+                }
                 namespace
                     .getattr(intern!(py, "PhraseQuery"))?
                     .call((query.terms, query.column.unwrap()), Some(&kwargs))