diff --git a/.bumpversion.toml b/.bumpversion.toml index decfe1972..70718edbb 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.41.0-beta.0" +current_version = "0.41.0-beta.2" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/Cargo.lock b/Cargo.lock index 58b4c6b84..8e44b93bd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3588,8 +3588,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-array", "rand 0.9.5", @@ -5095,26 +5095,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "jsonb" -version = "0.5.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb98fb29636087c40ad0d1274d9a30c0c1e83e03ae93f6e7e89247b37fcc6953" -dependencies = [ - "byteorder", - "ethnum", - "fast-float2", - "itoa", - "jiff", - "nom 8.0.0", - "num-traits", - "ordered-float 5.3.0", - "rand 0.9.5", - "serde", - "serde_json", - "zmij", -] - [[package]] name = "kanaria" version = "0.2.0" @@ -5143,8 +5123,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arc-swap", "arrow", @@ -5217,7 +5197,7 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "14.0.0-beta.7" +version = "14.0.0-beta.10" dependencies = [ "arrow-array", "arrow-buffer", @@ -5231,7 +5211,7 @@ dependencies = [ "futures", "getrandom 0.2.17", "half", - "jsonb", + "lance-jsonb", "num-traits", "rand 0.9.5", ] @@ -5239,7 +5219,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-array", "arrow-buffer", @@ -5253,7 +5233,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-array", "arrow-schema", @@ -5262,8 +5242,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrayref", "crunchy", @@ -5273,8 +5253,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-array", "arrow-buffer", @@ -5311,8 +5291,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow", "arrow-array", @@ -5329,10 +5309,10 @@ dependencies = [ "datafusion-physical-expr", "futures", "half", - "jsonb", "lance-arrow", "lance-core", "lance-geo", + "lance-jsonb", "log", "pin-project", "prost", @@ -5343,8 +5323,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow", "arrow-array", @@ -5361,8 +5341,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "proc-macro2", "quote", @@ -5371,8 +5351,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-arith", "arrow-array", @@ -5405,8 +5385,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-arith", "arrow-array", @@ -5438,8 +5418,8 @@ dependencies = [ [[package]] name = "lance-geo" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "datafusion", "geo-traits", @@ -5453,8 +5433,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arc-swap", "arrow", @@ -5485,7 +5465,6 @@ dependencies = [ "hex", "itertools 0.14.0", "jieba-rs", - "jsonb", "lance-arrow", "lance-arrow-stats", "lance-bitpacking", @@ -5496,6 +5475,7 @@ dependencies = [ "lance-geo", "lance-index-core", "lance-io", + "lance-jsonb", "lance-linalg", "lance-select", "lance-table", @@ -5525,8 +5505,8 @@ dependencies = [ [[package]] name = "lance-index-core" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-array", "arrow-schema", @@ -5549,8 +5529,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow", "arrow-array", @@ -5590,10 +5570,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "lance-jsonb" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" +dependencies = [ + "byteorder", + "ethnum", + "fast-float2", + "itoa", + "jiff", + "nom 8.0.0", + "num-traits", + "ordered-float 5.3.0", + "serde", + "serde_json", + "zmij", +] + [[package]] name = "lance-linalg" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-array", "arrow-schema", @@ -5607,8 +5605,8 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow", "async-trait", @@ -5622,8 +5620,8 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arc-swap", "arrow", @@ -5677,8 +5675,8 @@ dependencies = [ [[package]] name = "lance-select" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-array", "arrow-buffer", @@ -5692,8 +5690,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow", "arrow-array", @@ -5733,8 +5731,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "arrow-array", "arrow-schema", @@ -5747,8 +5745,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.10" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.10#fa3638fd816ab2ff9a2ecbba7e51465d8879958f" dependencies = [ "frostem", "icu_segmenter", @@ -5761,7 +5759,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.41.0-beta.0" +version = "0.41.0-beta.2" dependencies = [ "ahash", "anyhow", @@ -5858,7 +5856,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.41.0-beta.0" +version = "0.41.0-beta.2" dependencies = [ "arrow-array", "arrow-buffer", @@ -5883,7 +5881,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.41.0-beta.0" +version = "0.41.0-beta.2" dependencies = [ "arc-swap", "arrow", diff --git a/Cargo.toml b/Cargo.toml index 401fcf257..8f9441315 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,21 @@ categories = ["database-implementations"] rust-version = "1.94.0" [workspace.dependencies] -lance = { "version" = "=14.0.0-beta.7", default-features = false, "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=14.0.0-beta.7", default-features = false, "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=14.0.0-beta.7", default-features = false, "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { version = "=14.0.0-beta.7", path = "python/vendor/lance-arrow" } +lance = { "version" = "=14.0.0-beta.10", default-features = false, "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=14.0.0-beta.10", default-features = false, "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-jsonb = { version = "=14.0.0-beta.10", tag = "v14.0.0-beta.10", git = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=14.0.0-beta.10", default-features = false, "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=14.0.0-beta.10", "tag" = "v14.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { version = "=14.0.0-beta.10", path = "python/vendor/lance-arrow" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 4594b6dd6..88cfc4a30 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.41.0-beta.0 + 0.41.0-beta.2 ``` diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index a0a77fd7c..3f56bb600 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -460,17 +460,16 @@ exhaustiely searched to find matches. This parameter controls how many partitions should be searched. Increasing this value will increase the recall of your query but will -also increase the latency of your query. The default value is 20. This -default is good for many cases but the best value to use will depend on -your data and the recall that you need to achieve. +also increase the latency of your query. If this method is not called, +Lance's adaptive probe defaults are used. For best results we recommend tuning this parameter with a benchmark against your actual data to find the smallest possible value that will still give you the desired recall. -For more fine grained control over behavior when you have a very narrow filter -you can use `minimumNprobes` and `maximumNprobes`. This method sets both -the minimum and maximum to the same value. +The value is retained as `nprobes` through client and server request +construction. Lance sets both probe bounds to this value. Explicit minimum +or maximum settings can override their respective bounds. #### Parameters diff --git a/docs/src/js/interfaces/FtsOptions.md b/docs/src/js/interfaces/FtsOptions.md index 1d3c9ff82..654fa138f 100644 --- a/docs/src/js/interfaces/FtsOptions.md +++ b/docs/src/js/interfaces/FtsOptions.md @@ -37,6 +37,8 @@ The following tokenizers are available: "raw" - Raw tokenizer. This tokenizer does not split the text into tokens and indexes the entire text as a single token. +"code" - Code tokenizer. This tokenizer keeps snake_case and camelCase identifiers as single tokens. + "icu" - ICU dictionary-based word segmentation. "icu/split" - ICU segmentation with simple-style delimiter splitting. diff --git a/docs/src/js/interfaces/ListTablesOptions.md b/docs/src/js/interfaces/ListTablesOptions.md index 52ace47e5..b8dd8b5ef 100644 --- a/docs/src/js/interfaces/ListTablesOptions.md +++ b/docs/src/js/interfaces/ListTablesOptions.md @@ -15,6 +15,7 @@ optional limit: number; ``` An upper bound on how many tables to return. +Defaults to 100. Zero returns an empty page without a continuation token. A page may hold fewer than this and still not be the last one, so keep going while the response carries a page token rather than while pages are diff --git a/docs/src/js/interfaces/TableNamesOptions.md b/docs/src/js/interfaces/TableNamesOptions.md index 9254e9fa7..123aa6c9f 100644 --- a/docs/src/js/interfaces/TableNamesOptions.md +++ b/docs/src/js/interfaces/TableNamesOptions.md @@ -19,7 +19,7 @@ instead. optional limit: number; ``` -An optional limit to the number of results to return. +The maximum number of names to return. Omitted returns all names; zero returns none. *** diff --git a/docs/src/js/type-aliases/BaseTokenizer.md b/docs/src/js/type-aliases/BaseTokenizer.md index 75b377f27..5fd1b55d4 100644 --- a/docs/src/js/type-aliases/BaseTokenizer.md +++ b/docs/src/js/type-aliases/BaseTokenizer.md @@ -12,6 +12,7 @@ type BaseTokenizer: | "whitespace" | "raw" | "ngram" + | "code" | "icu" | "icu/split" | `jieba/${string}` diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 863edbe4a..3776283dc 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -86,6 +86,15 @@ print(query.describe().status) for batch in query.reader(): print(batch.num_rows) +# Values bind to $1 / $name placeholders and travel as Arrow, so a float32 +# stays a float32 and a vector stays a compact fixed-size list: +import numpy as np + +reader = db.execute_query( + "SELECT id FROM docs ORDER BY distance(vector, $vector) LIMIT $k", + parameters={"vector": np.random.rand(768).astype(np.float32), "k": 10}, +) + # The async connection exposes the same lifecycle without blocking: # async_db = await lancedb.connect_async( # "db://analytics", @@ -201,6 +210,8 @@ listing a storage directory. ::: lancedb.sql.QueryDescription +::: lancedb.sql.QueryParameters + ## Materialized Views (Synchronous) ::: lancedb.materialized_view.MaterializedView diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index 3cd3a970a..b10da3f6f 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -139,9 +139,9 @@ Learn how to configure, build, and search LanceDB vector indexes, including buil | Parameter | Description | | :--- | :--- | | `limit` | Number of results to return (the `k` in `k-ANN`). | -| `nprobes` | Shorthand that sets both `minimum_nprobes` and `maximum_nprobes` to the same value. LanceDB defaults this to 20 which is a good default for many cases. At very high scales you will likely need to experiment to find your preferred recall / latency balance. | -| `minimum_nprobes` | Minimum number of partitions scanned. | -| `maximum_nprobes` | Maximum number of partitions scanned. Only scans more than `minimum_nprobes` if an initial pass does not return enough results — useful for narrow filters. Set to `0` to remove the cap. | +| `nprobes` | Sets both `minimum_nprobes` and `maximum_nprobes` to the same value in Lance. Leave unset to use Lance's adaptive defaults. At very high scales you will likely need to experiment to find your preferred recall / latency balance. | +| `minimum_nprobes` | Minimum number of partitions scanned. Leave unset to use Lance's default. | +| `maximum_nprobes` | Maximum number of partitions scanned. Only scans more than `minimum_nprobes` if an initial pass does not return enough results — useful for narrow filters. Leave unset, or set to `0`, to remove the cap. | | `ef` | HNSW search-time exploration factor. Start around `1.5 * k` and increase up to `10 * k` for higher recall. | | `refine_factor` | Reads and reranks additional candidates in memory to recover recall lost to quantization. Start between `5` and `50` (a factor of `25` reads and reranks the top `25 * k` candidates) and measure the latency cost, typically a few milliseconds, against the recall gained. | | `distance_range(lower_bound, upper_bound)` | Return only rows whose distance falls within `[lower_bound, upper_bound)`. Either bound is optional. Useful for near-duplicate detection or "close-enough" matching. | @@ -165,8 +165,8 @@ older client. | Index type | Guidance | | :--- | :--- | -| `IVF_HNSW_FLAT`, `IVF_HNSW_SQ` | Keep the default `nprobes`, then tune `ef` first. Expect higher latency variance under filtered search. | -| `IVF_RQ`, `IVF_PQ` | Keep the default `nprobes` and raise when recall is insufficient. If using a selective filter you may need to increase `maximum_nprobes` if you do not received enough results. | +| `IVF_HNSW_FLAT`, `IVF_HNSW_SQ` | Leave the probe settings unset, then tune `ef` first. Expect higher latency variance under filtered search. | +| `IVF_RQ`, `IVF_PQ` | Leave the probe settings unset and configure them only when recall is insufficient. If using a selective filter you may need to increase `maximum_nprobes` if you do not receive enough results. | diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index 922bac384..987d51345 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.0 + 0.41.0-beta.2 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 5f21306f4..964dfed0c 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.0 + 0.41.0-beta.2 pom ${project.artifactId} LanceDB Java SDK Parent POM @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 14.0.0-beta.7 + 14.0.0-beta.10 false 2.30.0 1.7 diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 74749b79b..c81fdb459 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.41.0-beta.0" +version = "0.41.0-beta.2" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index 816f8c3de..895c2a7c5 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -import { readdirSync } from "fs"; +import { mkdirSync, readdirSync } from "fs"; +import { join } from "path"; import { Field, Float64, Schema } from "apache-arrow"; import * as tmp from "tmp"; import { @@ -180,6 +181,32 @@ describe("given a connection", () => { expect(seen).toEqual(created); }); + it("should list all names and use a default page size of 100", async () => { + const names = Array.from( + { length: 130 }, + (_, i) => `t${i.toString().padStart(3, "0")}`, + ); + for (const name of names) { + mkdirSync(join(tmpDir.name, `${name}.lance`)); + } + await expect(db.tableNames()).resolves.toEqual(names); + await expect(db.tableNames({ limit: 0 })).resolves.toEqual([]); + const first = await db.listTables(); + expect(first.tables).toEqual(names.slice(0, 100)); + expect(first.pageToken).toBeDefined(); + const second = await db.listTables({ pageToken: first.pageToken }); + expect(second.tables).toEqual(names.slice(100)); + expect(second.pageToken).toBeUndefined(); + await expect(db.listTables({ limit: 200 })).resolves.toEqual({ + tables: names, + pageToken: undefined, + }); + await expect(db.listTables({ limit: 0 })).resolves.toEqual({ + tables: [], + pageToken: undefined, + }); + }); + it("should list tables in a namespace", async () => { const db = await connect(tmpDir.name, { // biome-ignore lint/style/useNamingConvention: opaque backend property key, must match Rust diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 3f2d4a436..6911b24b9 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -28,6 +28,7 @@ import { StaticHeaderProvider, } from "../lancedb/header"; import { Index } from "../lancedb/indices"; +import { LocalTable } from "../lancedb/table"; // Test-only header providers class CustomProvider extends HeaderProvider { @@ -93,6 +94,58 @@ async function withMockDatabase( } describe("remote connection", () => { + it("reports unsupported operations with backend-neutral errors", async () => { + await withMockDatabase( + (req, res) => { + expect(req.url).toBe("/v1/table/test/describe/"); + res.writeHead(200, { "Content-Type": "application/json" }).end( + JSON.stringify({ + version: 1, + schema: { + fields: [ + { name: "id", type: { type: "int64" }, nullable: false }, + { + name: "vector", + type: { + type: "fixed_size_list", + fields: [ + { name: "item", type: { type: "float" }, nullable: true }, + ], + length: 2, + }, + nullable: false, + }, + ], + }, + }), + ); + }, + async (db) => { + const table = (await db.openTable("test")) as LocalTable; + const cases: [() => Promise, string][] = [ + [() => table.usesV2ManifestPaths(), "uses_v2_manifest_paths"], + [() => table.migrateManifestPathsV2(), "migrate_manifest_paths_v2"], + [ + () => table.setUnenforcedPrimaryKey("id"), + "set_unenforced_primary_key", + ], + [() => table.optimize(), "optimize"], + [ + () => table.createIndex("vector", { config: Index.hnswPq() }), + "IVF_HNSW_PQ", + ], + ]; + for (const [run, operation] of cases) { + let message = `LanceDBError: not supported: ${operation} is not supported for remote tables.`; + if (operation === "IVF_HNSW_PQ") { + message += " Please use IVF_HNSW_SQ instead."; + } + await expect(run()).rejects.toThrow(new Error(message)); + } + }, + ); + }); + it("rejects nprobes(0) before sending a query", async () => { const requests: string[] = []; await withMockDatabase( diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 5a27b852b..d2244b48c 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -1231,7 +1231,9 @@ describe("When creating an index", () => { expect(() => tbl.query().nearestTo(queryVec).minimumNprobes(0)).toThrow( "Invalid input, minimum_nprobes must be greater than 0", ); - expect(() => tbl.query().nearestTo(queryVec).maximumNprobes(5)).toThrow( + expect(() => + tbl.query().nearestTo(queryVec).minimumNprobes(10).maximumNprobes(5), + ).toThrow( "Invalid input, maximum_nprobes must be greater than or equal to minimum_nprobes", ); @@ -4050,6 +4052,28 @@ describe.each(arrowVersions)( ); }); + test.each(["Klingon", "english"])( + "rejects unsupported full text language %s with a catchable error", + (language) => { + expect(() => Index.fts({ language })).toThrow( + new Error( + `LanceDB does not support the requested language: '${language}'`, + ), + ); + }, + ); + + test("full text search with a supported language", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable("test", [{ text: "running" }]); + await table.createIndex("text", { + config: Index.fts({ language: "English", stem: true }), + }); + + const results = await table.search("run").toArray(); + expect(results.map((row) => row.text)).toEqual(["running"]); + }); + test("full text search without lowercase", async () => { const db = await connect(tmpDir.name); const data = [ @@ -4180,6 +4204,20 @@ describe.each(arrowVersions)( expect(mustNotResults.length).toBe(1); }); + test("full text search code tokenizer", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable("test", [ + { id: 1, text: "def getUserName(user_id): return user_id" }, + { id: 2, text: "def parseFile(file_path): return file_path" }, + ]); + await table.createIndex("text", { + config: Index.fts({ baseTokenizer: "code" }), + }); + + const results = await table.search("getUserName", "fts").toArray(); + expect(results.map((row) => row.id)).toEqual([1]); + }); + test("full text search ngram", async () => { const db = await connect(tmpDir.name); const data = [ diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 5ba531eed..11341a8c6 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -150,7 +150,7 @@ export interface TableNamesOptions { * the last table name from the previous page. */ startAfter?: string; - /** An optional limit to the number of results to return. */ + /** The maximum number of names to return. Omitted returns all names; zero returns none. */ limit?: number; } @@ -164,6 +164,7 @@ export interface ListTablesOptions { pageToken?: string; /** * An upper bound on how many tables to return. + * Defaults to 100. Zero returns an empty page without a continuation token. * * A page may hold fewer than this and still not be the last one, so keep * going while the response carries a page token rather than while pages are diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index bc9e19300..43d1cf1b5 100644 --- a/nodejs/lancedb/indices.ts +++ b/nodejs/lancedb/indices.ts @@ -491,6 +491,7 @@ export type BaseTokenizer = | "whitespace" | "raw" | "ngram" + | "code" | "icu" | "icu/split" | `jieba/${string}` @@ -520,6 +521,8 @@ export interface FtsOptions { * * "raw" - Raw tokenizer. This tokenizer does not split the text into tokens and indexes the entire text as a single token. * + * "code" - Code tokenizer. This tokenizer keeps snake_case and camelCase identifiers as single tokens. + * * "icu" - ICU dictionary-based word segmentation. * * "icu/split" - ICU segmentation with simple-style delimiter splitting. diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index f370e1b31..0f0eaf610 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -532,17 +532,16 @@ export class VectorQuery extends StandardQueryBase { * partitions should be searched. * * Increasing this value will increase the recall of your query but will - * also increase the latency of your query. The default value is 20. This - * default is good for many cases but the best value to use will depend on - * your data and the recall that you need to achieve. + * also increase the latency of your query. If this method is not called, + * Lance's adaptive probe defaults are used. * * For best results we recommend tuning this parameter with a benchmark against * your actual data to find the smallest possible value that will still give * you the desired recall. * - * For more fine grained control over behavior when you have a very narrow filter - * you can use `minimumNprobes` and `maximumNprobes`. This method sets both - * the minimum and maximum to the same value. + * The value is retained as `nprobes` through client and server request + * construction. Lance sets both probe bounds to this value. Explicit minimum + * or maximum settings can override their respective bounds. */ nprobes(nprobes: number): VectorQuery { this.doVectorCall((inner) => inner.nprobes(nprobes)); diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 9686e2bf3..90106b6d4 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.41.0-beta.0", + "version": "0.41.0-beta.2", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index 2b9f5c27a..73bafc124 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.41.0-beta.0", + "version": "0.41.0-beta.2", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index bdd01b6a0..dea0f49a2 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.41.0-beta.0", + "version": "0.41.0-beta.2", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index b38c05b2e..0ae15f6d5 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.41.0-beta.0", + "version": "0.41.0-beta.2", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index 59405732b..b3a17754b 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.41.0-beta.0", + "version": "0.41.0-beta.2", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index 7c5e6b900..4681f63c9 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.41.0-beta.0", + "version": "0.41.0-beta.2", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 8f60e7eea..bcf00b687 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.41.0-beta.0", + "version": "0.41.0-beta.2", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 8d85e501a..b4b64eec1 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.41.0-beta.0", + "version": "0.41.0-beta.2", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/nodejs/src/index.rs b/nodejs/src/index.rs index 762f74870..e5884aa4d 100644 --- a/nodejs/src/index.rs +++ b/nodejs/src/index.rs @@ -239,7 +239,12 @@ impl Index { opts = opts.base_tokenizer(base_tokenizer); } if let Some(language) = language { - opts = opts.language(&language).unwrap(); + opts = opts.language(&language).map_err(|_| { + napi::Error::from_reason(format!( + "LanceDB does not support the requested language: '{}'", + language + )) + })?; } if let Some(max_token_length) = max_token_length { opts = opts.max_token_length(Some(max_token_length as usize)); diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index ddf5be9f4..32cdcce3c 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -836,7 +836,10 @@ impl Table { pub async fn uses_v2_manifest_paths(&self) -> napi::Result { self.inner_ref()? .as_native() - .ok_or_else(|| napi::Error::from_reason("This cannot be run on a remote table"))? + .ok_or_else(|| lancedb::Error::NotSupported { + message: "uses_v2_manifest_paths is not supported for remote tables.".into(), + }) + .default_error()? .uses_v2_manifest_paths() .await .default_error() @@ -846,7 +849,10 @@ impl Table { pub async fn migrate_manifest_paths_v2(&self) -> napi::Result<()> { self.inner_ref()? .as_native() - .ok_or_else(|| napi::Error::from_reason("This cannot be run on a remote table"))? + .ok_or_else(|| lancedb::Error::NotSupported { + message: "migrate_manifest_paths_v2 is not supported for remote tables.".into(), + }) + .default_error()? .migrate_manifest_paths_v2() .await .default_error() diff --git a/python/Cargo.toml b/python/Cargo.toml index 28aac5fba..c4cdb7bbc 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.41.0-beta.0" +version = "0.41.0-beta.2" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/python/pyproject.toml b/python/pyproject.toml index 2d753c139..cde7247f3 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "pyarrow>=16", "pydantic>=2.7.4,<3", "tqdm>=4.27.0", - "lance-namespace>=0.3.2" + "lance-namespace>=0.13.0" ] description = "lancedb" authors = [{ name = "LanceDB Devs", email = "dev@lancedb.com" }] @@ -63,7 +63,7 @@ tests = [ "polars>=0.19, <=1.32.3", "pyarrow<25", "pyarrow-stubs>=16.0", - "pylance==14.0.0b6", + "pylance==14.0.0b10", "requests>=2.31.0", "datafusion>=54,<55", "opentelemetry-sdk>=1.30.0", diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 213794ef3..39353f270 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -321,8 +321,8 @@ def tokenize( language: str = "English", max_token_length: Optional[int] = 40, lower_case: bool = True, - stem: bool = True, - remove_stop_words: bool = True, + stem: Optional[bool] = None, + remove_stop_words: Optional[bool] = None, custom_stop_words: Optional[List[str]] = None, ascii_folding: bool = True, ngram_min_length: int = 3, @@ -333,6 +333,8 @@ def tokenize( This does not require an FTS index. The tokenizer options match :class:`lancedb.index.FTS`. ``custom_stop_words`` accepts a list of strings. + When ``stem`` or ``remove_stop_words`` is None, the base tokenizer's default + applies: False for ``code`` and ``ngram``, True otherwise. """ return _tokenize( diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index f7477e891..cfaa8c68a 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -58,8 +58,8 @@ def tokenize( language: str = "English", max_token_length: Optional[int] = 40, lower_case: bool = True, - stem: bool = True, - remove_stop_words: bool = True, + stem: Optional[bool] = None, + remove_stop_words: Optional[bool] = None, custom_stop_words: Optional[List[str]] = None, ascii_folding: bool = True, ngram_min_length: int = 3, @@ -217,6 +217,7 @@ class Connection(object): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[pa.RecordBatch] = None, ) -> SqlQuery: ... async def describe_query(self, query_id: UUID) -> QueryDescription: ... async def create_table( @@ -422,6 +423,7 @@ class Table: def name(self) -> str: ... def __repr__(self) -> str: ... def is_open(self) -> bool: ... + def _is_native(self) -> bool: ... def close(self) -> None: ... async def schema(self) -> pa.Schema: ... async def add( @@ -776,6 +778,7 @@ class PyQueryRequest: use_lsm: Optional[bool] column: Optional[str] query_vector: Optional[List[pa.Array]] + nprobes: Optional[int] minimum_nprobes: Optional[int] maximum_nprobes: Optional[int] lower_bound: Optional[float] diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index f890c5a14..47d24b762 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -53,7 +53,7 @@ from .functions import FunctionVersion, UdfDefinition from .job import AsyncJob, Job, _typed_job from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery -from .sql import QueryDescription +from .sql import QueryDescription, QueryParameters, to_parameter_batch from .materialized_view import ( AsyncMaterializedView, MaterializedView, @@ -317,7 +317,7 @@ class DBConnection(EnforceOverrides): page_token: Optional[str] = None, limit: Optional[int] = None, ) -> ListTablesResponse: - """List all tables in this database with pagination support. + """List a page of tables in this database. Parameters ---------- @@ -328,7 +328,8 @@ class DBConnection(EnforceOverrides): Token for pagination. Use the token from a previous response to get the next page of results. limit: int, optional - The maximum number of results to return. + The maximum number of results to return, default 100. + Zero returns an empty page without a continuation token. Returns ------- @@ -343,7 +344,7 @@ class DBConnection(EnforceOverrides): def table_names( self, page_token: Optional[str] = None, - limit: int = 10, + limit: Optional[int] = None, *, namespace_path: Optional[List[str]] = None, ) -> Iterable[str]: @@ -357,8 +358,9 @@ class DBConnection(EnforceOverrides): page_token: str, optional The token to use for pagination. If not present, start from the beginning. Typically, this token is last table name from the previous page. - limit: int, default 10 - The size of the page to return. + limit: int, optional + The maximum number of tables to return. None returns all tables; + zero returns an empty list. Returns ------- @@ -1178,6 +1180,7 @@ class DBConnection(EnforceOverrides): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> pa.RecordBatchReader: """Execute SQL and return a blocking Arrow reader. @@ -1188,6 +1191,7 @@ class DBConnection(EnforceOverrides): return self.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=parameters, ).reader() def execute_query_async( @@ -1195,9 +1199,32 @@ class DBConnection(EnforceOverrides): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> SqlQuery: """Start executing SQL and return its query handle. + ``parameters`` binds values to the statement's ``$1`` / ``$name`` + placeholders. They travel as Arrow rather than as SQL text, so a + float keeps its exact value and type and a vector stays a compact + fixed-size list. A list binds by position, a dict binds by name, and + a one-row ``pyarrow.RecordBatch`` or ``pyarrow.Table`` binds both + ways. A 1-D numpy or pyarrow array becomes a fixed-size list of its + element type, the shape of a vector column: + + .. code-block:: python + + db.execute_query( + "SELECT id FROM docs ORDER BY distance(vector, $1) LIMIT $2", + parameters=[np.array([0.1, 0.2, 0.3], dtype=np.float32), 10], + ) + + Parameters are supported in queries, not in DDL or DML. A + parameterized query runs on the call that returns its rows: + cancelling it, or dropping its handle before reading, stops it on the + server. Its status follows that stream as well: ``describe()`` + reports ``running`` until the reader has received every row, so open + ``reader()`` directly rather than waiting for ``finished``. + Local connections do not support SQL. """ raise NotImplementedError("SQL is not supported for this connection type") @@ -1465,7 +1492,7 @@ class LanceDBConnection(DBConnection): page_token: Optional[str] = None, limit: Optional[int] = None, ) -> ListTablesResponse: - """List all tables in this database with pagination support. + """List a page of tables in this database. Parameters ---------- @@ -1476,7 +1503,8 @@ class LanceDBConnection(DBConnection): Token for pagination. Use the token from a previous response to get the next page of results. limit: int, optional - The maximum number of results to return. + The maximum number of results to return, default 100. + Zero returns an empty page without a continuation token. Returns ------- @@ -1495,7 +1523,7 @@ class LanceDBConnection(DBConnection): def table_names( self, page_token: Optional[str] = None, - limit: int = 10, + limit: Optional[int] = None, *, namespace_path: Optional[List[str]] = None, ) -> Iterable[str]: @@ -1510,8 +1538,9 @@ class LanceDBConnection(DBConnection): The namespace to list tables in. page_token: str, optional The token to use for pagination. - limit: int, default 10 - The maximum number of tables to return. + limit: int, optional + The maximum number of tables to return. None returns all tables; + zero returns an empty list. Returns ------- @@ -2341,7 +2370,7 @@ class AsyncConnection(object): page_token: Optional[str] = None, limit: Optional[int] = None, ) -> ListTablesResponse: - """List all tables in this database with pagination support. + """List a page of tables in this database. Parameters ---------- @@ -2352,7 +2381,8 @@ class AsyncConnection(object): Token for pagination. Use the token from a previous response to get the next page of results. limit: int, optional - The maximum number of results to return. + The maximum number of results to return, default 100. + Zero returns an empty page without a continuation token. Returns ------- @@ -2389,8 +2419,9 @@ class AsyncConnection(object): This can be combined with limit to implement pagination by setting this to the last table name from the previous page. - limit: int, default 10 - The number of results to return. + limit: int, optional + The maximum number of tables to return. None returns all tables; + zero returns an empty list. Returns ------- @@ -3267,6 +3298,7 @@ class AsyncConnection(object): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> AsyncRecordBatchReader: """Execute SQL and return an asynchronous Arrow reader. @@ -3277,6 +3309,7 @@ class AsyncConnection(object): submitted = await self.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=parameters, ) return await submitted.reader() @@ -3285,17 +3318,21 @@ class AsyncConnection(object): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> AsyncSqlQuery: """Start executing SQL and return its query handle. The database from ``connect_async`` is used for unqualified database - references. The namespace defaults to ``["public"]``. Local - connections raise ``NotImplementedError``. + references. The namespace defaults to ``["public"]``. ``parameters`` + binds values to ``$1`` / ``$name`` placeholders; see + [DBConnection.execute_query_async][lancedb.db.DBConnection.execute_query_async]. + Local connections raise ``NotImplementedError``. """ return AsyncSqlQuery( await self._inner.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=to_parameter_batch(parameters), ) ) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index c170c7375..445d31027 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -372,7 +372,7 @@ class FunctionVersion(_RemoteValue): image: FunctionImage signature: FunctionSignature secret_bindings: tuple[SecretBinding, ...] = () - created_at: str + created_at_millis: int metadata: Mapping[str, str] disabled: bool diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index 78ee72e57..cf3bf8e4f 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -112,9 +112,42 @@ class Fm: class FTS: """Describe a FTS index configuration. - `FTS` is a full-text search index that can be used on `String` columns + `FTS` supports string columns, lists of strings, and JSON columns created + from PyArrow ``pa.json_()``. JSON columns use Lance's JSON-aware tokenizer + and are stored as ``LargeBinary`` with the ``lance.json`` logical type. + Arbitrary binary columns are not supported. - For example, it works with `title`, `description`, `content`, etc. + JSON queries use ``path,str,value`` triples. Paths are relative to the JSON + column and use dots for nested object keys. For example, + ``meta.tag,str,nature`` matches the string value at ``$.meta.tag``; + it does not search other paths or JSON key names. The value is processed + with the configured text tokenizer. Use ``str``, not ``string``. + + Create and query a JSON index through the modern API:: + + import pyarrow as pa + from lancedb.index import FTS + + table = db.create_table( + "documents", + pa.table({"doc": pa.array( + ['{"meta":{"tag":"nature"}}'], type=pa.json_() + )}), + ) + table.create_index("doc", config=FTS()) + table.search( + "meta.tag,str,nature", query_type="fts", fts_columns="doc" + ).to_list() + + Create the index on the JSON column itself, not a JSON subpath such as + ``doc.meta.tag``. JSON-path scalar indexes do not provide FTS path scoping. + Plain text such as ``nature`` is not a valid JSON FTS query. The current + Lance tokenizer can panic on malformed triples; see + https://github.com/lance-format/lance/issues/8812. + + Remote tables use the same ``FTS()`` configuration and require a server + that supports JSON FTS. Wait for ``table.wait_for_index(["doc_idx"])`` + before searching a newly created remote index with its default name. Examples -------- @@ -138,6 +171,7 @@ class FTS: - "simple": Splits text by whitespace and punctuation. - "whitespace": Split text by whitespace, but not punctuation. - "raw": No tokenization. The entire text is treated as a single token. + - "code": Tokenizes source code and identifiers. - "ngram": N-gram tokenizer for substring-style matching. - "icu": ICU dictionary-based word segmentation. - "icu/split": ICU segmentation with simple-style delimiter splitting. @@ -151,12 +185,16 @@ class FTS: ignored. lower_case : bool, default True Whether to convert the token to lower case. This makes queries case-insensitive. - stem : bool, default True + stem : bool, optional Whether to stem the token. Stemming reduces words to their root form. For example, in English "running" and "runs" would both be reduced to "run". - remove_stop_words : bool, default True + ``None`` uses the base tokenizer's default: False for ``code`` and + ``ngram``, True otherwise. + remove_stop_words : bool, optional Whether to remove stop words. Stop words are common words that are often removed from text before indexing. For example, in English "the" and "and". + ``None`` uses the base tokenizer's default: False for ``code`` and + ``ngram``, True otherwise. custom_stop_words : list of str, optional Custom words replace the built-in language stop words and only take effect when ``remove_stop_words`` is True. ``None`` uses @@ -197,8 +235,8 @@ class FTS: language: str = "English" max_token_length: Optional[int] = 40 lower_case: bool = True - stem: bool = True - remove_stop_words: bool = True + stem: Optional[bool] = None + remove_stop_words: Optional[bool] = None ascii_folding: bool = True ngram_min_length: int = 3 ngram_max_length: int = 3 diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index a3571f64b..ed87e441f 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -53,7 +53,7 @@ from lancedb.db import AsyncConnection, DBConnection from lancedb.job import AsyncJob, Job from lancedb.listing import AsyncListing, Listing from lancedb.sql import AsyncQuery as AsyncSqlQuery -from lancedb.sql import QueryDescription +from lancedb.sql import QueryDescription, QueryParameters from lance_namespace import ( LanceNamespace, connect as namespace_connect, @@ -190,8 +190,12 @@ def _query_to_namespace_request( kwargs["with_row_id"] = query.with_row_id if query.ef is not None: kwargs["ef"] = query.ef + if query.nprobes is not None: + kwargs["nprobes"] = query.nprobes if query.minimum_nprobes is not None: - kwargs["nprobes"] = query.minimum_nprobes + kwargs["minimum_nprobes"] = query.minimum_nprobes + if query.maximum_nprobes is not None: + kwargs["maximum_nprobes"] = query.maximum_nprobes if query.refine_factor is not None: kwargs["refine_factor"] = query.refine_factor if query.lower_bound is not None: @@ -205,6 +209,13 @@ def _query_to_namespace_request( if query.distance_type is not None: kwargs["distance_type"] = query.distance_type + if kwargs.get("maximum_nprobes") == 0: + # The generated 0.13 model rejects LanceDB's unbounded sentinel; + # model_copy preserves direct wire pass-through to the Lance receiver. + maximum_nprobes = kwargs.pop("maximum_nprobes") + return QueryTableRequest(**kwargs).model_copy( + update={"maximum_nprobes": maximum_nprobes} + ) return QueryTableRequest(**kwargs) @@ -522,13 +533,15 @@ class LanceNamespaceDBConnection(DBConnection): def table_names( self, page_token: Optional[str] = None, - limit: int = 10, + limit: Optional[int] = None, *, namespace_path: Optional[List[str]] = None, ) -> Iterable[str]: """ List table names in the database. + If limit is None, return all table names; zero returns an empty list. + .. deprecated:: Use :meth:`list_tables` instead, which provides proper pagination support. """ @@ -1157,13 +1170,15 @@ class AsyncLanceNamespaceDBConnection: async def table_names( self, page_token: Optional[str] = None, - limit: int = 10, + limit: Optional[int] = None, *, namespace_path: Optional[List[str]] = None, ) -> Iterable[str]: """ List table names in the namespace. + If limit is None, return all table names; zero returns an empty list. + .. deprecated:: Use :meth:`list_tables` instead, which provides proper pagination support. """ @@ -1584,11 +1599,13 @@ class AsyncLanceNamespaceDBConnection: query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> AsyncRecordBatchReader: """Execute SQL when supported by the underlying connection.""" return await self._inner.execute_query( query, default_namespace_path=default_namespace_path, + parameters=parameters, ) async def execute_query_async( @@ -1596,6 +1613,7 @@ class AsyncLanceNamespaceDBConnection: query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> AsyncSqlQuery: """Start executing SQL when supported by the underlying connection. @@ -1604,6 +1622,7 @@ class AsyncLanceNamespaceDBConnection: return await self._inner.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=parameters, ) async def describe_query(self, query_id: UUID) -> QueryDescription: diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index b210c6497..c7664ab78 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -39,7 +39,7 @@ from .expr import Expr from .rerankers.base import Reranker from .rerankers.rrf import RRFReranker from .rerankers.util import check_reranker_result -from .schema import is_blob_like_field, schema_has_blob_field +from .schema import blob_column_paths, is_blob_like_field, schema_has_blob_field from .util import _validate_query_vector, flatten_columns from . import _wal_hybrid # WAL-PK-FUSION: delete. from ._blob import ( @@ -51,6 +51,7 @@ from ._blob import ( finalize_blob_query_table, replace_v2_blob_columns_with_bytes, replace_v2_blob_columns_with_bytes_sync, + strip_auto_row_ids, validate_blob_mode, ) from .types import BlobMode, QueryProjection @@ -479,7 +480,7 @@ class MatchQuery(FullTextQuery): boost : float, default 1.0 The boost factor for the query. The score of each matching document is multiplied by this value. - fuzziness : int, optional + fuzziness : int or None, default 0 The maximum edit distance for each term in the match query. Defaults to 0 (exact match). If None, fuzziness is applied automatically by the rules: @@ -506,7 +507,7 @@ class MatchQuery(FullTextQuery): query: str column: str boost: float = pydantic.Field(1.0, kw_only=True) - fuzziness: int = pydantic.Field(0, kw_only=True) + fuzziness: Optional[int] = pydantic.Field(0, kw_only=True) 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) @@ -692,8 +693,9 @@ class Query(pydantic.BaseModel): if True then apply the filter after vector / FTS search. This is ignored for plain SQL filtering. nprobes : Optional[int] - The number of IVF partitions to search. If this is None then a default - number of partitions will be used. + The legacy number of IVF partitions to search. Lance sets both probe + bounds to this value. If this is None then Lance's default probe settings + will be used. - A higher number makes search more accurate but also slower. @@ -772,20 +774,23 @@ class Query(pydantic.BaseModel): # distance type to use for vector search distance_type: Optional[str] = None + # legacy number of IVF partitions to search + # + # Lance sets both probe bounds to this value. Explicit bounds can override it. + nprobes: Optional[int] = None + # which columns to return in the results (dict values may be str or Expr) columns: QueryProjection = None # minimum number of IVF partitions to search # - # If None then a default value (20) will be used. + # If None then Lance's default will be used. minimum_nprobes: Optional[int] = None # maximum number of IVF partitions to search # - # If None then a default value (20) will be used. - # - # If 0 then no limit will be applied and all partitions could be searched - # if needed to satisfy the limit. + # If None then Lance's default will be used. If 0 then no limit will be applied + # and all partitions could be searched if needed to satisfy the limit. maximum_nprobes: Optional[int] = None # lower bound for distance search @@ -846,6 +851,7 @@ class Query(pydantic.BaseModel): query.vector_column = req.column query.vector = req.query_vector query.distance_type = req.distance_type + query.nprobes = req.nprobes query.minimum_nprobes = req.minimum_nprobes query.maximum_nprobes = req.maximum_nprobes query.lower_bound = req.lower_bound @@ -1626,6 +1632,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): super().__init__(table) self._query = query self._distance_type = None + self._nprobes = None self._minimum_nprobes = None self._maximum_nprobes = None self._lower_bound = None @@ -1690,9 +1697,9 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): See discussion in [Querying an ANN Index](https://lancedb.com/docs/indexing/) for tuning advice. - This method sets both the minimum and maximum number of probes to the same - value. See `minimum_nprobes` and `maximum_nprobes` for more fine-grained - control. + The value is retained as `nprobes` through client and server request + construction. Lance sets both probe bounds to this value. Explicit minimum + or maximum settings can override their respective bounds. Parameters ---------- @@ -1706,8 +1713,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): """ if nprobes <= 0: raise ValueError("Invalid input, nprobes must be greater than 0") - self._minimum_nprobes = nprobes - self._maximum_nprobes = nprobes + self._nprobes = nprobes return self def minimum_nprobes(self, minimum_nprobes: int) -> LanceVectorQueryBuilder: @@ -1849,6 +1855,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): limit=self._limit, distance_type=self._distance_type, columns=self._columns, + nprobes=self._nprobes, minimum_nprobes=self._minimum_nprobes, maximum_nprobes=self._maximum_nprobes, lower_bound=self._lower_bound, @@ -2207,6 +2214,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): self._fts_columns = fts_columns self._norm = None self._reranker = None + self._nprobes = None self._minimum_nprobes = None self._maximum_nprobes = None self._refine_factor = None @@ -2498,6 +2506,9 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): Higher values will yield better recall (more likely to find vectors if they exist) at the expense of latency. + The value is retained as `nprobes` until Lance sets both probe bounds to it. + Explicit minimum or maximum settings can override their respective bounds. + Parameters ---------- nprobes: int @@ -2510,8 +2521,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): """ if nprobes <= 0: raise ValueError("Invalid input, nprobes must be greater than 0") - self._minimum_nprobes = nprobes - self._maximum_nprobes = nprobes + self._nprobes = nprobes return self def minimum_nprobes(self, minimum_nprobes: int) -> LanceHybridQueryBuilder: @@ -2790,6 +2800,8 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): self._fts_query.phrase_query(True) if self._distance_type: self._vector_query.metric(self._distance_type) + if self._nprobes is not None: + self._vector_query.nprobes(self._nprobes) if self._minimum_nprobes is not None: self._vector_query.minimum_nprobes(self._minimum_nprobes) if self._maximum_nprobes is not None: @@ -3073,14 +3085,18 @@ class AsyncQueryBase(object): If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. blob_mode: str, default "lazy" - Controls how blob columns are returned for plain scan queries. - Vector, FTS, hybrid, and other non-native query shapes keep the - existing Arrow conversion path and only support blob descriptions. + Controls how blob columns are returned. Remote queries support + "descriptions"; "bytes" and "lazy" are not yet supported. **kwargs Forwarded to pyarrow.Table.to_pandas after query execution and optional flattening. """ validate_blob_mode(blob_mode) + if self._table is not None and not self._table._inner._is_native(): + return await self._remote_to_pandas( + flatten=flatten, timeout=timeout, blob_mode=blob_mode, **kwargs + ) + if hasattr(self._inner, "output_schema"): schema = await self.output_schema() if _blob_mode_requires_native_pandas(blob_mode, schema): @@ -3108,6 +3124,31 @@ class AsyncQueryBase(object): ) return tbl.to_pandas(**kwargs) + async def _remote_to_pandas( + self, + *, + flatten: Optional[Union[int, bool]], + timeout: Optional[timedelta], + blob_mode: BlobMode, + **kwargs, + ) -> "pd.DataFrame": + # A live remote table can be replaced between query and blob fetch. + # Row ids and version numbers do not identify the producing table. + tbl = await self.to_arrow(timeout=timeout) + projected_blob_paths = set(blob_column_paths(tbl.schema)) + if not projected_blob_paths: + return flatten_columns(tbl, flatten).to_pandas(**kwargs) + + if blob_mode == "descriptions": + tbl = strip_auto_row_ids(tbl, self._blob_paths) + return flatten_columns(tbl, flatten).to_pandas(**kwargs) + + raise NotImplementedError( + f"remote to_pandas(blob_mode={blob_mode!r}) cannot safely materialize " + "blob columns without a stable table snapshot; " + "use blob_mode='descriptions'" + ) + async def _plain_scan_to_pandas( self, blob_mode: BlobMode, @@ -3660,7 +3701,7 @@ class AsyncVectorQueryBase: def nprobes(self, nprobes: int) -> Self: """ - Set the number of partitions to search (probe) + Set the legacy IVF probe parameter The number of probes must be greater than 0. @@ -3672,16 +3713,19 @@ class AsyncVectorQueryBase: The partition whose centroids are closest to the query vector will be exhaustiely searched to find matches. This parameter controls how many - partitions should be searched. + partitions will be searched. Increasing this value will increase the recall of your query but will - also increase the latency of your query. The default value is 20. This - default is good for many cases but the best value to use will depend on - your data and the recall that you need to achieve. + also increase the latency of your query. If this method is not called, + Lance's adaptive probe defaults are used. For best results we recommend tuning this parameter with a benchmark against your actual data to find the smallest possible value that will still give you the desired recall. + + LanceDB retains this as `nprobes` through local and remote request + construction. Lance sets both probe bounds to this value. Explicit minimum + or maximum settings can override their respective bounds. """ self._inner.nprobes(nprobes) return self diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 2688c3efc..0f3d35fc7 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -41,7 +41,7 @@ from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job from ..sql import Query as SqlQuery -from ..sql import QueryDescription +from ..sql import QueryDescription, QueryParameters from ..materialized_view import MaterializedView, MaterializedViewSource, SelectArg from ..secrets import EnvVarSecret, SecretInfo from ..view import ViewDescription @@ -379,7 +379,7 @@ class RemoteDBConnection(DBConnection): page_token: Optional[str] = None, limit: Optional[int] = None, ) -> ListTablesResponse: - """List all tables in this database with pagination support. + """List a page of tables in this database. Parameters ---------- @@ -390,7 +390,8 @@ class RemoteDBConnection(DBConnection): Token for pagination. Use the token from a previous response to get the next page of results. limit: int, optional - The maximum number of results to return. + The maximum number of results to return, default 100. + Zero returns an empty page without a continuation token. Returns ------- @@ -409,7 +410,7 @@ class RemoteDBConnection(DBConnection): def table_names( self, page_token: Optional[str] = None, - limit: int = 10, + limit: Optional[int] = None, *, namespace_path: Optional[List[str]] = None, ) -> Iterable[str]: @@ -425,8 +426,9 @@ class RemoteDBConnection(DBConnection): Empty list represents root namespace. page_token: str The last token to start the new page. - limit: int, default 10 - The maximum number of tables to return for each page. + limit: int, optional + The maximum number of tables to return. None returns all tables; + zero returns an empty list. Returns ------- @@ -486,12 +488,12 @@ class RemoteDBConnection(DBConnection): namespace_path = [] if storage_options is not None: logging.info( - "storage_options is ignored in LanceDb Cloud" + "storage_options is ignored for remote tables" " (storage is managed; set storage_options on connect() instead)" ) if index_cache_size is not None: logging.info( - "index_cache_size is ignored in LanceDb Cloud" + "index_cache_size is ignored for remote tables" " (there is no local cache to configure)" ) @@ -1090,18 +1092,22 @@ class RemoteDBConnection(DBConnection): query: str, *, default_namespace_path: Optional[List[str]] = None, + parameters: Optional[QueryParameters] = None, ) -> SqlQuery: """Start executing SQL through this remote connection. Unqualified tables use this connection's database and the ``["public"]`` namespace by default. Fully qualified table names may reference other databases available to the same deployment. + ``parameters`` binds values to ``$1`` / ``$name`` placeholders; see + [DBConnection.execute_query_async][lancedb.db.DBConnection.execute_query_async]. """ return SqlQuery( LOOP.run( self._conn.execute_query_async( query, default_namespace_path=default_namespace_path, + parameters=parameters, ) ) ) diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 57d01e4c9..7f03f8fa9 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -42,6 +42,7 @@ from lancedb.index import ( BTree, Bitmap, HnswFlat, + HnswPq, HnswSq, IvfFlat, IvfPq, @@ -239,12 +240,12 @@ class RemoteTable(Table): return LOOP.run(self._table.list_versions()) def to_arrow(self) -> pa.Table: - """to_arrow() is not yet supported on LanceDB cloud.""" - raise NotImplementedError("to_arrow() is not yet supported on LanceDB cloud.") + """to_arrow() is not supported for remote tables.""" + raise NotImplementedError("to_arrow() is not supported for remote tables.") def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs): - """to_pandas() is not yet supported on LanceDB cloud.""" - raise NotImplementedError("to_pandas() is not yet supported on LanceDB cloud.") + """to_pandas() is not supported for remote tables.""" + raise NotImplementedError("to_pandas() is not supported for remote tables.") def checkout(self, version: Union[int, str]): result = LOOP.run(self._table.checkout(version)) @@ -353,8 +354,8 @@ class RemoteTable(Table): language: str = "English", max_token_length: Optional[int] = 40, lower_case: bool = True, - stem: bool = True, - remove_stop_words: bool = True, + stem: Optional[bool] = None, + remove_stop_words: Optional[bool] = None, custom_stop_words: Optional[List[str]] = None, ascii_folding: bool = True, ngram_min_length: int = 3, @@ -501,14 +502,8 @@ class RemoteTable(Table): if accelerator is not None: logging.warning( - "GPU accelerator is not yet supported on LanceDB cloud." - "If you have 100M+ vectors to index," - "please contact us at contact@lancedb.com" - ) - if replace is not None: - logging.warning( - "replace is not supported on LanceDB cloud." - "Existing indexes will always be replaced." + "GPU accelerator is not supported for remote tables " + "and will be ignored." ) idx_type = index_type.upper() @@ -528,9 +523,11 @@ class RemoteTable(Table): elif idx_type == "IVF_SQ": config = IvfSq(distance_type=metric, num_partitions=num_partitions) elif idx_type == "IVF_HNSW_PQ": - raise ValueError( - "IVF_HNSW_PQ is not supported on LanceDB cloud." - "Please use IVF_HNSW_SQ instead." + config = HnswPq( + distance_type=metric, + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, ) elif idx_type == "IVF_HNSW_SQ": config = HnswSq(distance_type=metric, num_partitions=num_partitions) @@ -659,7 +656,7 @@ class RemoteTable(Table): capped at the number of CPU cores. Lower this if bulk ingestion is using too much memory. allow_external_blob_outside_bases: bool, default False - Not supported on LanceDB Cloud. Setting this raises. + Not supported for remote tables. Setting this raises. Returns ------- @@ -689,7 +686,7 @@ class RemoteTable(Table): vector_column_name: Optional[str] = None, query_type="auto", fts_columns: Optional[Union[str, List[str]]] = None, - fast_search: bool = False, + fast_search: Optional[bool] = None, ) -> LanceVectorQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support @@ -746,7 +743,7 @@ class RemoteTable(Table): Skip a flat search of unindexed data. This may improve search performance but search results will not include unindexed data. - - *default False*. + - *default None*, which leaves fast search disabled. Returns ------- @@ -781,7 +778,7 @@ class RemoteTable(Table): query, query_type, vector_column_name=vector_column_name, - fts_columns=fts_columns, + fts_columns=fts_columns or [], fast_search=fast_search, ) @@ -956,26 +953,20 @@ class RemoteTable(Table): def cleanup_old_versions(self, *_): """ - cleanup_old_versions() is a no-op on LanceDB Cloud. - - Tables are automatically cleaned up and optimized. + cleanup_old_versions() is a no-op for remote tables. """ warnings.warn( - "cleanup_old_versions() is a no-op on LanceDB Cloud. " - "Tables are automatically cleaned up and optimized.", + "cleanup_old_versions() is a no-op for remote tables.", stacklevel=2, ) pass def compact_files(self, *_): """ - compact_files() is a no-op on LanceDB Cloud. - - Tables are automatically compacted and optimized. + compact_files() is a no-op for remote tables. """ warnings.warn( - "compact_files() is a no-op on LanceDB Cloud. " - "Tables are automatically compacted and optimized.", + "compact_files() is a no-op for remote tables.", stacklevel=2, ) pass @@ -986,17 +977,13 @@ class RemoteTable(Table): cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, ): - """ - optimize() is a no-op on LanceDB Cloud. - - Indices are optimized automatically. - """ - warnings.warn( - "optimize() is a no-op on LanceDB Cloud. " - "Indices are optimized automatically.", - stacklevel=2, + """optimize() is not supported for remote tables.""" + return LOOP.run( + self._table.optimize( + cleanup_older_than=cleanup_older_than, + delete_unverified=delete_unverified, + ) ) - pass def count_rows(self, filter: Optional[str] = None) -> int: return LOOP.run(self._table.count_rows(filter)) @@ -1044,7 +1031,7 @@ class RemoteTable(Table): return LOOP.run(self._table.drop_columns(columns)) def set_unenforced_primary_key(self, columns: Union[str, Iterable[str]]) -> None: - """Not supported on LanceDB Cloud.""" + """Not supported for remote tables.""" return LOOP.run(self._table.set_unenforced_primary_key(columns)) def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None: @@ -1090,7 +1077,7 @@ class RemoteTable(Table): ) def close_lsm_writers(self) -> None: - """No-op on LanceDB Cloud (no local shard writers).""" + """No-op for remote tables (no local shard writers).""" return LOOP.run(self._table.close_lsm_writers()) def drop_index(self, index_name: str): @@ -1158,14 +1145,10 @@ class RemoteTable(Table): return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) def uses_v2_manifest_paths(self) -> bool: - raise NotImplementedError( - "uses_v2_manifest_paths() is not supported on the LanceDB Cloud" - ) + return LOOP.run(self._table.uses_v2_manifest_paths()) def migrate_v2_manifest_paths(self): - raise NotImplementedError( - "migrate_v2_manifest_paths() is not supported on the LanceDB Cloud" - ) + return LOOP.run(self._table.migrate_manifest_paths_v2()) def blob_columns(self) -> list[str]: return LOOP.run(self._table.blob_columns()) @@ -1176,9 +1159,7 @@ class RemoteTable(Table): return LOOP.run(self._table.fetch_blobs(column, row_ids)) def fetch_blob_ranges(self, column: str, requests) -> pa.LargeBinaryArray: - raise NotImplementedError( - "fetch_blob_ranges() is not supported on LanceDB Cloud" - ) + return LOOP.run(self._table.fetch_blob_ranges(column, list(requests))) def fetch_blob_files( self, column: str, row_ids: Union[list[int], pa.Table] diff --git a/python/python/lancedb/sql.py b/python/python/lancedb/sql.py index 41bbb5328..05b82b0d3 100644 --- a/python/python/lancedb/sql.py +++ b/python/python/lancedb/sql.py @@ -3,8 +3,10 @@ """Handles to SQL queries running on a remote database.""" +from typing import Any, Mapping, Optional, Sequence, Union from uuid import UUID +import numpy as np import pyarrow as pa from lancedb.background_loop import LOOP @@ -14,6 +16,82 @@ from .arrow import AsyncRecordBatchReader QueryDescription = _lancedb.QueryDescription +QueryParameters = Union[pa.RecordBatch, pa.Table, Sequence[Any], Mapping[str, Any]] +"""Values for a statement's ``$1`` / ``$name`` placeholders. + +A sequence binds by position (``$1``, ``$2``, ...), a mapping binds by name +(``$name``), and a one-row ``pyarrow.RecordBatch`` or ``pyarrow.Table`` binds +both ways: column ``i`` is ``$`` and, when its name is not a number, +also ``$``. +""" + + +def _parameter_array(value: Any) -> pa.Array: + """One parameter as a one-element array of its own type.""" + if isinstance(value, pa.Scalar): + return pa.array([value.as_py()], type=value.type) + # A 1-D array is a vector: keep its element type and make it the + # fixed-size list a vector column stores, rather than a list of Python + # floats that would come back as float64. + if isinstance(value, pa.ChunkedArray): + value = value.combine_chunks() + if isinstance(value, np.ndarray): + if value.ndim == 0: + return pa.array(value.reshape(1)) + if value.ndim != 1: + raise ValueError( + "a numpy query parameter must be a scalar or a 1-D vector, " + f"got {value.ndim} dimensions" + ) + value = pa.array(value) + if isinstance(value, pa.Array): + return pa.FixedSizeListArray.from_arrays(value, len(value)) + return pa.array([value]) + + +def to_parameter_batch( + parameters: Optional[QueryParameters], +) -> Optional[pa.RecordBatch]: + """Turn ``parameters`` into the one-row batch a parameterized query sends. + + Scalars keep the type pyarrow infers for them (numpy scalars keep their + dtype), a ``pyarrow.Scalar`` keeps its own type, and a 1-D numpy or + pyarrow array becomes a ``FixedSizeList`` of its element type, which is + what a vector column stores. + """ + if parameters is None: + return None + batch: pa.RecordBatch + if isinstance(parameters, pa.RecordBatch): + batch = parameters + elif isinstance(parameters, pa.Table): + batches = parameters.combine_chunks().to_batches() + batch = ( + batches[0] + if len(batches) == 1 + else pa.RecordBatch.from_pylist([], schema=parameters.schema) + ) + elif isinstance(parameters, Mapping): + batch = pa.RecordBatch.from_arrays( + [_parameter_array(value) for value in parameters.values()], + names=[str(name) for name in parameters.keys()], + ) + elif isinstance(parameters, Sequence) and not isinstance(parameters, (str, bytes)): + batch = pa.RecordBatch.from_arrays( + [_parameter_array(value) for value in parameters], + names=[str(position + 1) for position in range(len(parameters))], + ) + else: + raise TypeError( + "query parameters must be a sequence, a mapping, or a one-row " + f"pyarrow RecordBatch or Table, not {type(parameters).__name__}" + ) + if batch.num_rows != 1: + raise ValueError( + f"query parameters must be exactly one row, got {batch.num_rows}" + ) + return batch + class AsyncQuery: """A handle to a submitted SQL query on an asynchronous connection.""" @@ -85,4 +163,4 @@ class Query: LOOP.run(self._inner.cancel()) -__all__ = ["AsyncQuery", "Query", "QueryDescription"] +__all__ = ["AsyncQuery", "Query", "QueryDescription", "QueryParameters"] diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 2275b0a93..21838ecef 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1602,8 +1602,8 @@ class Table(ABC): language: str = "English", max_token_length: Optional[int] = 40, lower_case: bool = True, - stem: bool = True, - remove_stop_words: bool = True, + stem: Optional[bool] = None, + remove_stop_words: Optional[bool] = None, custom_stop_words: Optional[List[str]] = None, ascii_folding: bool = True, ngram_min_length: int = 3, @@ -1653,6 +1653,7 @@ class Table(ABC): - "simple": Splits text by whitespace and punctuation. - "whitespace": Split text by whitespace, but not punctuation. - "raw": No tokenization. The entire text is treated as a single token. + - "code": Tokenizes source code and identifiers. - "ngram": N-Gram tokenizer. - "icu": ICU dictionary-based word segmentation. - "icu/split": ICU segmentation with simple-style delimiter splitting. @@ -1667,12 +1668,16 @@ class Table(ABC): lower_case : bool, default True Whether to convert the token to lower case. This makes queries case-insensitive. - stem : bool, default True + stem : bool, optional Whether to stem the token. Stemming reduces words to their root form. For example, in English "running" and "runs" would both be reduced to "run". - remove_stop_words : bool, default True + ``None`` uses the base tokenizer's default: False for ``code`` and + ``ngram``, True otherwise. + remove_stop_words : bool, optional Whether to remove stop words. Stop words are common words that are often removed from text before indexing. For example, in English "the" and "and". + ``None`` uses the base tokenizer's default: False for ``code`` and + ``ngram``, True otherwise. custom_stop_words : list of str, optional Custom words that replace the built-in language stop words. ``None`` uses the built-in list; an empty list explicitly uses no stop words. @@ -3808,8 +3813,8 @@ class LanceTable(Table): language: str = "English", max_token_length: Optional[int] = 40, lower_case: bool = True, - stem: bool = True, - remove_stop_words: bool = True, + stem: Optional[bool] = None, + remove_stop_words: Optional[bool] = None, custom_stop_words: Optional[List[str]] = None, ascii_folding: bool = True, ngram_min_length: int = 3, @@ -5597,7 +5602,11 @@ class AsyncTable: return AsyncQuery(self._inner.query(), self) async def to_lance(self, **kwargs) -> lance.LanceDataset: - """Return the Lance dataset backing this table. + """Return the Lance dataset backing a local table. + + Remote tables cannot be opened as Lance datasets by the client. Use + [to_arrow][lancedb.table.AsyncTable.to_arrow] or + [query][lancedb.table.AsyncTable.query] to read them through the server. Parameters ---------- @@ -5614,6 +5623,12 @@ class AsyncTable: >>> async def get_lance_dataset(table): ... return await table.to_lance() """ + if not self._inner._is_native(): + raise NotImplementedError( + "to_lance() is not supported for remote tables; " + "query the server instead" + ) + try: import lance except ImportError: @@ -5643,7 +5658,8 @@ class AsyncTable: Parameters ---------- blob_mode: str, default "lazy" - Controls how Lance blob columns are returned. + Controls how Lance blob columns are returned. Remote tables support + "descriptions"; "bytes" and "lazy" are not yet supported. **kwargs Forwarded to PyArrow / Lance pandas conversion. @@ -5661,6 +5677,9 @@ class AsyncTable: ) return arrow_tbl.to_pandas(**kwargs) + if not self._inner._is_native(): + return await self.query().to_pandas(blob_mode=blob_mode, **kwargs) + if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory": return (await self.to_arrow()).to_pandas(**kwargs) if blob_mode == "bytes" and blob_v2_column_paths(schema): @@ -6420,9 +6439,10 @@ class AsyncTable: ) if query.distance_type is not None: async_query = async_query.distance_type(query.distance_type) + if query.nprobes is not None: + async_query = async_query.nprobes(query.nprobes) if query.minimum_nprobes is not None and query.maximum_nprobes is not None: - # Set both to the minimum first to avoid min > max error. - async_query = async_query.nprobes( + async_query = async_query.minimum_nprobes( query.minimum_nprobes ).maximum_nprobes(query.maximum_nprobes) elif query.minimum_nprobes is not None: diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index ebd86bdc3..7d37910a3 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -245,6 +245,40 @@ def test_db_contains_and_len_include_all_table_name_pages(tmp_db: lancedb.DBConn assert "does_not_exist" not in tmp_db +@pytest.mark.asyncio +async def test_table_listing_defaults_and_zero_limit(tmp_path): + names = [f"t{i:03}" for i in range(130)] + for name in names: + (tmp_path / f"{name}.lance").mkdir() + + db = lancedb.connect(tmp_path) + assert db.table_names() == names + assert db.table_names(limit=None) == names + assert db.table_names(limit=0) == [] + first = db.list_tables() + assert first.tables == names[:100] + assert first.page_token + second = db.list_tables(page_token=first.page_token) + assert second.tables == names[100:] + assert second.page_token is None + assert db.list_tables(limit=200).tables == names + empty = db.list_tables(limit=0) + assert empty.tables == [] + assert empty.page_token is None + + adb = await lancedb.connect_async(tmp_path) + assert await adb.table_names() == names + assert await adb.table_names(limit=0) == [] + first = await adb.list_tables() + assert first.tables == names[:100] + second = await adb.list_tables(page_token=first.page_token) + assert second.tables == names[100:] + assert second.page_token is None + empty = await adb.list_tables(limit=0) + assert empty.tables == [] + assert empty.page_token is None + + def test_db_contains_stops_after_matching_table_page( tmp_db: lancedb.DBConnection, monkeypatch ): diff --git a/python/python/tests/test_e2e_remote_db.py b/python/python/tests/test_e2e_remote_db.py index d4364bc43..74d5cd7de 100644 --- a/python/python/tests/test_e2e_remote_db.py +++ b/python/python/tests/test_e2e_remote_db.py @@ -2,13 +2,78 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import os +from datetime import timedelta +from uuid import uuid4 + +import lancedb import numpy as np +import pyarrow as pa import pytest from lancedb import LanceDBConnection +from lancedb.index import FTS # TODO: setup integ test mark and script +@pytest.mark.skipif( + not os.environ.get("LANCEDB_URI"), reason="Set LANCEDB_URI to a JSON FTS server" +) +def test_json_fts_against_server(): + """Requires a writable database and capacity to run the distributed indexer. + + Configure LANCEDB_URI, LANCEDB_API_KEY, and optionally LANCEDB_REGION and + LANCEDB_HOST_OVERRIDE before running this test. + """ + uri = os.environ["LANCEDB_URI"] + assert uri.startswith("db://"), "LANCEDB_URI must identify a remote database" + conn = lancedb.connect( + uri, + api_key=os.environ.get("LANCEDB_API_KEY"), + region=os.environ.get("LANCEDB_REGION", "us-east-1"), + host_override=os.environ.get("LANCEDB_HOST_OVERRIDE"), + ) + name = f"json_fts_{uuid4().hex}" + table = conn.create_table( + name, + pa.table( + { + "id": [0, 1, 2], + "doc": pa.array( + [ + '{"meta":{"tag":"nature"},"author":"alice"}', + '{"meta":{"tag":"city"},"author":"nature"}', + '{"meta":{"tag":"nature"},"author":"bob"}', + ], + type=pa.json_(), + ), + "raw": pa.array([b"nature"] * 3, type=pa.large_binary()), + } + ), + ) + try: + table.create_index("doc", config=FTS(), name="doc_fts") + table.wait_for_index(["doc_fts"], timeout=timedelta(minutes=5)) + stats = table.index_stats("doc_fts") + assert stats.num_indexed_rows == 3 + assert stats.num_unindexed_rows == 0 + assert stats.index_type == "FTS" + for query, expected in [ + ("meta.tag,str,nature", [0, 2]), + ("meta.tag,str,alice", []), + ("meta.tag,str,tag", []), + ]: + search = table.search(query, query_type="fts", fts_columns="doc") + assert sorted(row["id"] for row in search.to_list()) == expected + plan = search.explain_plan(True) + assert "MatchQuery" in plan, plan + assert "FlatMatchQuery" not in plan, plan + with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"): + table.create_index("raw", config=FTS()) + finally: + conn.drop_table(name) + + @pytest.mark.skip(reason="Need to set up a local server") def test_against_local_server(): conn = LanceDBConnection("lancedb+http://localhost:10024") diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index c18893f1b..975310ed4 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -1555,7 +1555,7 @@ def _mock_remote_function_catalog(): )["image"], "signature": body["signature"], "secret_bindings": body.get("secret_bindings", []), - "created_at": "2026-08-21T00:00:00Z", + "created_at_millis": 1787270400000, # 2026-08-21T00:00:00Z } response = {"job_id": "job-register"} status = 202 diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index 5107ed445..d7fa2e0ec 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -13,6 +13,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import json import os import random import shutil @@ -44,6 +45,75 @@ from utils import exception_output TEST_LANGUAGE_MODEL_HOME = Path(__file__).parent / "models" +@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type") +@pytest.mark.parametrize("block_size", [128, 256]) +def test_json_fts(tmp_path, block_size): + db = ldb.connect(tmp_path) + data = pa.table( + { + "id": [0, 1, 2, 3], + "doc": pa.array( + [ + '{"title":"brown fox","meta":{"tag":"nature"},"author":"alice"}', + '{"title":"blue car","meta":{"tag":"city"}}', + '{"title":"brown bear","meta":{"tag":"nature"}}', + '{"title":"red fox","author":"alice"}', + ], + type=pa.json_(), + ), + } + ) + table = db.create_table("json_fts", data) + table.create_index("doc", config=FTS(block_size=block_size), name="doc_fts") + table = db.open_table("json_fts") + index = table.list_indices()[0] + assert index.name == "doc_fts" + assert index.index_type == "FTS" + assert index.columns == ["doc"] + stats = table.index_stats("doc_fts") + assert stats.num_indexed_rows == 4 + assert stats.num_unindexed_rows == 0 + + for query, expected in [ + ("title,str,brown", [0, 2]), + ("meta.tag,str,nature", [0, 2]), + ("author,str,alice", [0, 3]), + ("title,str,alice", []), + ("title,str,title", []), + ]: + search = table.search(query, query_type="fts", fts_columns="doc") + assert sorted(row["id"] for row in search.to_list()) == expected + plan = search.explain_plan(True) + assert "MatchQuery" in plan, plan + assert "FlatMatchQuery" not in plan, plan + + +@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type") +@pytest.mark.asyncio +async def test_json_fts_async(tmp_path): + db = await ldb.connect_async(tmp_path) + table = await db.create_table( + "json_fts", + pa.table({"doc": pa.array(['{"meta":{"tag":"nature"}}'], type=pa.json_())}), + ) + await table.create_index("doc", config=FTS()) + search = table.query().nearest_to_text(MatchQuery("meta.tag,str,nature", "doc")) + assert len(await search.to_list()) == 1 + plan = await search.explain_plan(True) + assert "MatchQuery" in plan, plan + assert "FlatMatchQuery" not in plan, plan + + +@pytest.mark.parametrize("dtype", [pa.binary(), pa.large_binary()]) +def test_fts_rejects_raw_binary(tmp_path, dtype): + table = ldb.connect(tmp_path).create_table( + "binary_fts", pa.table({"doc": pa.array([b'{"title":"brown"}'], type=dtype)}) + ) + with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"): + table.create_index("doc", config=FTS()) + assert table.list_indices() == [] + + @pytest.fixture def table(tmp_path) -> ldb.table.LanceTable: # Use local random state to avoid affecting other tests @@ -894,6 +964,78 @@ def test_tokenize_uses_simple_index_tokenizer(mem_db: DBConnection): ] +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize( + ("base_tokenizer", "options", "expected_filters", "expected_tokens"), + [ + ( + "code", + {}, + (False, False), + [("running", 0), ("in", 1), ("cafe", 2)], + ), + ( + "code", + {"stem": None, "remove_stop_words": None}, + (False, False), + [("running", 0), ("in", 1), ("cafe", 2)], + ), + ( + "code", + {"stem": True}, + (True, False), + [("run", 0), ("in", 1), ("cafe", 2)], + ), + ( + "code", + {"remove_stop_words": True}, + (False, True), + [("running", 0), ("cafe", 2)], + ), + ( + "code", + {"stem": True, "remove_stop_words": True}, + (True, True), + [("run", 0), ("cafe", 2)], + ), + ( + "simple", + {}, + (True, True), + [("run", 0), ("cafe", 2)], + ), + ( + "simple", + {"stem": False, "remove_stop_words": False}, + (False, False), + [("running", 0), ("in", 1), ("cafe", 2)], + ), + ], +) +def test_fts_tokenizer_filter_defaults( + mem_db, legacy, base_tokenizer, options, expected_filters, expected_tokens +): + text = "Running in café" + table = mem_db.create_table("test_filter_defaults", data=[{"text": text}]) + if legacy: + with pytest.warns(DeprecationWarning, match="create_fts_index"): + table.create_fts_index("text", base_tokenizer=base_tokenizer, **options) + else: + table.create_index("text", config=FTS(base_tokenizer=base_tokenizer, **options)) + + details = table.list_indices()[0].index_details + assert (details["stem"], details["remove_stop_words"]) == expected_filters + assert [ + (token.text, token.position) for token in table.tokenize(text, column="text") + ] == expected_tokens + assert [ + (token.text, token.position) + for token in ldb.tokenize(text, base_tokenizer=base_tokenizer, **options) + ] == expected_tokens + if base_tokenizer == "code" and not expected_filters[1]: + assert table.search("in", query_type="fts").to_list()[0]["text"] == text + + def test_tokenize_uses_icu_index_tokenizer_by_name(mem_db: DBConnection): data = pa.table({"text": ["Hello, こんにちは世界!"]}) table = mem_db.create_table("test_tokenize_icu", data=data) @@ -1137,6 +1279,34 @@ def test_fts_lindera_tokenizer( assert [row["text"] for row in results] == ["成田国際空港"] +@pytest.mark.parametrize("fuzziness", [None, 0, 1, 2]) +def test_match_query_fuzziness_to_json(fuzziness): + query = MatchQuery("puppy", "text", fuzziness=fuzziness) + + assert query.fuzziness == fuzziness + assert json.loads(query.to_json())["match"]["fuzziness"] == fuzziness + + +@pytest.mark.parametrize( + "options, expected", + [ + pytest.param({}, ["ab"], id="default-exact"), + pytest.param({"fuzziness": 0}, ["ab"], id="explicit-exact"), + pytest.param({"fuzziness": None}, ["ab"], id="automatic"), + pytest.param({"fuzziness": 1}, ["ab", "ac"], id="one-edit"), + pytest.param({"fuzziness": 2}, ["ab", "ac", "zz"], id="two-edits"), + ], +) +def test_match_query_fuzziness_search(tmp_db, options, expected): + # Automatic fuzziness is zero for two-character terms. + table = tmp_db.create_table("fuzzy", pa.table({"text": ["ab", "ac", "zz"]})) + table.create_index("text", config=FTS(stem=False, remove_stop_words=False)) + + results = table.search(MatchQuery("ab", "text", **options)).to_list() + + assert sorted(row["text"] for row in results) == sorted(expected) + + def test_fts_query_to_json(): """Test that FTS query to_json() produces valid JSON strings with exact format.""" diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index c5f58308d..8d670a136 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -355,6 +355,20 @@ def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table): ) +def test_hybrid_query_nprobes_remains_independent(sync_table: Table): + query = sync_table.search(query_type="hybrid").nprobes(20) + + assert query._nprobes == 20 + assert query._minimum_nprobes is None + assert query._maximum_nprobes is None + + query = query.minimum_nprobes(5).nprobes(20) + + assert query._nprobes == 20 + assert query._minimum_nprobes == 5 + assert query._maximum_nprobes is None + + def test_hybrid_query_distance_range(sync_table: Table): reranker = RRFReranker(return_score="all") result = ( diff --git a/python/python/tests/test_namespace.py b/python/python/tests/test_namespace.py index d9af24a90..76aa5fc55 100644 --- a/python/python/tests/test_namespace.py +++ b/python/python/tests/test_namespace.py @@ -6,12 +6,17 @@ import tempfile import shutil import importlib +import http.server +import json +import threading +from urllib.parse import parse_qs, unquote, urlparse import pytest import pyarrow as pa import lancedb from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError from lancedb.listing import AsyncListing, Listing -from lancedb.namespace import _MAX_QUERY_K +from lancedb.namespace import _MAX_QUERY_K, _query_to_namespace_request +from lancedb.query import Query from lancedb.table import AsyncTable, LanceTable @@ -73,6 +78,132 @@ def _namespace_lance_table(namespace_client: _NamespaceClient) -> LanceTable: return table +@pytest.mark.asyncio +@pytest.mark.parametrize("namespace_path", [[], ["test_ns"]]) +async def test_namespace_table_names_defaults_and_limits(tmp_path, namespace_path): + db = lancedb.connect_namespace("dir", {"root": str(tmp_path)}) + if namespace_path: + db.create_namespace(namespace_path) + names = [f"t{idx:03}" for idx in range(12)] + schema = pa.schema([pa.field("id", pa.int64())]) + for name in names: + db.create_table(name, schema=schema, namespace_path=namespace_path) + + assert list(db.table_names(namespace_path=namespace_path)) == names + assert list(db.table_names(limit=None, namespace_path=namespace_path)) == names + assert list(db.table_names(limit=0, namespace_path=namespace_path)) == [] + assert list(db.table_names(limit=2, namespace_path=namespace_path)) == names[:2] + assert ( + list( + db.table_names(page_token=names[1], limit=2, namespace_path=namespace_path) + ) + == names[2:4] + ) + + adb = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)}) + assert list(await adb.table_names(namespace_path=namespace_path)) == names + assert ( + list(await adb.table_names(limit=None, namespace_path=namespace_path)) == names + ) + assert list(await adb.table_names(limit=0, namespace_path=namespace_path)) == [] + assert ( + list(await adb.table_names(limit=2, namespace_path=namespace_path)) == names[:2] + ) + assert ( + list( + await adb.table_names( + page_token=names[1], limit=2, namespace_path=namespace_path + ) + ) + == names[2:4] + ) + + +@pytest.fixture +def paginated_rest_namespace(): + names = [f"t{idx:03}" for idx in reversed(range(12))] + requests = [] + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + url = urlparse(self.path) + query = parse_qs(url.query) + requests.append((unquote(url.path), query)) + token = query.get("page_token", [None])[0] + if token not in {None, "opaque-next"}: + self.send_error(400, "Unknown page token") + return + start = 0 if token is None else 10 + body = {"tables": names[start : start + 10]} + if start == 0: + body["page_token"] = "opaque-next" + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(body).encode()) + + def log_message(self, *args): + pass + + with http.server.HTTPServer(("127.0.0.1", 0), Handler) as server: + thread = threading.Thread(target=server.serve_forever) + thread.start() + try: + yield {"uri": f"http://127.0.0.1:{server.server_port}"}, requests + finally: + server.shutdown() + thread.join() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("namespace_path", [[], ["n"]]) +async def test_rest_namespace_table_names_collects_pages( + paginated_rest_namespace, namespace_path +): + properties, requests = paginated_rest_namespace + names = [f"t{idx:03}" for idx in range(12)] + db = lancedb.connect_namespace("rest", properties) + first = db.list_tables(namespace_path=namespace_path) + second = db.list_tables(namespace_path=namespace_path, page_token=first.page_token) + assert len(first.tables) == 10 + assert first.page_token == "opaque-next" + assert sorted(first.tables + second.tables) == names + assert second.page_token is None + + assert list(db.table_names(namespace_path=namespace_path)) == names + assert list(db.table_names(limit=None, namespace_path=namespace_path)) == names + assert list(db.table_names(limit=11, namespace_path=namespace_path)) == names[:11] + assert ( + list(db.table_names(page_token="t005", limit=2, namespace_path=namespace_path)) + == names[6:8] + ) + + adb = lancedb.connect_namespace_async("rest", properties) + assert list(await adb.table_names(namespace_path=namespace_path)) == names + assert ( + list(await adb.table_names(limit=None, namespace_path=namespace_path)) == names + ) + assert ( + list(await adb.table_names(limit=11, namespace_path=namespace_path)) + == names[:11] + ) + assert ( + list( + await adb.table_names( + page_token="t005", limit=2, namespace_path=namespace_path + ) + ) + == names[6:8] + ) + + request_count = len(requests) + assert list(db.table_names(limit=0, namespace_path=namespace_path)) == [] + assert list(await adb.table_names(limit=0, namespace_path=namespace_path)) == [] + assert len(requests) == request_count + identifier = "$" if not namespace_path else "n" + assert all(path == f"/v1/namespace/{identifier}/table/list" for path, _ in requests) + + class TestNamespaceConnection: """Test namespace-based LanceDB connection using DirectoryNamespace.""" @@ -981,6 +1112,29 @@ class TestPushdownOperations: ] assert all(r.k <= 2**31 - 1 for r in namespace_client.requests) + def test_probe_fields_pass_through_independently(self): + request = _query_to_namespace_request( + ["geneva", "hist"], + Query( + vector=[1.0, 2.0], + nprobes=20, + minimum_nprobes=3, + maximum_nprobes=10, + ), + ) + + assert request.nprobes == 20 + assert request.minimum_nprobes == 3 + assert request.maximum_nprobes == 10 + + request = _query_to_namespace_request( + ["geneva", "hist"], + Query(vector=[1.0, 2.0], nprobes=5, maximum_nprobes=0), + ) + assert request.nprobes == 5 + assert request.maximum_nprobes == 0 + assert request.to_dict()["maximum_nprobes"] == 0 + @pytest.mark.asyncio class TestAsyncPushdownOperations: diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index d3fedbcc8..daee0dd07 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -995,12 +995,22 @@ def test_invalid_nprobes_sync(table): ValueError, match="maximum_nprobes must be greater than or equal to minimum_nprobes", ): - LanceVectorQueryBuilder(table, [0, 0], "vector").maximum_nprobes(5).to_list() + ( + LanceVectorQueryBuilder(table, [0, 0], "vector") + .minimum_nprobes(10) + .maximum_nprobes(5) + .to_list() + ) with pytest.raises( ValueError, - match="minimum_nprobes must be less than or equal to maximum_nprobes", + match="maximum_nprobes must be greater than or equal to minimum_nprobes", ): - LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(100).to_list() + ( + LanceVectorQueryBuilder(table, [0, 0], "vector") + .maximum_nprobes(5) + .minimum_nprobes(100) + .to_list() + ) def test_nprobes_works_sync(table): @@ -1039,12 +1049,22 @@ async def test_invalid_nprobes_async(table_async: AsyncTable): ValueError, match="maximum_nprobes must be greater than or equal to minimum_nprobes", ): - await table_async.vector_search([0, 0]).maximum_nprobes(5).to_list() + await ( + table_async.vector_search([0, 0]) + .minimum_nprobes(10) + .maximum_nprobes(5) + .to_list() + ) with pytest.raises( ValueError, match="minimum_nprobes must be less than or equal to maximum_nprobes", ): - await table_async.vector_search([0, 0]).minimum_nprobes(100).to_list() + await ( + table_async.vector_search([0, 0]) + .maximum_nprobes(5) + .minimum_nprobes(100) + .to_list() + ) def test_query_builder_with_prefilter(table): @@ -1635,8 +1655,9 @@ def test_query_serialization_sync(table: lancedb.table.Table): q, vector_column="vector", vector=[5.0, 6.0], - minimum_nprobes=10, - maximum_nprobes=10, + nprobes=10, + minimum_nprobes=None, + maximum_nprobes=None, refine_factor=5, ) @@ -1654,8 +1675,19 @@ def test_query_serialization_sync(table: lancedb.table.Table): q, vector_column="vector", vector=[5.0, 6.0], - minimum_nprobes=50, - maximum_nprobes=50, + nprobes=50, + minimum_nprobes=None, + maximum_nprobes=None, + ) + + q = table.search([5.0, 6.0]).minimum_nprobes(5).nprobes(50).to_query_object() + check_set_props( + q, + vector_column="vector", + vector=[5.0, 6.0], + nprobes=50, + minimum_nprobes=5, + maximum_nprobes=None, ) q = table.search([5.0, 6.0]).maximum_nprobes(10).to_query_object() @@ -1667,6 +1699,15 @@ def test_query_serialization_sync(table: lancedb.table.Table): minimum_nprobes=None, ) + q = table.search([5.0, 6.0]).nprobes(5).maximum_nprobes(0).to_query_object() + check_set_props( + q, + vector_column="vector", + vector=[5.0, 6.0], + nprobes=5, + maximum_nprobes=0, + ) + q = table.search([5.0, 6.0]).distance_range(0.0, 1.0).to_query_object() check_set_props( q, vector_column="vector", vector=[5.0, 6.0], lower_bound=0.0, upper_bound=1.0 @@ -1716,8 +1757,8 @@ async def test_query_serialization_async(table_async: AsyncTable): limit=10, vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, ) @@ -1727,8 +1768,8 @@ async def test_query_serialization_async(table_async: AsyncTable): q, vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1739,8 +1780,9 @@ async def test_query_serialization_async(table_async: AsyncTable): q, vector=sample_vector, postfilter=False, - minimum_nprobes=50, - maximum_nprobes=50, + nprobes=50, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1759,8 +1801,8 @@ async def test_query_serialization_async(table_async: AsyncTable): filter="id = 1", postfilter=True, vector=sample_vector, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, ) @@ -1774,8 +1816,9 @@ async def test_query_serialization_async(table_async: AsyncTable): check_set_props( q, vector=sample_vector, - minimum_nprobes=10, - maximum_nprobes=10, + nprobes=10, + minimum_nprobes=None, + maximum_nprobes=None, refine_factor=5, postfilter=False, with_row_id=False, @@ -1788,7 +1831,24 @@ async def test_query_serialization_async(table_async: AsyncTable): q, vector=sample_vector, minimum_nprobes=5, - maximum_nprobes=20, + maximum_nprobes=None, + postfilter=False, + with_row_id=False, + bypass_vector_index=False, + limit=10, + ) + + q = ( + (await table_async.search([5.0, 6.0])) + .nprobes(5) + .maximum_nprobes(0) + .to_query_object() + ) + check_set_props( + q, + vector=sample_vector, + nprobes=5, + maximum_nprobes=0, postfilter=False, with_row_id=False, bypass_vector_index=False, @@ -1806,8 +1866,8 @@ async def test_query_serialization_async(table_async: AsyncTable): lower_bound=0.0, upper_bound=1.0, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1819,8 +1879,8 @@ async def test_query_serialization_async(table_async: AsyncTable): distance_type="cosine", vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1832,8 +1892,8 @@ async def test_query_serialization_async(table_async: AsyncTable): ef=7, vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1845,8 +1905,8 @@ async def test_query_serialization_async(table_async: AsyncTable): bypass_vector_index=True, vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, limit=10, ) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index b72965faa..3e66ac0e0 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -13,13 +13,19 @@ import threading import time from unittest.mock import MagicMock, patch import uuid +from urllib.parse import parse_qs, urlparse from packaging.version import Version import lancedb import numpy as np from lancedb.conftest import MockNonNormTextEmbeddingFunction, MockTextEmbeddingFunction from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry -from lancedb.query import AsyncQuery, ColumnOrdering, LanceVectorQueryBuilder +from lancedb.query import ( + AsyncQuery, + ColumnOrdering, + LanceVectorQueryBuilder, + MatchQuery, +) from lancedb.remote import ClientConfig from lancedb.remote.errors import HttpError, RetryError import pytest @@ -100,6 +106,67 @@ async def mock_lancedb_connection_async(handler, **client_config): handle.join() +def paginated_table_listing_handler(request): + url = urlparse(request.path) + assert url.path in { + "/v1/table/", + "/v1/namespace/$/table/list", + "/v1/namespace/n/table/list", + } + query = parse_qs(url.query) + limit = int(query.get("limit", [10])[0]) + assert limit > 0 + token = query.get("page_token", ["page-0"])[0] + start = int(token.removeprefix("page-")) + end = min(start + limit, 130) + body = {"tables": [f"t{i:03}" for i in range(start, end)]} + if end < 130: + body["page_token"] = f"page-{end}" + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(body).encode()) + + +@pytest.mark.parametrize("namespace", [[], ["n"]]) +def test_remote_table_listing_defaults_and_zero_limit(namespace): + names = [f"t{i:03}" for i in range(130)] + with mock_lancedb_connection(paginated_table_listing_handler) as db: + assert db.table_names(namespace_path=namespace) == names + assert db.table_names(limit=None, namespace_path=namespace) == names + assert db.table_names(limit=0, namespace_path=namespace) == [] + first = db.list_tables(namespace_path=namespace) + assert first.tables == names[:100] + assert first.page_token == "page-100" + second = db.list_tables(namespace_path=namespace, page_token=first.page_token) + assert second.tables == names[100:] + assert second.page_token is None + assert db.list_tables(namespace_path=namespace, limit=200).tables == names + empty = db.list_tables(namespace_path=namespace, limit=0) + assert empty.tables == [] + assert empty.page_token is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("namespace", [[], ["n"]]) +async def test_async_remote_table_listing_defaults_and_zero_limit(namespace): + names = [f"t{i:03}" for i in range(130)] + async with mock_lancedb_connection_async(paginated_table_listing_handler) as db: + assert await db.table_names(namespace_path=namespace) == names + assert await db.table_names(limit=0, namespace_path=namespace) == [] + first = await db.list_tables(namespace_path=namespace) + assert first.tables == names[:100] + assert first.page_token == "page-100" + second = await db.list_tables( + namespace_path=namespace, page_token=first.page_token + ) + assert second.tables == names[100:] + assert second.page_token is None + empty = await db.list_tables(namespace_path=namespace, limit=0) + assert empty.tables == [] + assert empty.page_token is None + + @pytest.mark.asyncio async def test_async_remote_db(): def handler(request): @@ -1018,6 +1085,77 @@ def test_table_create_indices(): table.drop_index("custom_fts_idx") +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize( + ("base_tokenizer", "options", "expected_filters"), + [ + ("code", {}, (False, False)), + ("code", {"stem": None, "remove_stop_words": None}, (False, False)), + ("code", {"stem": True}, (True, False)), + ("code", {"remove_stop_words": True}, (False, True)), + ("code", {"stem": True, "remove_stop_words": True}, (True, True)), + ("simple", {}, (True, True)), + ("simple", {"stem": False, "remove_stop_words": False}, (False, False)), + ], +) +def test_remote_fts_tokenizer_filter_defaults( + legacy, base_tokenizer, options, expected_filters +): + from lancedb.index import FTS + + received_requests = [] + + def handler(request): + if request.path == "/v1/table/test/create_index/": + content_len = int(request.headers["Content-Length"]) + received_requests.append(json.loads(request.rfile.read(content_len))) + request.send_response(200) + request.end_headers() + elif request.path == "/v1/table/test/create/?mode=create": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b"{}") + elif request.path == "/v1/table/test/describe/": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write( + json.dumps( + { + "version": 1, + "schema": { + "fields": [ + { + "name": "text", + "type": {"type": "string"}, + "nullable": False, + } + ] + }, + } + ).encode() + ) + else: + request.send_response(404) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + table = db.create_table("test", [{"text": "Running in cafes"}]) + if legacy: + with pytest.warns(DeprecationWarning, match="create_fts_index"): + table.create_fts_index("text", base_tokenizer=base_tokenizer, **options) + else: + table.create_index( + "text", config=FTS(base_tokenizer=base_tokenizer, **options) + ) + + assert len(received_requests) == 1 + params = received_requests[0] + assert params["base_tokenizer"] == base_tokenizer + assert (params["stem"], params["remove_stop_words"]) == expected_filters + + def test_remote_create_index_async_returns_job(): from lancedb.index import BTree @@ -1272,7 +1410,73 @@ def test_remote_job_wait_raises_on_failure(): job.wait() -def test_remote_create_index_new_api(): +@pytest.mark.parametrize("server_version", ["0.2.0", "0.6.0"]) +def test_remote_json_fts_config(server_version): + from lancedb.index import FTS + + received_requests = [] + + def handler(request): + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.send_header("phalanx-version", server_version) + request.end_headers() + if request.path == "/v1/table/test/describe/": + fields = [ + { + "name": "doc", + "type": {"type": "large_binary"}, + "nullable": True, + "metadata": {"ARROW:extension:name": "lance.json"}, + }, + { + "name": "raw", + "type": {"type": "large_binary"}, + "nullable": True, + }, + ] + request.wfile.write( + json.dumps({"version": 1, "schema": {"fields": fields}}).encode() + ) + elif request.path == "/v1/table/test/create_index/": + received_requests.append(read_json_body(request)) + request.wfile.write(b"{}") + else: + raise AssertionError(f"Unexpected request: {request.path}") + + with mock_lancedb_connection(handler) as db: + table = db.open_table("test") + table.create_index( + "doc", + config=FTS( + with_position=True, + base_tokenizer="whitespace", + stem=False, + remove_stop_words=False, + lower_case=False, + block_size=256, + ), + name="doc_fts", + replace=False, + ) + with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"): + table.create_index("raw", config=FTS()) + + assert len(received_requests) == 1 + request = received_requests[0] + assert request["column"] == "doc" + assert request["index_type"] == "FTS" + assert request["name"] == "doc_fts" + assert request["replace"] is False + assert request["with_position"] is True + assert request["base_tokenizer"] == "whitespace" + assert request["stem"] is False + assert request["remove_stop_words"] is False + assert request["lower_case"] is False + assert request["block_size"] == 256 + + +def test_remote_create_index_new_api(caplog): received_requests = [] def handler(request): @@ -1333,7 +1537,8 @@ def test_remote_create_index_new_api(): from lancedb.index import BTree, FTS, IvfPq, IvfRq with mock_lancedb_connection(handler) as db: - table = db.create_table("test", [{"id": 1}]) + table = db.create_table("test", [{"id": 1}], embedding_functions=[]) + assert "embedding_functions is not yet supported" not in caplog.text # New API: column-first, config= kwarg. Should NOT emit DeprecationWarning. import warnings as _warnings @@ -1355,6 +1560,8 @@ def test_remote_create_index_new_api(): vector_column_name="vector", index_type="IVF_RQ", num_partitions=8, + replace=False, + accelerator="cuda", ) assert len(received_requests) == 6 @@ -1368,6 +1575,160 @@ def test_remote_create_index_new_api(): ] assert received_requests[2]["block_size"] == 256 assert received_requests[4]["replace"] is False + assert received_requests[5]["replace"] is False + assert "replace is not supported" not in caplog.text + assert caplog.messages == [ + "GPU accelerator is not supported for remote tables and will be ignored." + ] + + +def remote_unsupported_handler(request): + assert request.path == "/v1/table/test/describe/" + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write( + json.dumps( + dict( + version=1, + schema=dict( + fields=[ + dict(name="id", type={"type": "int64"}, nullable=False), + dict( + name="vector", + type={ + "type": "fixed_size_list", + "fields": [ + dict( + name="item", + type={"type": "float"}, + nullable=True, + ) + ], + "length": 2, + }, + nullable=False, + ), + ] + ), + ) + ).encode() + ) + + +def remote_unsupported_operations(): + from lancedb.index import HnswPq + + return [ + ("uses_v2_manifest_paths", (), {}, "uses_v2_manifest_paths"), + ("migrate_v2_manifest_paths", (), {}, "migrate_manifest_paths_v2"), + ("set_unenforced_primary_key", (["id"],), {}, "set_unenforced_primary_key"), + ("fetch_blob_ranges", ("image", [(1, 0, 1)]), {}, "fetch_blob_ranges"), + ("optimize", (), {}, "optimize"), + ("create_index", ("vector",), {"config": HnswPq()}, "IVF_HNSW_PQ"), + ] + + +@pytest.mark.parametrize( + "method,args,kwargs,operation", remote_unsupported_operations() +) +def test_remote_unsupported_operations_sync(method, args, kwargs, operation): + with mock_lancedb_connection(remote_unsupported_handler) as db: + table = db.open_table("test") + with pytest.raises(NotImplementedError) as error: + getattr(table, method)(*args, **kwargs) + expected = f"LanceDBError: not supported: {operation} is not supported " + expected += "for remote tables." + if operation == "IVF_HNSW_PQ": + expected += " Please use IVF_HNSW_SQ instead." + assert str(error.value) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method,args,kwargs,operation", remote_unsupported_operations() +) +async def test_remote_unsupported_operations_async(method, args, kwargs, operation): + async with mock_lancedb_connection_async(remote_unsupported_handler) as db: + table = await db.open_table("test") + if method == "migrate_v2_manifest_paths": + method = "migrate_manifest_paths_v2" + with pytest.raises(NotImplementedError) as error: + await getattr(table, method)(*args, **kwargs) + expected = f"LanceDBError: not supported: {operation} is not supported " + expected += "for remote tables." + if operation == "IVF_HNSW_PQ": + expected += " Please use IVF_HNSW_SQ instead." + assert str(error.value) == expected + + +def test_remote_unsupported_legacy_hnsw_pq(): + with mock_lancedb_connection(remote_unsupported_handler) as db: + table = db.open_table("test") + with pytest.warns(DeprecationWarning, match="create_index"): + with pytest.raises( + NotImplementedError, + match="IVF_HNSW_PQ is not supported for remote tables. " + "Please use IVF_HNSW_SQ instead.", + ): + table.create_index( + vector_column_name="vector", index_type="IVF_HNSW_PQ" + ) + + +def test_remote_python_only_messages(caplog): + with mock_lancedb_connection(remote_unsupported_handler) as db: + with caplog.at_level("INFO"): + table = db.open_table("test", storage_options={}, index_cache_size=1) + assert "storage_options is ignored for remote tables" in caplog.text + assert "index_cache_size is ignored for remote tables" in caplog.text + for method in ["to_arrow", "to_pandas"]: + with pytest.raises(NotImplementedError) as error: + getattr(table, method)() + assert str(error.value) == f"{method}() is not supported for remote tables." + for method in ["compact_files", "cleanup_old_versions"]: + with pytest.warns(UserWarning) as warning: + getattr(table, method)() + assert ( + str(warning[0].message) == f"{method}() is a no-op for remote tables." + ) + + +def test_remote_embedding_config_does_not_warn_or_drop_metadata(caplog): + received_tables = [] + + def handler(request): + assert request.path == "/v1/table/test/create/?mode=create" + if request.headers.get("Transfer-Encoding") == "chunked": + body = bytearray() + while size := int(request.rfile.readline().split(b";")[0], 16): + body.extend(request.rfile.read(size)) + assert request.rfile.read(2) == b"\r\n" + assert request.rfile.read(2) == b"\r\n" + else: + body = request.rfile.read(int(request.headers["Content-Length"])) + received_tables.append(pa.ipc.open_stream(body).read_all()) + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b"{}") + + config = EmbeddingFunctionConfig( + source_column="text", + vector_column="vector", + function=MockNonNormTextEmbeddingFunction.create(), + ) + schema = pa.schema( + [pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 10))] + ) + with mock_lancedb_connection(handler) as db: + db.create_table("test", schema=schema, embedding_functions=[config]) + assert caplog.messages == [] + assert len(received_tables) == 1 + metadata = EmbeddingFunctionRegistry.get_instance().get_table_metadata([config]) + assert ( + received_tables[0].schema.metadata == pa.schema([], metadata=metadata).metadata + ) def test_table_wait_for_index_timeout(): @@ -1674,6 +2035,108 @@ def query_test_table( yield table +@pytest.mark.parametrize( + "query,search_kwargs,expected_column", + [ + pytest.param([0.5] * 8, {}, "vector", id="vector-8"), + pytest.param([0.5] * 4, {}, "vec2", id="vector-4"), + pytest.param(np.array([0.5] * 4), {}, "vec2", id="numpy-vector"), + pytest.param([[0.5] * 4, [0.2] * 4], {}, "vec2", id="batch-vectors"), + pytest.param( + [0.5] * 4, {"vector_column_name": "vec2"}, "vec2", id="explicit-column" + ), + pytest.param(None, {}, None, id="scan"), + pytest.param("hello", {"query_type": "fts"}, None, id="fts"), + pytest.param( + "hello", {"query_type": "fts", "fts_columns": "text"}, None, id="fts-column" + ), + pytest.param(MatchQuery("hello", "text"), {}, None, id="structured-fts"), + ], +) +def test_query_object_sync_matches_local(mem_db, query, search_kwargs, expected_column): + vector_columns = {"vector": 8, "vec2": 4} + schema = pa.schema( + [ + pa.field(name, pa.list_(pa.float32(), dim)) + for name, dim in vector_columns.items() + ] + ) + local = mem_db.create_table("test", schema=schema) + local_query = local.search(query, **search_kwargs).limit(2).to_query_object() + assert local_query.vector_column == expected_column + assert local_query.fast_search is None + + with query_test_table(None, vector_columns=vector_columns) as remote: + remote_query = remote.search(query, **search_kwargs).limit(2).to_query_object() + assert remote_query == local_query + + +@pytest.mark.parametrize("fast_search", [None, False, True]) +@pytest.mark.parametrize("query,query_type", [([1, 2, 3], "vector"), ("hello", "fts")]) +def test_query_object_sync_preserves_fast_search(fast_search, query, query_type): + vector_columns = {"vector": 3} if query_type == "vector" else {} + with query_test_table(None, vector_columns=vector_columns) as table: + query_obj = table.search( + query, query_type=query_type, fast_search=fast_search + ).to_query_object() + assert query_obj.fast_search is fast_search + + +def test_query_sync_rejects_ambiguous_vector_column(mem_db): + vector_columns = {"vector": 4, "vec2": 4} + schema = pa.schema( + [ + pa.field(name, pa.list_(pa.float32(), dim)) + for name, dim in vector_columns.items() + ] + ) + local = mem_db.create_table("test", schema=schema) + with pytest.raises(ValueError, match="Candidates:.*vector.*vec2") as local_error: + local.search([0.5] * 4) + + with query_test_table(None, vector_columns=vector_columns) as remote: + with pytest.raises(ValueError) as remote_error: + remote.search([0.5] * 4) + assert str(remote_error.value) == str(local_error.value) + + +@pytest.mark.parametrize("vector_columns", [{}, {"v1": 4, "v2": 8}]) +@pytest.mark.parametrize("query_type", ["auto", "fts"]) +def test_query_sync_fts_without_embeddings(vector_columns, query_type): + seen = [] + + def handler(body): + seen.append(body) + assert body["full_text_query"] == {"query": "hello", "columns": []} + assert body["vector"] == [] + assert "vector_column" not in body + return pa.table({"id": [1]}) + + with query_test_table(handler, vector_columns=vector_columns) as table: + assert table.search("hello", query_type=query_type).to_list() == [{"id": 1}] + assert len(seen) == 1 + + +def test_query_sync_auto_with_embedding_function(): + def handler(body): + assert body["vector_column"] == "vector" + assert body["vector"] == [0.0] * 10 + assert "full_text_query" not in body + return pa.table({"id": [1]}) + + config = EmbeddingFunctionConfig( + source_column="text", + vector_column="vector", + function=MockTextEmbeddingFunction.create(), + ) + with query_test_table( + handler, vector_columns={"vector": 10}, embedding_functions=[config] + ) as table: + query = table.search("hello") + assert query.to_query_object().vector_column == "vector" + assert query.to_list() == [{"id": 1}] + + def test_head(): def handler(body): assert body == { @@ -1701,9 +2164,6 @@ def test_query_sync_minimal(): "ef": None, "vector": [1.0, 2.0, 3.0], "vector_column": "vector", - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, "version": None, } @@ -1875,8 +2335,6 @@ def test_query_sync_maximal(): "refine_factor": 10, "vector": [1.0, 2.0, 3.0], "nprobes": 5, - "minimum_nprobes": 5, - "maximum_nprobes": 5, "lower_bound": None, "upper_bound": None, "ef": None, @@ -1972,7 +2430,7 @@ async def test_query_async_nprobes_zero(hybrid): assert requests == ["/v1/table/test/describe/"] -def test_query_sync_nprobes(): +def test_query_sync_probe_bounds(): def handler(body): assert body == { "k": 10, @@ -1984,7 +2442,6 @@ def test_query_sync_nprobes(): "upper_bound": None, "ef": None, "vector": [1.0, 2.0, 3.0], - "nprobes": 5, "minimum_nprobes": 5, "maximum_nprobes": 15, "version": None, @@ -2001,7 +2458,28 @@ def test_query_sync_nprobes(): ) -def test_query_sync_no_max_nprobes(): +def test_query_sync_nprobes_passthrough(): + def handler(body): + assert body["nprobes"] == 20 + assert "minimum_nprobes" not in body + assert "maximum_nprobes" not in body + return pa.table({"id": [1]}) + + with query_test_table(handler) as table: + table.search([1, 2, 3]).nprobes(20).to_list() + + +def test_query_sync_nprobes_with_unbounded_maximum(): + def handler(body): + assert body["nprobes"] == 20 + assert body["maximum_nprobes"] == 0 + return pa.table({"id": [1]}) + + with query_test_table(handler) as table: + table.search([1, 2, 3]).nprobes(20).maximum_nprobes(0).to_list() + + +def test_query_sync_minimum_nprobes_only(): def handler(body): assert body == { "k": 10, @@ -2013,9 +2491,7 @@ def test_query_sync_no_max_nprobes(): "upper_bound": None, "ef": None, "vector": [1.0, 2.0, 3.0], - "nprobes": 5, "minimum_nprobes": 5, - "maximum_nprobes": 0, "version": None, } @@ -2025,7 +2501,6 @@ def test_query_sync_no_max_nprobes(): ( table.search([1, 2, 3], vector_column_name="vector2", fast_search=True) .minimum_nprobes(5) - .maximum_nprobes(0) .to_list() ) @@ -2180,9 +2655,6 @@ def test_query_sync_hybrid(): "refine_factor": None, "vector": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], "vector_column": "vector", - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, "lower_bound": None, "upper_bound": None, "ef": None, @@ -2878,14 +3350,18 @@ def blob_remote_table(*, server_version=Version("0.5.0")): elif request.path == "/v1/table/test/query/": content_len = int(request.headers.get("Content-Length", 0)) body = json.loads(request.rfile.read(content_len)) - assert body["with_row_id"] is True - if body["columns"] == ["image"]: - # fetch_blob_files sizes its handles from a descriptor take. - assert body["filter"].startswith("_rowid IN") - response_table = blob_descriptor_take_table() + columns = body.get("columns") + assert columns in (None, ["id", "image"], ["id"], ["image"]) + if columns == ["id"]: + response_table = blob_query_response_table().select(["id"]) else: - assert body["columns"] == ["id", "image"] - response_table = blob_query_response_table() + assert body["with_row_id"] is True + if columns == ["image"]: + # fetch_blob_files sizes its handles from a descriptor take. + assert body["filter"].startswith("_rowid IN") + response_table = blob_descriptor_take_table() + else: + response_table = blob_query_response_table() request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.file") request.end_headers() @@ -2919,6 +3395,86 @@ def test_remote_blob_columns_and_fetch(): assert blobs.to_pylist() == [b"alpha", None, b"gamma"] +@pytest.mark.asyncio +async def test_async_remote_to_lance_rejects_storage_access(monkeypatch): + with blob_remote_table() as remote_table: + table = remote_table._table + + async def fail_uri(): + raise AssertionError("remote table URI must not be requested") + + monkeypatch.setattr(table, "uri", fail_uri) + with pytest.raises(NotImplementedError, match="remote tables"): + await table.to_lance() + + +@pytest.mark.asyncio +async def test_async_remote_blob_descriptions_to_pandas_uses_server(monkeypatch): + with blob_remote_table() as remote_table: + table = remote_table._table + + async def fail_uri(): + raise AssertionError("remote table URI must not be requested") + + async def fail_to_lance(): + raise AssertionError("remote table storage must not be opened") + + monkeypatch.setattr(table, "uri", fail_uri) + monkeypatch.setattr(table, "to_lance", fail_to_lance) + + query_df = await ( + table.query().select(["id", "image"]).to_pandas(blob_mode="descriptions") + ) + table_df = await table.to_pandas(blob_mode="descriptions") + + for df in (query_df, table_df): + assert df["id"].tolist() == [1, 2, 3] + assert "_rowid" not in df.columns + images = df["image"].tolist() + assert images[1] is None + assert images[0]["size"] == 5 + assert "_lance_row_id" not in images[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blob_mode", ["bytes", "lazy"]) +async def test_async_remote_blob_pandas_rejects_unpinned_fetch(monkeypatch, blob_mode): + with blob_remote_table() as remote_table: + table = remote_table._table + + async def fail_uri(): + raise AssertionError("remote table URI must not be requested") + + async def fail_to_lance(): + raise AssertionError("remote table storage must not be opened") + + async def fail_fetch(*args, **kwargs): + raise AssertionError("blob fetch must not use an unpinned remote table") + + monkeypatch.setattr(table, "uri", fail_uri) + monkeypatch.setattr(table, "to_lance", fail_to_lance) + monkeypatch.setattr(table, "fetch_blobs", fail_fetch) + monkeypatch.setattr(table, "fetch_blob_files", fail_fetch) + + pandas_kwargs = {} if blob_mode == "lazy" else {"blob_mode": blob_mode} + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await table.query().select(["id", "image"]).to_pandas(**pandas_kwargs) + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await table.to_pandas(**pandas_kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blob_mode", ["bytes", "lazy"]) +async def test_async_remote_nonblob_projection_to_pandas(blob_mode): + with blob_remote_table() as remote_table: + df = await ( + remote_table._table.query().select(["id"]).to_pandas(blob_mode=blob_mode) + ) + + assert df["id"].tolist() == [1, 2, 3] + assert list(df.columns) == ["id"] + + def test_remote_blob_files_are_lazy_seekable_handles(): with blob_remote_table() as table: files = table.fetch_blob_files("image", [10, 20, 30, 40]) @@ -2995,9 +3551,15 @@ def test_remote_blob_query_survives_a_server_that_ignores_the_row_id_request(): def test_remote_blob_byte_apis_not_supported_on_old_server(): with blob_remote_table(server_version=Version("0.1.0")) as table: assert table.blob_columns() == ["image"] - with pytest.raises(NotImplementedError, match="not supported"): + with pytest.raises( + NotImplementedError, + match="fetch_blobs is not supported by this LanceDB server", + ): table.fetch_blobs("image", [1]) - with pytest.raises(NotImplementedError, match="not supported"): + with pytest.raises( + NotImplementedError, + match="fetch_blob_files requires LanceDB server 0.5.0 or newer", + ): table.fetch_blob_files("image", [1]) diff --git a/python/python/tests/test_sql.py b/python/python/tests/test_sql.py index eeadb3a33..e6f5fd3c7 100644 --- a/python/python/tests/test_sql.py +++ b/python/python/tests/test_sql.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import threading from uuid import UUID +import numpy as np import pytest import pyarrow as pa @@ -11,7 +13,7 @@ from lancedb import _lancedb from lancedb.arrow import AsyncRecordBatchReader from lancedb.db import AsyncConnection from lancedb.remote.db import RemoteDBConnection -from lancedb.sql import AsyncQuery, Query +from lancedb.sql import AsyncQuery, Query, to_parameter_batch NIL_QUERY_ID = UUID(int=0) @@ -24,12 +26,20 @@ class FakeNativeQuery: class FakeNativeConnection: - async def execute_query_async(self, query, *, default_namespace_path=None): + def __init__(self): + self.parameters = [] + + async def execute_query_async( + self, query, *, default_namespace_path=None, parameters=None + ): + self.parameters.append(parameters) return FakeNativeQuery() class FakeAsyncConnection: - async def execute_query_async(self, query, *, default_namespace_path=None): + async def execute_query_async( + self, query, *, default_namespace_path=None, parameters=None + ): return AsyncQuery(FakeNativeQuery()) @@ -160,3 +170,100 @@ def test_execute_query_async_rejects_invalid_namespace_components( remote_connection().execute_query_async( "SELECT 1", default_namespace_path=default_namespace_path ) + + +def test_positional_parameters_keep_their_types(): + vector = np.array([0.1, 0.2, 0.3], dtype=np.float32) + batch = to_parameter_batch( + [42, 1.5, "a", np.float32(0.25), vector, pa.scalar(7, pa.int16())] + ) + assert batch.num_rows == 1 + assert batch.schema.names == ["1", "2", "3", "4", "5", "6"] + assert batch.schema.types == [ + pa.int64(), + pa.float64(), + pa.string(), + pa.float32(), + pa.list_(pa.float32(), 3), + pa.int16(), + ] + sent = batch.column(4).flatten().to_numpy() + assert sent.dtype == np.float32 + assert sent.tobytes() == vector.tobytes() + + +def test_named_parameters_and_batches(): + half = pa.array(np.array([1.0, 2.0], dtype=np.float16)) + batch = to_parameter_batch({"vector": half, "k": 3}) + assert batch.schema.names == ["vector", "k"] + assert batch.schema.field("vector").type == pa.list_(pa.float16(), 2) + + row = pa.record_batch({"id": [1]}) + assert to_parameter_batch(row) is row + assert to_parameter_batch(pa.table({"id": [1]})).equals(row) + assert to_parameter_batch(None) is None + + +@pytest.mark.parametrize( + "parameters", + [pa.record_batch({"id": [1, 2]}), pa.table({"id": pa.array([], pa.int64())})], +) +def test_parameters_are_exactly_one_row(parameters): + with pytest.raises(ValueError, match="exactly one row"): + to_parameter_batch(parameters) + + +def test_parameters_must_be_a_collection(): + with pytest.raises(TypeError, match="query parameters"): + to_parameter_batch("SELECT 1") + + +@pytest.mark.asyncio +async def test_async_execute_query_sends_parameters_as_a_batch(): + native = FakeNativeConnection() + connection = AsyncConnection(native) + await connection.execute_query("SELECT $1", parameters=[1]) + await connection.execute_query_async("SELECT 1") + sent, unparameterized = native.parameters + assert sent.equals(pa.record_batch({"1": [1]})) + assert unparameterized is None + + +def test_parameters_reach_the_server_as_the_values_that_were_sent(): + """A round trip through the bindings against a Flight server that answers + each exchange with the parameter row it received.""" + flight = pytest.importorskip("pyarrow.flight") + + class EchoServer(flight.FlightServerBase): + def __init__(self): + super().__init__("grpc://127.0.0.1:0") + self.commands = [] + self.lock = threading.Lock() + + def do_exchange(self, context, descriptor, reader, writer): + parameters = reader.read_all() + with self.lock: + self.commands.append(descriptor.command) + writer.begin(parameters.schema) + writer.write_table(parameters) + writer.close() + + vector = np.array([0.1, 1e-45, -0.0], dtype=np.float32) + with EchoServer() as server: + db = lancedb.connect( + "db://analytics", + api_key="test-key", + host_override="http://localhost:10024", + sql_host_override=f"grpc://127.0.0.1:{server.port}", + ) + echoed = db.execute_query( + "SELECT echo", parameters={"vector": vector, "k": 10} + ).read_all() + (command,) = server.commands + + assert b"CommandStatementQuery" in command + assert b"SELECT echo" in command + assert echoed.schema.field("vector").type == pa.list_(pa.float32(), 3) + received = echoed.column("vector").combine_chunks().flatten().to_numpy() + assert received.view(np.uint32).tolist() == vector.view(np.uint32).tolist() + assert echoed.column("k").to_pylist() == [10] diff --git a/python/src/connection.rs b/python/src/connection.rs index 71ba70e0e..ec0768eb2 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -14,9 +14,10 @@ use crate::{ table::Table, }; use arrow::{ + array::RecordBatch, datatypes::Schema, ffi_stream::ArrowArrayStreamReader, - pyarrow::{FromPyArrow, ToPyArrow}, + pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; use lancedb::{ connection::Connection as LanceConnection, @@ -148,18 +149,22 @@ impl Connection { self.get_inner().map(|inner| inner.uri().to_string()) } - #[pyo3(signature = (query, *, default_namespace_path=None))] + #[pyo3(signature = (query, *, default_namespace_path=None, parameters=None))] pub fn execute_query_async<'a>( self_: PyRef<'a, Self>, query: String, default_namespace_path: Option>, + parameters: Option>, ) -> PyResult> { let inner = self_.get_inner()?.clone(); let default_namespace_path = parse_default_namespace_path(default_namespace_path)?; future_into_py(self_.py(), async move { - let operation = inner + let mut operation = inner .execute_query_async(query) .default_namespace_path(default_namespace_path); + if let Some(PyArrowType(parameters)) = parameters { + operation = operation.parameters(parameters); + } operation .execute() .await diff --git a/python/src/index.rs b/python/src/index.rs index 54b15f55e..d4e16a67e 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -54,8 +54,6 @@ pub fn extract_index_params(source: &Option>) -> PyResult>) -> PyResult, lower_case: bool, - stem: bool, - remove_stop_words: bool, + stem: Option, + remove_stop_words: Option, custom_stop_words: Option>, ascii_folding: bool, ngram_min_length: u32, diff --git a/python/src/query.rs b/python/src/query.rs index 1079fb314..88b7243d4 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -336,9 +336,8 @@ pub struct PyQueryRequest { pub use_lsm: Option, pub column: Option, pub query_vector: Option, + pub nprobes: Option, pub minimum_nprobes: Option, - // None means user did not set it and default should be used (currently 20) - // Some(0) means user set it to None and there is no limit pub maximum_nprobes: Option, pub lower_bound: Option, pub upper_bound: Option, @@ -369,6 +368,7 @@ impl From for PyQueryRequest { use_lsm: query_request.use_lsm, column: None, query_vector: None, + nprobes: None, minimum_nprobes: None, maximum_nprobes: None, lower_bound: None, @@ -396,11 +396,9 @@ impl From for PyQueryRequest { use_lsm: vector_query.base.use_lsm, column: vector_query.column, query_vector: Some(PyQueryVectors(vector_query.query_vector)), - minimum_nprobes: Some(vector_query.minimum_nprobes), - maximum_nprobes: match vector_query.maximum_nprobes { - None => Some(0), - Some(value) => Some(value), - }, + nprobes: vector_query.nprobes, + minimum_nprobes: vector_query.minimum_nprobes, + maximum_nprobes: vector_query.maximum_nprobes, lower_bound: vector_query.lower_bound, upper_bound: vector_query.upper_bound, ef: vector_query.ef, diff --git a/python/src/table.rs b/python/src/table.rs index 0daa09263..d82ebbaaf 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -823,8 +823,8 @@ impl From for FtsToken { language = "English".to_string(), max_token_length = Some(40), lower_case = true, - stem = true, - remove_stop_words = true, + stem = None, + remove_stop_words = None, custom_stop_words = None, ascii_folding = true, ngram_min_length = 3, @@ -838,15 +838,15 @@ pub fn tokenize( language: String, max_token_length: Option, lower_case: bool, - stem: bool, - remove_stop_words: bool, + stem: Option, + remove_stop_words: Option, custom_stop_words: Option>, ascii_folding: bool, ngram_min_length: u32, ngram_max_length: u32, prefix_only: bool, ) -> PyResult> { - let params = FtsIndexBuilder::default() + let mut params = FtsIndexBuilder::default() .base_tokenizer(base_tokenizer) .language(&language) .map_err(|_| { @@ -857,13 +857,17 @@ pub fn tokenize( })? .max_token_length(max_token_length.map(|value| value as usize)) .lower_case(lower_case) - .stem(stem) - .remove_stop_words(remove_stop_words) .ascii_folding(ascii_folding) .ngram_min_length(ngram_min_length) .ngram_max_length(ngram_max_length) .ngram_prefix_only(prefix_only) .custom_stop_words(custom_stop_words); + if let Some(stem) = stem { + params = params.stem(stem); + } + if let Some(remove_stop_words) = remove_stop_words { + params = params.remove_stop_words(remove_stop_words); + } let tokens = lancedb_tokenize(&query, ¶ms).infer_error()?; Ok(tokens.into_iter().map(FtsToken::from).collect()) } @@ -913,6 +917,11 @@ impl Table { self.inner.is_some() } + /// Whether this table has a local Lance dataset that the client may open. + pub fn _is_native(&self) -> PyResult { + Ok(self.inner_ref()?.as_native().is_some()) + } + /// Closes the table, releasing any resources associated with it. pub fn close(&mut self) { self.inner.take(); @@ -1708,7 +1717,10 @@ impl Table { future_into_py(self_.py(), async move { inner .as_native() - .ok_or_else(|| PyValueError::new_err("This cannot be run on a remote table"))? + .ok_or_else(|| lancedb::Error::NotSupported { + message: "uses_v2_manifest_paths is not supported for remote tables.".into(), + }) + .infer_error()? .uses_v2_manifest_paths() .await .infer_error() @@ -1720,7 +1732,10 @@ impl Table { future_into_py(self_.py(), async move { inner .as_native() - .ok_or_else(|| PyValueError::new_err("This cannot be run on a remote table"))? + .ok_or_else(|| lancedb::Error::NotSupported { + message: "migrate_manifest_paths_v2 is not supported for remote tables.".into(), + }) + .infer_error()? .migrate_manifest_paths_v2() .await .infer_error() diff --git a/python/uv.lock b/python/uv.lock index f80541711..b013e5982 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -1849,19 +1849,19 @@ wheels = [ [[package]] name = "lance-namespace" -version = "0.11.1" +version = "0.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "lance-namespace-urllib3-client" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bf/93/da5f7fcac690db9b282a3439ed9e34960c147619a0d6e1f4eb8cd240e7a5/lance_namespace-0.11.1.tar.gz", hash = "sha256:f67cfbbe0647b7cb42f23b673e7edf8a75b7d8a047265a916492f8d247ee1bc2", size = 11631, upload-time = "2026-08-18T17:40:06.294Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c0/e7/d5d46594678ee479c0eda830c47b2f5c46133bd100e84e9aa6e01306eca7/lance_namespace-0.13.0.tar.gz", hash = "sha256:24554a0997bdb39595c6e4cb3ac6722069f6cd3bd1a74e7acbba7bfaf774a40d", size = 11636, upload-time = "2026-09-14T02:47:03.871523Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/bc/601f2b3cc4cfa0070d858a33223bc823fffdd7981a25c45984a5216ca952/lance_namespace-0.11.1-py3-none-any.whl", hash = "sha256:07643fce9a42ad4d58cc8bf91e3f592bc7f4cbd8d0ad5233223506debf67551c", size = 13507, upload-time = "2026-08-18T17:40:03.561Z" }, + { url = "https://files.pythonhosted.org/packages/e7/47/cfa33cca1ba7c749fd2918cfc5c8ded788f378cc6424d23e4fead5a14125/lance_namespace-0.13.0-py3-none-any.whl", hash = "sha256:438c7b17aef421c21c138196e715f2510d62f07e865372047f87c1e75e618c7a", size = 13507, upload-time = "2026-09-14T02:47:01.765655Z" }, ] [[package]] name = "lance-namespace-urllib3-client" -version = "0.11.1" +version = "0.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pydantic" }, @@ -1869,9 +1869,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e4/c5/2bdd0ff98b469894c8a73be809d26ffdad5402517b0e5f9e758026cba29e/lance_namespace_urllib3_client-0.11.1.tar.gz", hash = "sha256:145a9e9424d7597487249b5b95ee274423bf2910e1a9160b6a07b676b61ea46a", size = 237345, upload-time = "2026-08-18T17:40:07.308Z" } +sdist = { url = "https://files.pythonhosted.org/packages/69/25/9aaa4a5e8999693fb0f227c2c0c4b97bd8f0539066408cdff0b89d41b2d5/lance_namespace_urllib3_client-0.13.0.tar.gz", hash = "sha256:1e8a79c6e4e6277033597fd76aa0e2f33d909ca1436c935936e9e569221f43ef", size = 238348, upload-time = "2026-09-14T02:47:04.521878Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/2a/eaaefd55d1190291207049fedc6b3eb22b506e57d6de91bae46bbaaa9c60/lance_namespace_urllib3_client-0.11.1-py3-none-any.whl", hash = "sha256:36537f529294da6d884ba0fe783704483f0a75463497c7705fd083a4d0257990", size = 406311, upload-time = "2026-08-18T17:40:04.842Z" }, + { url = "https://files.pythonhosted.org/packages/60/4c/7b8f0712a7fe1b8655b711342bc4989696635a79e7027ee56ba3cae23e99/lance_namespace_urllib3_client-0.13.0-py3-none-any.whl", hash = "sha256:fb361eb4f6c7f2d1f9e92809609657b73e38241393e9c4516e12d6b6ee643a0a", size = 406848, upload-time = "2026-09-14T02:47:02.715954Z" }, ] [[package]] @@ -1980,7 +1980,7 @@ requires-dist = [ { name = "huggingface-hub", marker = "extra == 'embeddings'", specifier = ">=0.19.0" }, { name = "ibm-watsonx-ai", marker = "python_full_version >= '3.10' and extra == 'embeddings'", specifier = ">=1.1.2" }, { name = "instructorembedding", marker = "extra == 'embeddings'", specifier = ">=1.0.1" }, - { name = "lance-namespace", specifier = ">=0.3.2" }, + { name = "lance-namespace", specifier = ">=0.13.0" }, { name = "mkdocs", marker = "extra == 'docs'" }, { name = "mkdocs-jupyter", marker = "extra == 'docs'" }, { name = "mkdocs-material", marker = "extra == 'docs'" }, @@ -2005,7 +2005,7 @@ requires-dist = [ { name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" }, { name = "pydantic", specifier = ">=2.7.4,<3" }, { name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" }, - { name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b6" }, + { name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b10" }, { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" }, { name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" }, { name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" }, @@ -3912,7 +3912,7 @@ crypto = [ [[package]] name = "pylance" -version = "14.0.0b6" +version = "14.0.0b10" source = { registry = "https://pypi.fury.io/lance-format/" } dependencies = [ { name = "lance-namespace" }, @@ -3922,12 +3922,12 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://pypi.fury.io/lance-format/-/ver_DuQEn/pylance-14.0.0b6-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3eade86fa30d9ff3248ee7ffb0b27ce29495394d12327da4d22b5c54bc2b59fe" }, - { url = "https://pypi.fury.io/lance-format/-/ver_K5sqj/pylance-14.0.0b6-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:10e291974ee738a1d5d7e9dad17a634394a91cb60080b5ceac209077fbaec43c" }, - { url = "https://pypi.fury.io/lance-format/-/ver_EnbkT/pylance-14.0.0b6-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:61bac10f5081c1210edd017f81c8c784d800b89446bd0ed05d4a24e7001965bc" }, - { url = "https://pypi.fury.io/lance-format/-/ver_19ffYg/pylance-14.0.0b6-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:3c627ad4a9870611becf84409d8e5ecddc6be64e1dd53bee4bd5fb10ab561e5c" }, - { url = "https://pypi.fury.io/lance-format/-/ver_2eLLY5/pylance-14.0.0b6-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:bc603a7dd690233f1fe000600171e1c092f153bc9060cdf198825c4bbe1418aa" }, - { url = "https://pypi.fury.io/lance-format/-/ver_t8ggc/pylance-14.0.0b6-cp310-abi3-win_amd64.whl", hash = "sha256:1d32df32a3d157611450def07eae8ace3a7401286eac68cad9a24e0d5deed2db" }, + { url = "https://pypi.fury.io/lance-format/-/ver_4hvE/pylance-14.0.0b10-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:415b6b1b9fe735fbae7c755eb9774ae2fe0feddc90bf80bee24087cbfbfdf8c1" }, + { url = "https://pypi.fury.io/lance-format/-/ver_pM1km/pylance-14.0.0b10-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:87317cad22f4d35fea11eb70aa5c0d71a0e1ea3b96de38f3deca933daa749fff" }, + { url = "https://pypi.fury.io/lance-format/-/ver_24jdyj/pylance-14.0.0b10-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6da554a2fc2197ba254676b12824d8120b48d65937697e0661a61302458a4740" }, + { url = "https://pypi.fury.io/lance-format/-/ver_XclaZ/pylance-14.0.0b10-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:36c80a801675bb9dcbd681fa99d0cdd5ca642dfb3ca08d71ac980278a2f7748b" }, + { url = "https://pypi.fury.io/lance-format/-/ver_L0yh/pylance-14.0.0b10-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:e50ad0b866faefb2fa3ae87328b93fa0d06b2bdd5bb5d0c1a6c74b907cc2fd93" }, + { url = "https://pypi.fury.io/lance-format/-/ver_2gbNr8/pylance-14.0.0b10-cp310-abi3-win_amd64.whl", hash = "sha256:3d86451fe00f71510005bb700737fb1a4f9361554541cbf802f3c5072554d014" }, ] [[package]] diff --git a/python/vendor/lance-arrow/Cargo.toml b/python/vendor/lance-arrow/Cargo.toml index e244fac98..cce94be63 100644 --- a/python/vendor/lance-arrow/Cargo.toml +++ b/python/vendor/lance-arrow/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lance-arrow" -version = "14.0.0-beta.7" +version = "14.0.0-beta.10" edition = "2024" authors = ["Lance Devs "] license = "Apache-2.0" @@ -25,7 +25,7 @@ bytemuck = { version = "1", default-features = false, features = ["extern_crate_ bytes.workspace = true futures.workspace = true half = { workspace = true, features = ["std", "bytemuck"] } -jsonb = { version = "0.5.3", default-features = false, features = ["databend"] } +lance-jsonb.workspace = true num-traits.workspace = true rand = { workspace = true, features = ["small_rng"] } diff --git a/python/vendor/lance-arrow/README.lancedb.md b/python/vendor/lance-arrow/README.lancedb.md index 5cfa885ae..b0100c79e 100644 --- a/python/vendor/lance-arrow/README.lancedb.md +++ b/python/vendor/lance-arrow/README.lancedb.md @@ -1,5 +1,5 @@ -This directory contains `rust/lance-arrow` from Lance tag `v14.0.0-beta.7` -(commit `0fd26effc79c62c244a5baf92cdda06e301baf21`), licensed under Apache-2.0. +This directory contains `rust/lance-arrow` from Lance tag `v14.0.0-beta.10` +(commit `fa3638fd816ab2ff9a2ecbba7e51465d8879958f`), licensed under Apache-2.0. LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps until a Lance release includes the fix. The patch is selected in the root diff --git a/python/vendor/lance-arrow/src/json.rs b/python/vendor/lance-arrow/src/json.rs index 9fad06b34..df908c4cb 100644 --- a/python/vendor/lance-arrow/src/json.rs +++ b/python/vendor/lance-arrow/src/json.rs @@ -14,8 +14,8 @@ use arrow_array::{ LargeStringArray, ListArray, MapArray, OffsetSizeTrait, RecordBatch, StringArray, StructArray, }; use arrow_schema::{ArrowError, DataType, Field as ArrowField, Fields, Schema}; -use jsonb::OwnedJsonb; -use jsonb::jsonpath::{JsonPath, Selector}; +use lance_jsonb::OwnedJsonb; +use lance_jsonb::jsonpath::{JsonPath, Selector}; use crate::ARROW_EXT_NAME_KEY; @@ -336,19 +336,19 @@ impl TryFrom for JsonArray { /// Encode JSON string to JSONB format pub fn encode_json(json_str: &str) -> Result, Box> { - let value = jsonb::parse_value(json_str.as_bytes())?; + let value = lance_jsonb::parse_value(json_str.as_bytes())?; Ok(value.to_vec()) } /// Decode JSONB bytes to JSON string pub fn decode_json(jsonb_bytes: &[u8]) -> String { - let raw_jsonb = jsonb::RawJsonb::new(jsonb_bytes); + let raw_jsonb = lance_jsonb::RawJsonb::new(jsonb_bytes); raw_jsonb.to_string() } /// Parse a JSONPath expression. pub fn parse_json_path(path: &str) -> Result, ArrowError> { - jsonb::jsonpath::parse_json_path(path.as_bytes()) + lance_jsonb::jsonpath::parse_json_path(path.as_bytes()) .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid JSONPath '{path}': {e}"))) } @@ -360,7 +360,7 @@ pub fn select_json_path( jsonb_bytes: &[u8], path: &JsonPath<'_>, ) -> Result, ArrowError> { - Selector::new(jsonb::RawJsonb::new(jsonb_bytes)) + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) .select_value(path) .map_err(|e| { ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}")) @@ -369,7 +369,7 @@ pub fn select_json_path( /// Whether `path` matches any value in a JSONB value. pub fn json_path_exists(jsonb_bytes: &[u8], path: &JsonPath<'_>) -> Result { - Selector::new(jsonb::RawJsonb::new(jsonb_bytes)) + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) .exists(path) .map_err(|e| { ArrowError::InvalidArgumentError(format!("Failed to match JSONPath {path}: {e}")) @@ -381,7 +381,7 @@ pub fn select_json_path_values( jsonb_bytes: &[u8], path: &JsonPath<'_>, ) -> Result, ArrowError> { - Selector::new(jsonb::RawJsonb::new(jsonb_bytes)) + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) .select_values(path) .map_err(|e| { ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}")) diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 906ad9136..b15c24ca5 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.41.0-beta.0" +version = "0.41.0-beta.2" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index b299cf4ff..8a0e489c2 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -23,8 +23,8 @@ use crate::connection::create_table::CreateTableBuilder; use crate::data::scannable::Scannable; use crate::database::listing::ListingDatabase; use crate::database::{ - CloneTableRequest, Database, DatabaseOptions, JobInfo, OpenTableRequest, PauseJobStatus, - ReadConsistency, ResumeJobStatus, TableNamesRequest, + CloneTableRequest, Database, DatabaseOptions, ExecuteQueryRequest, JobInfo, OpenTableRequest, + PauseJobStatus, ReadConsistency, ResumeJobStatus, TableNamesRequest, }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; @@ -38,10 +38,7 @@ use crate::remote::{ }, }; use crate::secrets::SecretInfo; -use crate::utils::{ - validate_namespace, validate_secret_component, validate_secret_reference, - validate_view_reference, -}; +use crate::utils::{validate_namespace, validate_secret_reference, validate_view_reference}; use crate::view::ViewDescription; use lance::io::ObjectStoreParams; pub use lance_file::version::LanceFileVersion; @@ -106,7 +103,9 @@ impl TableNamesBuilder { self } - /// The maximum number of table names to return + /// The maximum number of table names to return. + /// + /// Without a limit, all names are returned. Zero returns an empty list. pub fn limit(mut self, limit: u32) -> Self { self.request.limit = Some(limit); self @@ -335,16 +334,14 @@ pub struct CloneTableBuilder { /// Builder for asynchronously executing a SQL statement on a remote database. pub struct ExecuteQueryAsyncBuilder { parent: Arc, - query: String, - default_namespace_path: Vec, + request: ExecuteQueryRequest, } impl ExecuteQueryAsyncBuilder { fn new(parent: Arc, query: String) -> Self { Self { parent, - query, - default_namespace_path: vec!["public".to_string()], + request: ExecuteQueryRequest::new(query), } } @@ -357,15 +354,63 @@ impl ExecuteQueryAsyncBuilder { I: IntoIterator, S: Into, { - self.default_namespace_path = path.into_iter().map(Into::into).collect(); + self.request.default_namespace_path = path.into_iter().map(Into::into).collect(); + self + } + + /// Bind values to the statement's placeholders. + /// + /// `parameters` is a single-row batch. Column `i` binds `$`, and a + /// column whose name is not a number also binds `$`, so + /// `WHERE id = $1` and `WHERE id = $id` both work. The values travel as + /// Arrow rather than as SQL text, so a float keeps its exact bits and type + /// and a vector stays a compact `FixedSizeList`. + /// + /// Every placeholder needs a value and every parameter must be used by a + /// placeholder. Parameters are supported in queries (`SELECT`, and + /// `EXPLAIN` of one), not in DDL or DML, nor in table-function arguments. + /// Do not follow a named placeholder directly with `$`: `$a$` starts a + /// dollar-quoted string. + /// + /// A parameterized statement runs on the call that returns its rows rather + /// than detached from it: cancelling the query, or dropping its handle + /// before reading it, stops the statement on the server. Its status follows + /// that stream too -- it reports `Running` until the reader has received + /// every row -- so open the reader directly rather than waiting for + /// `Finished`. + /// + /// # Example + /// + /// ```no_run + /// # use std::sync::Arc; + /// # use arrow_array::{ArrayRef, FixedSizeListArray, Int64Array, RecordBatch}; + /// # use arrow_array::types::Float32Type; + /// # async fn query(db: &lancedb::Connection) -> lancedb::Result<()> { + /// let vector = FixedSizeListArray::from_iter_primitive::( + /// [Some([0.1_f32, 0.2, 0.3].map(Some))], + /// 3, + /// ); + /// let parameters = RecordBatch::try_from_iter([ + /// ("vector", Arc::new(vector) as ArrayRef), + /// ("k", Arc::new(Int64Array::from(vec![10])) as ArrayRef), + /// ])?; + /// let query = db + /// .execute_query_async("SELECT id FROM docs ORDER BY distance(vector, $vector) LIMIT $k") + /// .parameters(parameters) + /// .execute() + /// .await?; + /// let rows = query.reader().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn parameters(mut self, parameters: RecordBatch) -> Self { + self.request.parameters = Some(parameters); self } /// Start the statement and return its asynchronous query handle. pub async fn execute(self) -> Result { - self.parent - .execute_query_async(&self.query, &self.default_namespace_path) - .await + self.parent.execute_query_async(self.request).await } } @@ -456,8 +501,9 @@ impl Connection { /// /// The query can reference tables in other databases with SQL dot notation. /// Use [`ExecuteQueryAsyncBuilder::default_namespace_path`] to avoid qualifying - /// tables in the default namespace. Local connections return - /// [`Error::NotSupported`]. + /// tables in the default namespace, and [`ExecuteQueryAsyncBuilder::parameters`] + /// to bind Arrow values to `$1` / `$name` placeholders instead of writing them + /// into the SQL text. Local connections return [`Error::NotSupported`]. /// /// # Example /// @@ -506,7 +552,8 @@ impl Connection { /// under creation, may contain only uncommitted storage, or may be concurrently /// dropped before it is opened. /// - /// The parameters `page_token` and `limit` can be used to paginate the results + /// Without a limit, all names are returned. The parameters `start_after` and + /// `limit` can be used to paginate the results. pub fn table_names(&self) -> TableNamesBuilder { TableNamesBuilder::new(self.internal.clone()) } @@ -761,9 +808,7 @@ impl Connection { let database = database.clone(); let namespace_path = namespace_path.clone(); async move { - for segment in &namespace_path { - validate_secret_component("Secret namespace path segment", segment)?; - } + validate_namespace(&namespace_path)?; database.list_secrets(&namespace_path, options).await } }) @@ -1138,7 +1183,21 @@ impl Connection { } /// List tables with pagination support - pub async fn list_tables(&self, request: ListTablesRequest) -> Result { + /// + /// The default limit is 100 tables per page for both local and remote connections, + /// including namespaces. Zero returns an empty page without a continuation token. + /// Follow the response's opaque `page_token` until it is absent to retrieve every table; + /// a page can contain fewer than the limit even when more tables remain. + pub async fn list_tables(&self, mut request: ListTablesRequest) -> Result { + // Apply the SDK's page size here so every backend receives the same request. + let limit = request.limit.get_or_insert(100); + if *limit <= 0 { + return Ok(ListTablesResponse { + context: None, + tables: Vec::new(), + page_token: None, + }); + } self.internal.list_tables(request).await } @@ -2224,6 +2283,61 @@ mod tests { assert_eq!(seen, names); } + #[tokio::test] + async fn test_table_listing_defaults_and_zero_limit() { + let tempdir = tempfile::tempdir().unwrap(); + // Listing discovers physical entries without opening the datasets. + let names: Vec<_> = (0..130).map(|i| format!("t{i:03}")).collect(); + for name in &names { + std::fs::create_dir(tempdir.path().join(format!("{name}.lance"))).unwrap(); + } + let db = connect(tempdir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + assert_eq!(db.table_names().execute().await.unwrap(), names); + assert!( + db.table_names() + .limit(0) + .execute() + .await + .unwrap() + .is_empty() + ); + + let first = db.list_tables(ListTablesRequest::default()).await.unwrap(); + assert_eq!(first.tables, names[..100]); + assert!(first.page_token.is_some()); + let second = db + .list_tables(ListTablesRequest { + page_token: first.page_token, + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(second.tables, names[100..]); + assert!(second.page_token.is_none()); + + let all = db + .list_tables(ListTablesRequest { + limit: Some(200), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(all.tables, names); + assert!(all.page_token.is_none()); + let empty = db + .list_tables(ListTablesRequest { + limit: Some(0), + ..Default::default() + }) + .await + .unwrap(); + assert!(empty.tables.is_empty()); + assert!(empty.page_token.is_none()); + } + #[tokio::test] async fn test_open_table() { let tc = new_test_connection().await.unwrap(); diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index ad4e8118d..05386e0b7 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -18,6 +18,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use arrow_array::RecordBatch; use lance::dataset::ReadParams; use lance_namespace::LanceNamespace; use lance_namespace::models::{ @@ -191,6 +192,32 @@ impl CloneTableRequest { } } +/// A request to start a SQL statement on a remote database. +#[derive(Clone, Debug)] +pub struct ExecuteQueryRequest { + /// The SQL text. + pub query: String, + /// The namespace unqualified table names resolve in. Empty is `public`, + /// the SQL name of the root namespace. + pub default_namespace_path: Vec, + /// Values for the statement's placeholders, as a single-row batch. + /// + /// Column `i` binds `$`, and a column whose name is not a number + /// also binds `$`. See + /// [`ExecuteQueryAsyncBuilder::parameters`](crate::connection::ExecuteQueryAsyncBuilder::parameters). + pub parameters: Option, +} + +impl ExecuteQueryRequest { + pub fn new(query: impl Into) -> Self { + Self { + query: query.into(), + default_namespace_path: vec!["public".to_string()], + parameters: None, + } + } +} + /// How long until a change is reflected from one Table instance to another /// /// Tables are always internally consistent. If a write method is called on @@ -555,8 +582,7 @@ pub trait Database: /// Start executing a SQL statement on a remote database. async fn execute_query_async( &self, - _query: &str, - _default_namespace_path: &[String], + _request: ExecuteQueryRequest, ) -> Result { Err(crate::error::Error::NotSupported { message: "SQL is not supported by this database".to_string(), diff --git a/rust/lancedb/src/database/namespace.rs b/rust/lancedb/src/database/namespace.rs index 6bca29476..d93f73c64 100644 --- a/rust/lancedb/src/database/namespace.rs +++ b/rust/lancedb/src/database/namespace.rs @@ -288,16 +288,44 @@ impl Database for LanceNamespaceDatabase { } async fn table_names(&self, request: TableNamesRequest) -> Result> { - let ns_request = ListTablesRequest { - id: Some(request.namespace_path), - page_token: request.start_after, - limit: request.limit.map(|l| l as i32), - ..Default::default() - }; + if request.limit == Some(0) { + return Ok(Vec::new()); + } - let response = self.namespace.list_tables(ns_request).await?; + // Namespace tokens are opaque, and pages need not arrive in name order. + // Collect the listing before applying the name cursor and overall limit. + let mut names = Vec::new(); + let mut page_token = None; + let mut seen_tokens = HashSet::new(); + loop { + let response = self + .namespace + .list_tables(ListTablesRequest { + id: Some(request.namespace_path.clone()), + page_token, + ..Default::default() + }) + .await?; + names.extend(response.tables); + let Some(token) = response.page_token.filter(|token| !token.is_empty()) else { + break; + }; + if !seen_tokens.insert(token.clone()) { + return Err(Error::Runtime { + message: "Namespace table listing response repeated a page_token".into(), + }); + } + page_token = Some(token); + } - Ok(response.tables) + names.sort(); + if let Some(ref start_after) = request.start_after { + names.retain(|name| name > start_after); + } + if let Some(limit) = request.limit { + names.truncate(limit as usize); + } + Ok(names) } async fn list_tables(&self, request: ListTablesRequest) -> Result { @@ -589,6 +617,135 @@ mod tests { use futures::TryStreamExt; use tempfile::tempdir; + #[derive(Debug)] + struct ScriptedNamespace { + pages: Vec, + requests: Mutex>, + } + + #[async_trait] + impl LanceNamespace for ScriptedNamespace { + fn namespace_id(&self) -> String { + "scripted".into() + } + + async fn list_tables( + &self, + request: ListTablesRequest, + ) -> lance_core::Result { + let mut requests = self.requests.lock().unwrap(); + let page = self.pages[requests.len()].clone(); + requests.push(request); + Ok(page) + } + } + + fn namespace_with_pages( + pages: Vec, + ) -> (LanceNamespaceDatabase, Arc) { + let namespace = Arc::new(ScriptedNamespace { + pages, + requests: Mutex::new(Vec::new()), + }); + let db = LanceNamespaceDatabase::from_namespace_client( + namespace.clone(), + "scripted".into(), + HashMap::new(), + HashMap::new(), + None, + None, + HashSet::new(), + ); + (db, namespace) + } + + fn listing_page(names: &[&str], token: Option<&str>) -> ListTablesResponse { + ListTablesResponse { + tables: names.iter().map(|name| (*name).into()).collect(), + page_token: token.map(String::from), + ..Default::default() + } + } + + #[tokio::test] + #[allow(deprecated)] + async fn test_table_names_collects_namespace_pages() { + for namespace_path in [vec![], vec!["child".into()]] { + for (start_after, limit, expected) in [ + (None, None, vec!["a", "b", "c", "d"]), + (Some("a".into()), Some(2), vec!["b", "c"]), + ] { + let (db, namespace) = namespace_with_pages(vec![ + listing_page(&["d", "b"], Some("opaque-1")), + listing_page(&[], Some("opaque-2")), + listing_page(&["c", "a"], None), + ]); + let names = db + .table_names(TableNamesRequest { + namespace_path: namespace_path.clone(), + start_after, + limit, + }) + .await + .unwrap(); + assert_eq!(names, expected); + + let requests = namespace.requests.lock().unwrap(); + assert_eq!(requests.len(), 3); + for request in requests.iter() { + assert_eq!(request.id.as_ref(), Some(&namespace_path)); + assert_eq!(request.limit, None); + } + assert_eq!(requests[0].page_token, None); + assert_eq!(requests[1].page_token.as_deref(), Some("opaque-1")); + assert_eq!(requests[2].page_token.as_deref(), Some("opaque-2")); + } + } + } + + #[tokio::test] + #[allow(deprecated)] + async fn test_table_names_zero_limit_skips_namespace_request() { + let (db, namespace) = namespace_with_pages(vec![]); + assert!( + db.table_names(TableNamesRequest { + limit: Some(0), + ..Default::default() + }) + .await + .unwrap() + .is_empty() + ); + assert!(namespace.requests.lock().unwrap().is_empty()); + } + + #[tokio::test] + #[allow(deprecated)] + async fn test_table_names_stops_on_empty_namespace_token() { + let (db, namespace) = namespace_with_pages(vec![listing_page(&["a"], Some(""))]); + assert_eq!( + db.table_names(TableNamesRequest::default()).await.unwrap(), + vec!["a"] + ); + assert_eq!(namespace.requests.lock().unwrap().len(), 1); + } + + #[tokio::test] + #[allow(deprecated)] + async fn test_table_names_rejects_namespace_token_cycle() { + let (db, namespace) = namespace_with_pages(vec![ + listing_page(&["a"], Some("one")), + listing_page(&["b"], Some("two")), + listing_page(&["c"], Some("one")), + ]); + let error = db + .table_names(TableNamesRequest::default()) + .await + .unwrap_err(); + assert!(error.to_string().contains("repeated a page_token")); + assert_eq!(namespace.requests.lock().unwrap().len(), 3); + } + /// Helper function to create test data fn create_test_data() -> RecordBatch { let schema = Arc::new(Schema::new(vec![ diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index f05633281..98d915118 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -453,7 +453,8 @@ pub struct FunctionVersion { signature: FunctionSignature, #[serde(default, skip_serializing_if = "Vec::is_empty")] secret_bindings: Vec, - created_at: String, + /// When this version was created, in milliseconds since the epoch. + created_at_millis: i64, metadata: BTreeMap, disabled: bool, } @@ -507,8 +508,9 @@ impl FunctionVersion { &self.secret_bindings } - pub fn created_at(&self) -> &str { - &self.created_at + /// When this version was created, in milliseconds since the epoch. + pub fn created_at_millis(&self) -> i64 { + self.created_at_millis } } diff --git a/rust/lancedb/src/index.rs b/rust/lancedb/src/index.rs index cdc973044..82f8ad845 100644 --- a/rust/lancedb/src/index.rs +++ b/rust/lancedb/src/index.rs @@ -75,6 +75,17 @@ pub enum Index { /// Full text search index using BM25. /// + /// Supports strings, lists of strings, and `lance.json` logical fields + /// stored as `LargeBinary`. Raw binary fields are not supported. Lance + /// selects its JSON-aware tokenizer from the field metadata. + /// + /// JSON queries use `path,str,value` triples, where the dotted path is + /// relative to the indexed column. For example, `meta.tag,str,nature` + /// searches the string value at `$.meta.tag`, not other paths or key names. + /// Create the index on the JSON column, not on a JSON subpath. Plain text + /// queries are invalid for this tokenizer; malformed triples can panic in + /// the current Lance dependency (see [lance#8812](https://github.com/lance-format/lance/issues/8812)). + /// /// The posting block size defaults to 128. Supported values are 128 and 256; /// a value of 256 uses the experimental FTS V3 format and may introduce /// breaking changes. diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 6b20fb123..8fb56ceca 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1114,9 +1114,18 @@ pub struct VectorQueryRequest { pub column: Option, /// The vector(s) to search for pub query_vector: Vec>, - /// The minimum number of partitions to search - pub minimum_nprobes: usize, - /// The maximum number of partitions to search + /// The legacy number of partitions to search. + /// + /// Lance sets both probe bounds to this value. Explicit bounds can override it. + pub nprobes: Option, + /// The minimum number of partitions to search. + /// + /// If unset, Lance's default is used. + pub minimum_nprobes: Option, + /// The maximum number of partitions to search. + /// + /// If unset, Lance's default is used. Zero represents an explicitly unbounded + /// maximum on request and wire representations. pub maximum_nprobes: Option, /// The lower bound (inclusive) of the distance to search for. pub lower_bound: Option, @@ -1141,8 +1150,9 @@ impl Default for VectorQueryRequest { base: QueryRequest::default(), column: None, query_vector: Vec::new(), - minimum_nprobes: 20, - maximum_nprobes: Some(20), + nprobes: None, + minimum_nprobes: None, + maximum_nprobes: None, lower_bound: None, upper_bound: None, ef: None, @@ -1248,7 +1258,7 @@ impl VectorQuery { Ok(self) } - /// Set the number of partitions to search (probe) + /// Set the legacy IVF probe parameter /// /// This argument is only used when the vector column has an IVF PQ index. /// If there is no index then this value is ignored. @@ -1258,19 +1268,19 @@ impl VectorQuery { /// /// The partition whose centroids are closest to the query vector will be /// exhaustiely searched to find matches. This parameter controls how many - /// partitions should be searched. + /// partitions may be searched. /// /// Increasing this value will increase the recall of your query but will - /// also increase the latency of your query. The default value is 20. This - /// default is good for many cases but the best value to use will depend on - /// your data and the recall that you need to achieve. + /// also increase the latency of your query. If this method is not called, + /// Lance's adaptive probe defaults are used. /// /// For best results we recommend tuning this parameter with a benchmark against /// your actual data to find the smallest possible value that will still give /// you the desired recall. /// - /// This method sets both the minimum and maximum number of partitions to search. - /// For more fine-grained control see [`VectorQuery::minimum_nprobes`] and + /// LanceDB retains this as `nprobes` through local and remote request + /// construction. Lance sets both probe bounds to this value. For more + /// fine-grained control see [`VectorQuery::minimum_nprobes`] and /// [`VectorQuery::maximum_nprobes`]. /// /// Returns an error if `nprobes` is not greater than 0. @@ -1280,8 +1290,7 @@ impl VectorQuery { message: "nprobes must be greater than 0".to_string(), }); } - self.request.minimum_nprobes = nprobes; - self.request.maximum_nprobes = Some(nprobes); + self.request.nprobes = Some(nprobes); Ok(self) } @@ -1303,6 +1312,7 @@ impl VectorQuery { }); } if let Some(maximum_nprobes) = self.request.maximum_nprobes + && maximum_nprobes != 0 && minimum_nprobes > maximum_nprobes { return Err(Error::InvalidInput { @@ -1310,7 +1320,7 @@ impl VectorQuery { .to_string(), }); } - self.request.minimum_nprobes = minimum_nprobes; + self.request.minimum_nprobes = Some(minimum_nprobes); Ok(self) } @@ -1327,7 +1337,7 @@ impl VectorQuery { /// This can be useful when there is a narrow filter to allow these queries to /// spend more time searching and avoid potential false negatives. /// - /// Set to None to search all partitions, if needed, to satisfy the limit + /// Set to None to use Lance's default maximum probe behavior. pub fn maximum_nprobes(mut self, maximum_nprobes: Option) -> Result { if let Some(maximum_nprobes) = maximum_nprobes { if maximum_nprobes == 0 { @@ -1335,14 +1345,18 @@ impl VectorQuery { message: "maximum_nprobes must be greater than 0".to_string(), }); } - if maximum_nprobes < self.request.minimum_nprobes { + if let Some(minimum_nprobes) = self.request.minimum_nprobes + && maximum_nprobes < minimum_nprobes + { return Err(Error::InvalidInput { message: "maximum_nprobes must be greater than or equal to minimum_nprobes" .to_string(), }); } } - self.request.maximum_nprobes = maximum_nprobes; + // Preserve an explicit reset separately from an untouched maximum. The + // request protocols already use zero to represent an unbounded maximum. + self.request.maximum_nprobes = Some(maximum_nprobes.unwrap_or(0)); Ok(self) } @@ -2431,6 +2445,9 @@ mod tests { let vector = Float32Array::from_iter_values([0.1, 0.2]); let query = table.query().nearest_to(&[0.1, 0.2]).unwrap(); + assert_eq!(query.request.nprobes, None); + assert_eq!(query.request.minimum_nprobes, None); + assert_eq!(query.request.maximum_nprobes, None); assert_eq!( *query .request @@ -2442,6 +2459,37 @@ mod tests { vector ); + let minimum_only_query = query.clone().minimum_nprobes(20).unwrap(); + assert_eq!(minimum_only_query.request.minimum_nprobes, Some(20)); + assert_eq!(minimum_only_query.request.maximum_nprobes, None); + + let maximum_only_query = query.maximum_nprobes(Some(20)).unwrap(); + assert_eq!(maximum_only_query.request.minimum_nprobes, None); + assert_eq!(maximum_only_query.request.maximum_nprobes, Some(20)); + + let explicitly_unbounded_query = table + .query() + .nearest_to(&[0.1, 0.2]) + .unwrap() + .nprobes(20) + .unwrap() + .maximum_nprobes(None) + .unwrap(); + assert_eq!(explicitly_unbounded_query.request.nprobes, Some(20)); + assert_eq!(explicitly_unbounded_query.request.maximum_nprobes, Some(0)); + + let bounded_query = table + .query() + .nearest_to(&[0.1, 0.2]) + .unwrap() + .minimum_nprobes(5) + .unwrap() + .nprobes(20) + .unwrap(); + assert_eq!(bounded_query.request.nprobes, Some(20)); + assert_eq!(bounded_query.request.minimum_nprobes, Some(5)); + assert_eq!(bounded_query.request.maximum_nprobes, None); + let new_vector = Float32Array::from_iter_values([9.8, 8.7]); let query = table @@ -2450,7 +2498,7 @@ mod tests { .offset(1) .nearest_to(&[9.8, 8.7]) .unwrap() - .nprobes(1000) + .nprobes(20) .unwrap() .postfilter() .distance_type(DistanceType::Cosine) @@ -2469,8 +2517,9 @@ mod tests { ); assert_eq!(query.request.base.limit.unwrap(), 100); assert_eq!(query.request.base.offset.unwrap(), 1); - assert_eq!(query.request.minimum_nprobes, 1000); - assert_eq!(query.request.maximum_nprobes, Some(1000)); + assert_eq!(query.request.nprobes, Some(20)); + assert_eq!(query.request.minimum_nprobes, None); + assert_eq!(query.request.maximum_nprobes, None); assert!(query.request.use_index); assert_eq!(query.request.distance_type, Some(DistanceType::Cosine)); assert_eq!(query.request.approx_mode, Some(ApproxMode::Accurate)); @@ -2492,13 +2541,13 @@ mod tests { Err(Error::InvalidInput { message }) if message == "nprobes must be greater than 0" )); - // Fixed probe counts replace both bounds, even when moving below or - // above the previous range. + // Legacy probe counts remain independent until Lance resolves them. let mut query = query; for nprobes in [30, 1, 50] { query = query.nprobes(nprobes).unwrap(); - assert_eq!(query.request.minimum_nprobes, nprobes); - assert_eq!(query.request.maximum_nprobes, Some(nprobes)); + assert_eq!(query.request.nprobes, Some(nprobes)); + assert_eq!(query.request.minimum_nprobes, None); + assert_eq!(query.request.maximum_nprobes, None); } } diff --git a/rust/lancedb/src/remote/catalog.rs b/rust/lancedb/src/remote/catalog.rs index 72b77a87e..fa6b185cc 100644 --- a/rust/lancedb/src/remote/catalog.rs +++ b/rust/lancedb/src/remote/catalog.rs @@ -19,6 +19,7 @@ use crate::catalog::{ ListDatabasesResponse, }; use crate::database::Database; +use crate::utils::validate_database_name; use crate::{Error, Result}; /// Authentication and client settings shared by a catalog and its databases. @@ -154,29 +155,6 @@ impl RemoteCatalog { }) } - fn validate_name(&self, name: &str) -> Result<()> { - let delimiter = self - .options - .client_config - .id_delimiter - .as_deref() - .unwrap_or("$"); - if name.is_empty() - || name.trim() != name - || !name.is_ascii() - || name.chars().any(char::is_control) - || name.contains(delimiter) - || matches!(name, "." | "..") - { - return Err(Error::InvalidInput { - message: format!( - "Invalid database name '{name}': expected a nonempty ASCII name without surrounding whitespace, control characters, or namespace delimiter '{delimiter}'" - ), - }); - } - Ok(()) - } - fn database(&self, name: &str) -> Result> { Ok(Arc::new(RemoteDatabase::for_catalog( &self.endpoint, @@ -203,7 +181,7 @@ impl Catalog for RemoteCatalog { } async fn create_database(&self, request: CreateDatabaseRequest) -> Result> { - self.validate_name(&request.name)?; + validate_database_name(&request.name)?; self.root .create_namespace(CreateNamespaceRequest { id: Some(vec![request.name.clone()]), @@ -231,7 +209,7 @@ impl Catalog for RemoteCatalog { } async fn drop_database(&self, request: DropDatabaseRequest) -> Result<()> { - self.validate_name(&request.name)?; + validate_database_name(&request.name)?; let result = self .root .drop_namespace(DropNamespaceRequest { @@ -286,7 +264,7 @@ impl Catalog for RemoteCatalog { } async fn connect_database(&self, name: &str) -> Result> { - self.validate_name(name)?; + validate_database_name(name)?; self.root .describe_namespace(DescribeNamespaceRequest { id: Some(vec![name.into()]), @@ -676,7 +654,24 @@ mod tests { } let catalog = RemoteCatalog::try_new("http://127.0.0.1:1", RemoteCatalogOptions::default()).unwrap(); - for name in ["", "a$b", "\r\ninjected", "..", "café", " padded "] { + for name in [ + "", + "/", + "//", + "a$b", + "\r\ninjected", + "..", + "café", + " padded ", + "a b", + "a?b", + "/db", + "db/", + "a//b", + "a/../b", + "a/./b", + &"a".repeat(65), + ] { assert!(matches!( catalog.create_database(name.into()).await, Err(Error::InvalidInput { .. }) diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index 705506705..5f24d24c7 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -12,6 +12,7 @@ use std::{collections::HashMap, future::Future, str::FromStr, sync::Arc, time::D use crate::error::{Error, Result}; use crate::remote::db::RemoteOptions; use crate::remote::retry::{ResolvedRetryConfig, RetryCounter}; +use crate::utils::validate_database_name; const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id"); @@ -383,6 +384,7 @@ pub fn parse_db_url(db_url: &str) -> Result { message: format!("Invalid encoded database name: {err}"), })? .into_owned(); + validate_database_name(&db_name)?; let db_prefix = { let prefix = parsed_url.path().trim_start_matches('/'); if prefix.is_empty() { @@ -1204,6 +1206,29 @@ mod tests { assert_eq!(parsed.db_prefix.as_deref(), Some("prefix")); } + #[test] + fn test_parse_database_uri_validates_decoded_name() { + for name in [ + "a b", + "a$b", + "a?b", + "café", + "/db", + "db/", + "a//b", + "a/../b", + "a/./b", + &"a".repeat(65), + ] { + let uri = format!("db://{}", urlencoding::encode(name)); + assert!( + matches!(parse_db_url(&uri), Err(Error::InvalidInput { .. })), + "accepted {name:?}" + ); + } + assert!(parse_db_url(&format!("db://{}", "a".repeat(64))).is_ok()); + } + #[test] fn test_timeout_config_default() { let config = TimeoutConfig::default(); diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index c44f7558e..28582cdc3 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use async_trait::async_trait; @@ -20,8 +20,8 @@ use lance_namespace::models::{ use crate::Error; use crate::database::{ - CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, JobInfo, - OpenTableRequest, ReadConsistency, TableNamesRequest, + CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, + ExecuteQueryRequest, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest, }; use crate::error::Result; use crate::function::{ @@ -36,7 +36,10 @@ use crate::remote::util::stream_as_body; use crate::secrets::SecretBinding; use crate::secrets::SecretInfo; use crate::table::BaseTable; -use crate::utils::{reject_relative_segment, validate_table_name}; +use crate::utils::{ + validate_database_name, validate_function_name, validate_namespace, validate_secret_name, + validate_table_name, validate_view_name, +}; use crate::view::ViewDescription; use super::client::{ @@ -230,6 +233,7 @@ pub struct RemoteDatabase { /// TLS configuration for mTLS support tls_config: Option, sql_client: Option, + is_catalog_root: bool, } #[derive(Clone)] @@ -320,6 +324,9 @@ impl RemoteDatabase { name: Option<&str>, options: &super::catalog::RemoteCatalogOptions, ) -> Result { + if let Some(name) = name { + validate_database_name(name)?; + } let scope = super::catalog::ScopedHeaderProvider { provider: options.client_config.header_provider.clone(), database: name.map(str::to_string), @@ -348,6 +355,7 @@ impl RemoteDatabase { )?; if name.is_none() { db.sql_client = None; + db.is_catalog_root = true; } Ok(db) } @@ -422,11 +430,38 @@ impl RemoteDatabase { namespace_context_provider, tls_config: client_config.tls_config, sql_client: Some(sql_client), + is_catalog_root: false, }) } } impl RemoteDatabase { + // A catalog's root namespaces are database names, which may contain slashes. + fn namespace_identifier(&self, namespace: &[String]) -> Result { + match namespace.first() { + None => Ok(ID_DELIMITER.to_string()), + Some(first) => { + if self.is_catalog_root { + // When dealing with the catalog root, the first component is a + // database name. Those need to be validated differently because + // they can contain slashes. Also, we need to use percent encoding + // for the slashes. + validate_database_name(first)?; + validate_namespace(&namespace[1..])?; + Ok(namespace + .iter() + .map(|component| urlencoding::encode(component).into_owned()) + .collect::>() + .join(ID_DELIMITER)) + } else { + // If we're not dealing with the catalog root, build the namespace identifier in the + // usual way. + build_namespace_identifier(namespace) + } + } + } + } + /// Post a request whose body carries a credential. /// /// Shared by the create and alter verbs, which declare their own request @@ -537,6 +572,7 @@ mod test_utils { namespace_context_provider: None, tls_config: None, sql_client: None, + is_catalog_root: false, } } @@ -560,6 +596,7 @@ mod test_utils { namespace_context_provider, tls_config: config.tls_config.clone(), sql_client: None, + is_catalog_root: false, } } } @@ -581,93 +618,60 @@ impl From<&CreateTableMode> for &'static str { } } -/// The path segment addressing one object: its namespace path and its name. -/// -/// One builder for tables, Secrets, Functions and materialized views: the -/// identifier grammar belongs to the namespace spec, not to an object type. An -/// empty path addresses an object with no namespace. -/// -/// Components are checked for addressability, not a character set. The name's -/// grammar is the caller's, so a table reports [`Error::InvalidTableName`], a -/// Function admits names a table may not, and a catalog database carries the -/// `/` that [`RemoteCatalog`] allows. -/// -/// [`RemoteCatalog`]: super::catalog::RemoteCatalog -fn build_object_identifier(what: &str, name: &str, namespace: &[String]) -> Result { - for segment in namespace { - reject_unaddressable_component("namespace segment", segment)?; +fn build_compound_identifier(name: &str, namespace: &[String]) -> Result { + validate_namespace(namespace)?; + let mut identifier = namespace.join("$"); + if !namespace.is_empty() { + identifier.push('$'); } - reject_unaddressable_component(what, name)?; - Ok(join_identifier( - namespace.iter().map(String::as_str).chain([name]), - )) + identifier.push_str(name); + Ok(identifier) } -/// What a component may not be if the join is to survive being split back -/// apart: empty, a segment URL parsing resolves away, or the delimiter itself. -/// -/// Each erases a boundary no encoding of the joined form recovers. `["prod", -/// ""]` joins to `prod$`, which reads back as `["prod"]`, so a drop reaches the -/// parent of the namespace the caller named. -/// -/// Not a character set: per-component percent-encoding makes the wider set -/// safe, since a `/` in a name reaches the service as `%2F`, still one -/// segment. -fn reject_unaddressable_component(what: &str, value: &str) -> Result<()> { - if value.is_empty() { - return Err(Error::InvalidInput { - message: format!( - "{what} must not be empty: the identifier would carry two delimiters in a row, \ - and splitting it back apart would name a different object" - ), - }); - } - reject_relative_segment(what, value)?; - if value.contains(ID_DELIMITER) { - return Err(Error::InvalidInput { - message: format!( - "{what} '{value}' contains the identifier delimiter '{ID_DELIMITER}', so the \ - namespace path and the name it joins could not be told apart" - ), - }); - } - Ok(()) -} - -/// The path segment addressing one table. A wrapper for the error type: -/// callers match on [`Error::InvalidTableName`]. +/// Build a full path to a table. Note that we do not need to percent-encode +/// this when using it in a URL because the validation functions ensure that +/// only URL-safe characters are present. fn build_table_identifier(name: &str, namespace: &[String]) -> Result { validate_table_name(name)?; - build_object_identifier("table name", name, namespace) + build_compound_identifier(name, namespace) } -/// Join components into the `{id}` a route addresses: each percent-encoded, -/// then joined by the delimiter. -/// -/// Per component rather than over the joined string, so the delimiter stays a -/// delimiter and nothing inside a component can end the path segment. -/// -/// A second line, not the first: a component from the object charset is all -/// unreserved and encodes to itself, so the route reads as the caller wrote it. -/// It does not cover `.` and `..`, which are unreserved too and resolve away -/// after decoding -- [`build_object_identifier`] refuses those. -fn join_identifier<'a>(components: impl Iterator) -> String { - components - .map(|component| urlencoding::encode(component).into_owned()) - .collect::>() - .join(ID_DELIMITER) -} - -/// The path segment addressing one namespace. +/// Translate an array of namespace components into a namespace string. +/// Note that we do not need to percent-encode this when using it in a URL +/// because validate_namespace ensures that only URL-safe characters are +/// present. fn build_namespace_identifier(namespace: &[String]) -> Result { - for segment in namespace { - reject_unaddressable_component("namespace segment", segment)?; - } + validate_namespace(namespace)?; if namespace.is_empty() { // According to the namespace spec, use delimiter to represent root namespace - return Ok(ID_DELIMITER.to_string()); + Ok(ID_DELIMITER.to_string()) + } else { + Ok(namespace.join(ID_DELIMITER)) } - Ok(join_identifier(namespace.iter().map(String::as_str))) +} + +/// Build a full path to a function. Note that we do not need to percent-encode +/// this when using it in a URL because the validation functions ensure that +/// only URL-safe characters are present. +fn build_function_identifier(name: &str, namespace: &[String]) -> Result { + validate_function_name(name)?; + build_compound_identifier(name, namespace) +} + +/// Build a full path to a secret. Note that we do not need to percent-encode +/// this when using it in a URL because the validation functions ensure that +/// only URL-safe characters are present. +fn build_secret_identifier(name: &str, namespace: &[String]) -> Result { + validate_secret_name(name)?; + build_compound_identifier(name, namespace) +} + +/// Build a full path to a view. Note that we do not need to percent-encode +/// this when using it in a URL because the validation functions ensure that +/// only URL-safe characters are present. +fn build_view_identifier(name: &str, namespace: &[String]) -> Result { + validate_view_name(name)?; + build_compound_identifier(name, namespace) } /// Build a secure cache key using length prefixes. @@ -945,7 +949,7 @@ impl Database for RemoteDatabase { request: FunctionRegistrationRequest, namespace_path: &[String], ) -> Result> { - let function_id = build_object_identifier("Function name", &request.name, namespace_path)?; + let function_id = build_function_identifier(&request.name, namespace_path)?; let req = self .client .post(&format!("/v1/function/{function_id}/create")) @@ -976,7 +980,7 @@ impl Database for RemoteDatabase { version: &str, namespace_path: &[String], ) -> Result { - let function_id = build_object_identifier("Function name", name, namespace_path)?; + let function_id = build_function_identifier(name, namespace_path)?; let req = self .client .post(&format!("/v1/function/{function_id}/describe")) @@ -1036,7 +1040,7 @@ impl Database for RemoteDatabase { version: &str, namespace_path: &[String], ) -> Result<(bool, Job)> { - let function_id = build_object_identifier("Function name", name, namespace_path)?; + let function_id = build_function_identifier(name, namespace_path)?; let req = self .client .post(&format!("/v1/function/{function_id}/drop")) @@ -1081,7 +1085,7 @@ impl Database for RemoteDatabase { value: &str, namespace_path: &[String], ) -> Result<()> { - let secret_id = build_object_identifier("Secret name", name, namespace_path)?; + let secret_id = build_secret_identifier(name, namespace_path)?; self.post_secret_write( &format!("/v1/secret/{secret_id}/create"), &RemoteCreateSecretRequest { value }, @@ -1090,7 +1094,7 @@ impl Database for RemoteDatabase { } async fn alter_secret(&self, name: &str, value: &str, namespace_path: &[String]) -> Result<()> { - let secret_id = build_object_identifier("Secret name", name, namespace_path)?; + let secret_id = build_secret_identifier(name, namespace_path)?; self.post_secret_write( &format!("/v1/secret/{secret_id}/alter"), &RemoteAlterSecretRequest { value }, @@ -1127,7 +1131,7 @@ impl Database for RemoteDatabase { } async fn drop_secret(&self, name: &str, namespace_path: &[String]) -> Result<()> { - let secret_id = build_object_identifier("Secret name", name, namespace_path)?; + let secret_id = build_secret_identifier(name, namespace_path)?; let req = self.client.post(&format!("/v1/secret/{secret_id}/drop")); let (request_id, response) = self.client.send(req).await?; self.client.check_response(&request_id, response).await?; @@ -1135,7 +1139,7 @@ impl Database for RemoteDatabase { } async fn describe_secret(&self, name: &str, namespace_path: &[String]) -> Result { - let secret_id = build_object_identifier("Secret name", name, namespace_path)?; + let secret_id = build_secret_identifier(name, namespace_path)?; let req = self .client .post(&format!("/v1/secret/{secret_id}/describe")); @@ -1150,7 +1154,7 @@ impl Database for RemoteDatabase { query: &str, namespace_path: &[String], ) -> Result { - let view_id = build_object_identifier("View name", name, namespace_path)?; + let view_id = build_view_identifier(name, namespace_path)?; let req = self .client .post(&format!("/v1/view/{view_id}/create")) @@ -1167,7 +1171,7 @@ impl Database for RemoteDatabase { name: &str, namespace_path: &[String], ) -> Result { - let view_id = build_object_identifier("View name", name, namespace_path)?; + let view_id = build_view_identifier(name, namespace_path)?; let req = self.client.post(&format!("/v1/view/{view_id}/describe")); let (request_id, response) = self.client.send(req).await?; let response = self.client.check_response(&request_id, response).await?; @@ -1184,7 +1188,7 @@ impl Database for RemoteDatabase { } async fn drop_view_async(&self, name: &str, namespace_path: &[String]) -> Result { - let view_id = build_object_identifier("View name", name, namespace_path)?; + let view_id = build_view_identifier(name, namespace_path)?; let req = self.client.post(&format!("/v1/view/{view_id}/drop")); let (request_id, response) = self.client.send(req).await?; let response = self.client.check_response(&request_id, response).await?; @@ -1329,18 +1333,14 @@ impl Database for RemoteDatabase { }) } - async fn execute_query_async( - &self, - query: &str, - default_namespace_path: &[String], - ) -> Result { + async fn execute_query_async(&self, request: ExecuteQueryRequest) -> Result { let client = self .sql_client .as_ref() .ok_or_else(|| Error::NotSupported { message: "SQL is unavailable for this remote database client".to_string(), })?; - client.submit(query, default_namespace_path).await + client.execute(request).await } async fn describe_query(&self, query_id: uuid::Uuid) -> Result { @@ -1354,25 +1354,49 @@ impl Database for RemoteDatabase { } async fn table_names(&self, request: TableNamesRequest) -> Result> { + // A zero limit is an empty listing, even on servers that reject limit=0. + if request.limit == Some(0) { + return Ok(Vec::new()); + } let (tables, version) = if request.namespace_path.is_empty() { // The flat route resumes after a table name and orders by name, which is exactly - // what `start_after` means, so the server does the paging. - let mut req = self.client.get("/v1/table/"); - if let Some(limit) = request.limit { - req = req.query(&[("limit", limit)]); + // what `start_after` means. Subsequent cursors come from the server and may be + // opaque. Follow them until the requested limit or the end of the listing. + let mut tables = Vec::new(); + let mut version = None; + let mut page_token = request.start_after.clone(); + let mut seen_tokens = HashSet::new(); + if let Some(ref token) = page_token { + seen_tokens.insert(token.clone()); } - if let Some(ref start_after) = request.start_after { - req = req.query(&[("page_token", start_after)]); + loop { + let mut req = self.client.get("/v1/table/"); + if let Some(limit) = request.limit { + req = req.query(&[("limit", limit as usize - tables.len())]); + } + if let Some(ref token) = page_token { + req = req.query(&[("page_token", token)]); + } + let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?; + let rsp = self.client.check_response(&request_id, rsp).await?; + if version.is_none() { + version = Some(parse_server_version(&request_id, &rsp)?); + } + let response: ListTablesResponse = rsp.json().await.err_to_http(request_id)?; + tables.extend(response.tables); + if let Some(limit) = request.limit + && tables.len() >= limit as usize + { + tables.truncate(limit as usize); + break; + } + // Empty or repeated tokens must not restart the listing or loop forever. + match response.page_token.filter(|token| !token.is_empty()) { + Some(token) if seen_tokens.insert(token.clone()) => page_token = Some(token), + _ => break, + } } - let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?; - let rsp = self.client.check_response(&request_id, rsp).await?; - let version = parse_server_version(&request_id, &rsp)?; - let tables = rsp - .json::() - .await - .err_to_http(request_id)? - .tables; - (tables, version) + (tables, version.unwrap_or_default()) } else { self.table_names_in_namespace(&request).await? }; @@ -1387,7 +1411,7 @@ impl Database for RemoteDatabase { async fn list_tables(&self, request: ListTablesRequest) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let mut req = self .client .get(&format!("/v1/namespace/{}/table/list", namespace_id)); @@ -1654,7 +1678,7 @@ impl Database for RemoteDatabase { request: ListNamespacesRequest, ) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let mut req = self .client .get(&format!("/v1/namespace/{}/list", namespace_id)); @@ -1676,7 +1700,7 @@ impl Database for RemoteDatabase { request: CreateNamespaceRequest, ) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let mut req = self .client .post(&format!("/v1/namespace/{}/create", namespace_id)); @@ -1707,7 +1731,7 @@ impl Database for RemoteDatabase { async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let mut req = self .client .post(&format!("/v1/namespace/{}/drop", namespace_id)); @@ -1741,7 +1765,7 @@ impl Database for RemoteDatabase { request: DescribeNamespaceRequest, ) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let req = self .client .post(&format!("/v1/namespace/{}/describe", namespace_id)) @@ -1867,6 +1891,40 @@ mod tests { remote::{ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE}, }; + #[tokio::test] + async fn test_catalog_namespace_identifier_validates_all_namespace_components() { + let mut db = super::RemoteDatabase::new_mock(|_| -> http::Response { + panic!("identifier validation must not send a request") + }); + db.is_catalog_root = true; + + for (path, expected) in [ + (vec![], "$"), + (vec!["team/search"], "team%2Fsearch"), + (vec!["team/search", "ns"], "team%2Fsearch$ns"), + (vec!["team/search", "ns", "child"], "team%2Fsearch$ns$child"), + ] { + let path = path.into_iter().map(String::from).collect::>(); + assert_eq!(db.namespace_identifier(&path).unwrap(), expected); + } + + for path in [ + vec!["../search"], + vec!["team/search", "bad/name"], + vec!["team/search", "ns", "bad/name"], + vec!["team/search", "public", "child"], + ] { + let path = path.into_iter().map(String::from).collect::>(); + assert!( + matches!( + db.namespace_identifier(&path), + Err(Error::InvalidInput { .. }) + ), + "accepted {path:?}" + ); + } + } + #[test] fn test_cache_key_security() { // Test that cache keys are unique regardless of delimiter manipulation @@ -2140,6 +2198,133 @@ mod tests { assert_eq!(names, vec!["table1", "table2"]); } + #[rstest::rstest] + #[case(None, None)] + #[case(Some("t009"), None)] + #[case(None, Some(200))] + #[case(Some("t009"), Some(15))] + #[tokio::test] + async fn test_table_names_follows_server_pages( + #[case] start_after: Option<&str>, + #[case] limit: Option, + ) { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/table/"); + let query: HashMap<_, _> = request.url().query_pairs().collect(); + let start = query + .get("page_token") + .map(|token| token[1..4].parse::().unwrap() + 1) + .unwrap_or(0); + let limit = query + .get("limit") + .map(|limit| limit.parse::().unwrap()) + .unwrap_or(10); + // A server may return a short page even when more tables remain. + let end = (start + limit.min(10)).min(130); + let tables: Vec<_> = (start..end).map(|i| format!("t{i:03}")).collect(); + let page_token = (end < 130).then(|| format!("t{:03}.lance/", end - 1)); + http::Response::builder() + .status(200) + .body(serde_json::json!({"tables": tables, "page_token": page_token}).to_string()) + .unwrap() + }); + let mut op = conn.table_names(); + if let Some(start_after) = start_after { + op = op.start_after(start_after); + } + if let Some(limit) = limit { + op = op.limit(limit); + } + let start = if start_after.is_some() { 10 } else { 0 }; + let end = limit.map(|limit| (start + limit).min(130)).unwrap_or(130); + let expected: Vec<_> = (start..end).map(|i| format!("t{i:03}")).collect(); + assert_eq!(op.execute().await.unwrap(), expected); + } + + #[rstest::rstest] + #[case(vec![])] + #[case(vec!["ns".to_string()])] + #[tokio::test] + async fn test_table_listing_zero_limit_never_sends_a_request(#[case] namespace: Vec) { + let conn = Connection::new_with_handler(|_| -> http::Response { + panic!("a zero limit must not be sent to the server") + }); + assert!( + conn.table_names() + .namespace(namespace.clone()) + .limit(0) + .execute() + .await + .unwrap() + .is_empty() + ); + let page = conn + .list_tables(lance_namespace::models::ListTablesRequest { + id: Some(namespace), + limit: Some(0), + ..Default::default() + }) + .await + .unwrap(); + assert!(page.tables.is_empty()); + assert!(page.page_token.is_none()); + } + + #[rstest::rstest] + #[case(vec![], "/v1/namespace/$/table/list")] + #[case(vec!["ns".to_string()], "/v1/namespace/ns/table/list")] + #[tokio::test] + async fn test_list_tables_default_page_size( + #[case] namespace: Vec, + #[case] path: &'static str, + ) { + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.url().path(), path); + let query: HashMap<_, _> = request.url().query_pairs().collect(); + assert_eq!(query.get("limit").map(|limit| limit.as_ref()), Some("100")); + let start = match query.get("page_token") { + None => 0, + Some(token) => { + assert_eq!(token, "opaque-token"); + 100 + } + }; + let end = (start + 100).min(130); + let tables: Vec<_> = (start..end).map(|i| format!("t{i:03}")).collect(); + http::Response::builder() + .status(200) + .body( + serde_json::json!({ + "tables": tables, + "page_token": (end < 130).then_some("opaque-token") + }) + .to_string(), + ) + .unwrap() + }); + let first = conn + .list_tables(lance_namespace::models::ListTablesRequest { + id: Some(namespace.clone()), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(first.tables.len(), 100); + assert_eq!(first.page_token.as_deref(), Some("opaque-token")); + let second = conn + .list_tables(lance_namespace::models::ListTablesRequest { + id: Some(namespace), + page_token: first.page_token, + ..Default::default() + }) + .await + .unwrap(); + let expected: Vec<_> = (100..130).map(|i| format!("t{i:03}")).collect(); + assert_eq!(second.tables, expected); + assert!(second.page_token.is_none()); + } + #[tokio::test] async fn test_table_names_in_a_namespace_never_invents_a_page_token() { // The namespace route's token belongs to the store, so `table_names` cannot build one @@ -2185,8 +2370,11 @@ mod tests { assert_eq!(names, vec!["widgets"]); } + #[rstest::rstest] + #[case(vec![])] + #[case(vec!["ns".to_string()])] #[tokio::test] - async fn test_table_names_in_a_namespace_stops_on_a_repeated_token() { + async fn test_table_names_stops_on_a_repeated_token(#[case] namespace: Vec) { // A server that handed back the token it was given would never finish the walk. let conn = Connection::new_with_handler(|_request| { http::Response::builder() @@ -2197,7 +2385,7 @@ mod tests { let names = conn .table_names() - .namespace(vec!["ns".to_string()]) + .namespace(namespace) .execute() .await .unwrap(); @@ -2206,8 +2394,11 @@ mod tests { assert_eq!(names, vec!["a", "a"]); } + #[rstest::rstest] + #[case(vec![])] + #[case(vec!["ns".to_string()])] #[tokio::test] - async fn test_table_names_in_a_namespace_stops_on_an_empty_token() { + async fn test_table_names_stops_on_an_empty_token(#[case] namespace: Vec) { // An empty token ends the listing. Sending it back would ask a server that reads it // as "start from the beginning" for the first page a second time, and every name on // that page would be collected twice. @@ -2227,7 +2418,7 @@ mod tests { let names = conn .table_names() - .namespace(vec!["ns".to_string()]) + .namespace(namespace) .execute() .await .unwrap(); @@ -3892,7 +4083,7 @@ mod tests { .expect_err("an illegal component must be refused"); assert!(!*reached.lock().unwrap(), "{name:?} reached the transport"); assert!( - error.to_string().contains("Secret name"), + error.to_string().contains("secret name"), "{name:?}: {error}" ); } @@ -3923,7 +4114,7 @@ mod tests { .await .expect_err("a dot-only name must be refused"); assert!( - by_name.to_string().contains("relative path segments"), + by_name.to_string().contains("Illegal secret name"), "{by_name}" ); @@ -3933,7 +4124,9 @@ mod tests { .await .expect_err("a dot-only namespace segment must be refused"); assert!( - by_segment.to_string().contains("relative path segments"), + by_segment + .to_string() + .contains("Illegal namespace component"), "{by_segment}" ); assert!( @@ -4012,28 +4205,20 @@ mod tests { } } - /// A segment outside the table charset still addresses one segment: the - /// service decides whether it may exist, and percent-encoding is what keeps - /// the question reaching the right route. A catalog database is named this - /// way. + /// Slashes belong to catalog database names, not table namespace segments. #[tokio::test] - async fn test_a_namespace_segment_outside_the_charset_is_encoded_not_refused() { - use std::sync::{Arc, Mutex}; - let seen = Arc::new(Mutex::new(String::new())); - let path = seen.clone(); - let conn = Connection::new_with_handler(move |request| { - *path.lock().unwrap() = request.url().path().to_string(); - http::Response::builder().status(200).body("{}").unwrap() + async fn test_a_namespace_segment_outside_the_charset_is_refused() { + let conn = Connection::new_with_handler(|_| -> http::Response { + panic!("invalid namespace reached the transport") }); - conn.drop_table("t", &["team/search".to_string()]) - .await - .unwrap(); - assert_eq!(*seen.lock().unwrap(), "/v1/table/team%2Fsearch$t/drop/"); + assert!(matches!( + conn.drop_table("t", &["team/search".to_string()]).await, + Err(Error::InvalidInput { .. }) + )); } - /// A Function name is percent-encoded, which covers everything but the - /// relative segment: `..` is unreserved, so it survives encoding and is - /// then resolved away, posting a registration body to `/v1/create`. + /// Reject relative segments and their encoded forms before sending a request, + /// so a function name cannot change the route through URL normalization. #[tokio::test] async fn test_a_relative_segment_function_name_is_refused() { use std::sync::{Arc, Mutex}; @@ -4048,10 +4233,7 @@ mod tests { .drop_function(name, "fv_1", &[]) .await .expect_err("a dot-only Function name must be refused"); - assert!( - error.to_string().contains("relative path segments"), - "{error}" - ); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); assert!(!*reached.lock().unwrap(), "{name:?} reached the transport"); } } @@ -4432,7 +4614,7 @@ mod tests { }; // A view is not a table, so the refusal says so. assert!( - error.contains("view name") || error.contains("view namespace path segment"), + error.contains("view name") || error.contains("namespace component"), "{error}" ); } diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index e8ddcf6ff..a58a9afad 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -27,12 +27,17 @@ use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity use uuid::Uuid; use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; +use crate::database::ExecuteQueryRequest; use crate::error::{Error, Result}; use crate::job::Job; use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; +mod exchange; + +use exchange::{ExchangeQuery, ExchangeQueryHandle}; + const DEFAULT_SQL_PORT: u16 = 10025; const DEFAULT_SQL_TLS_PORT: u16 = 10026; const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(120); @@ -185,6 +190,23 @@ impl SqlClient { } } + pub(super) async fn execute(&self, request: ExecuteQueryRequest) -> Result { + match request.parameters { + Some(parameters) => { + self.submit_with_parameters( + &request.query, + parameters, + &request.default_namespace_path, + ) + .await + } + None => { + self.submit(&request.query, &request.default_namespace_path) + .await + } + } + } + pub(super) async fn submit( &self, query: &str, @@ -206,7 +228,8 @@ impl SqlClient { default_namespace_path.to_vec(), poll_info, )?); - self.queries.insert(query_id, query.clone()); + self.queries + .insert(query_id, RegisteredQuery::Polled(query.clone())); Ok(Query::new(Arc::new(RemoteQueryHandle::new(query)))) }) .await @@ -248,6 +271,35 @@ impl SqlClient { }) } + /// Values cannot ride on a poll, so a statement that has them is sent on + /// its own exchange instead. + pub(super) async fn submit_with_parameters( + &self, + query: &str, + parameters: RecordBatch, + default_namespace_path: &[String], + ) -> Result { + let timeout = self.inner.overall_timeout()?; + with_overall_timeout(timeout, "SQL query submission", async { + validate_namespace_path(default_namespace_path)?; + let query_id = Uuid::now_v7(); + let query = Arc::new( + ExchangeQuery::open( + query_id, + self.inner.clone(), + query, + parameters, + default_namespace_path, + ) + .await?, + ); + self.queries + .insert(query_id, RegisteredQuery::Exchanged(query.clone())); + Ok(Query::new(Arc::new(ExchangeQueryHandle::new(query)))) + }) + .await + } + pub(super) async fn describe(&self, query_id: Uuid) -> Result { let query = self .queries @@ -530,8 +582,40 @@ impl SqlClientInner { } } +/// A submitted query as the connection remembers it. +#[derive(Clone)] +enum RegisteredQuery { + Polled(Arc), + Exchanged(Arc), +} + +impl RegisteredQuery { + fn touch(&self) { + match self { + Self::Polled(query) => query.touch(), + Self::Exchanged(query) => query.touch(), + } + } + + /// Whether the registry may forget this query. It is abandoned when the + /// registry holds the only reference to it. + fn registry_expired(&self) -> bool { + match self { + Self::Polled(query) => query.registry_expired(Arc::strong_count(query) == 1), + Self::Exchanged(query) => query.registry_expired(Arc::strong_count(query) == 1), + } + } + + async fn describe(&self) -> Result { + match self { + Self::Polled(query) => query.describe().await, + Self::Exchanged(query) => query.describe(), + } + } +} + struct QueryRegistry { - queries: StdMutex>>, + queries: StdMutex>, } impl QueryRegistry { @@ -541,12 +625,12 @@ impl QueryRegistry { } } - fn insert(&self, id: Uuid, query: Arc) { + fn insert(&self, id: Uuid, query: RegisteredQuery) { self.remove_expired(); self.queries.lock().unwrap().insert(id, query); } - fn get(&self, id: Uuid) -> Option> { + fn get(&self, id: Uuid) -> Option { self.remove_expired(); let query = self.queries.lock().unwrap().get(&id).cloned(); if let Some(query) = &query { @@ -559,7 +643,7 @@ impl QueryRegistry { self.queries .lock() .unwrap() - .retain(|_, query| !query.registry_expired(Arc::strong_count(query) == 1)); + .retain(|_, query| !query.registry_expired()); } } diff --git a/rust/lancedb/src/remote/sql/exchange.rs b/rust/lancedb/src/remote/sql/exchange.rs new file mode 100644 index 000000000..523aaf5d7 --- /dev/null +++ b/rust/lancedb/src/remote/sql/exchange.rs @@ -0,0 +1,435 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! SQL statements sent with parameters. +//! +//! A statement without parameters is polled: `PollFlightInfo` starts it +//! detached from the connection and `DoGet` reads its endpoints. Neither call +//! carries Arrow data, so a statement's values would have to be spelled into +//! its text. A statement with parameters is sent on one `DoExchange` instead: +//! the command and the parameter schema, then one row, and the result rows +//! come back on the same call. The call is the statement's lifetime, so +//! cancelling it -- or dropping the handle before reading -- stops the +//! statement on the server. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex as StdMutex, OnceLock}; +use std::time::Instant; + +use arrow_array::RecordBatch; +use arrow_flight::FlightDescriptor; +use arrow_flight::encode::FlightDataEncoderBuilder; +use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; +use arrow_schema::SchemaRef; +use prost::Message; +use tokio::sync::{Notify, mpsc}; +use uuid::Uuid; + +use super::{ + ABANDONED_QUERY_RETENTION, DEFAULT_READ_TIMEOUT, ResultEndpointStream, ResultStartGuard, + SqlClientInner, TERMINAL_QUERY_RETENTION, resolve_timeout, sql_error, with_overall_timeout, +}; +use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; +use crate::error::{Error, Result}; +use crate::sql::{QueryDescription, QueryHandle, QueryStatus}; + +#[derive(Clone, Debug, PartialEq, Eq)] +enum Lifecycle { + Running, + /// Every row has been received. + Finished, + Cancelled, + Failed(String), +} + +/// A failed statement, keeping the message `describe` reports it with -- not +/// `Runtime error: …` again when `describe` wraps it in a runtime error. +fn failure(error: &Error) -> Lifecycle { + Lifecycle::Failed(match error { + Error::Runtime { message } => message.clone(), + other => other.to_string(), + }) +} + +/// A statement running on its own `DoExchange` call. +pub(super) struct ExchangeQuery { + id: Uuid, + client: Arc, + /// The call's rows, until a reader takes them or a cancellation drops them. + response: StdMutex>, + lifecycle: StdMutex, + cancelled: Notify, + terminal_at: OnceLock, + last_accessed: StdMutex, +} + +struct PreparedExchange { + schema: SchemaRef, + first: Option, + rest: Option, +} + +/// The call's rows while a reader waits for their first batch. +/// +/// A reader whose future is dropped mid-wait -- its caller gave up, or its +/// task was cancelled -- puts them back, so the next reader finds the query as +/// this one did. A query cancelled meanwhile has its rows dropped instead, +/// which ends the call. +struct Borrowed<'a> { + query: &'a ExchangeQuery, + rows: Option, +} + +impl Borrowed<'_> { + fn get(&mut self) -> &mut ResultEndpointStream { + self.rows.as_mut().expect("rows are borrowed until kept") + } + + fn keep(mut self) -> ResultEndpointStream { + self.rows.take().expect("rows are borrowed until kept") + } +} + +impl Drop for Borrowed<'_> { + fn drop(&mut self) { + let Some(rows) = self.rows.take() else { + return; + }; + // Checked under the response lock: `cancel` settles first and takes + // the response after, so rows put back here are never missed by it. + let mut response = self.query.response.lock().unwrap(); + if *self.query.lifecycle.lock().unwrap() == Lifecycle::Running { + *response = Some(rows); + } + } +} + +impl ExchangeQuery { + /// Send the statement and its parameters, returning once the server has + /// planned it: a statement it refuses is refused here, as a polled one is + /// on its first poll. + pub(super) async fn open( + id: Uuid, + client: Arc, + query: &str, + parameters: RecordBatch, + default_namespace_path: &[String], + ) -> Result { + if parameters.num_rows() != 1 { + return Err(Error::InvalidInput { + message: format!( + "SQL query parameters must be exactly one row, got {}", + parameters.num_rows() + ), + }); + } + let request_id = Uuid::new_v4().to_string(); + let read_timeout = resolve_timeout( + client.client_config.timeout_config.read_timeout, + "LANCE_CLIENT_READ_TIMEOUT", + Some(DEFAULT_READ_TIMEOUT), + )? + .unwrap(); + let descriptor = FlightDescriptor::new_cmd( + CommandStatementQuery { + query: query.to_string(), + transaction_id: None, + } + .as_any() + .encode_to_vec(), + ); + // The descriptor rides on the schema message, and the row follows it. + let request = FlightDataEncoderBuilder::new() + .with_flight_descriptor(Some(descriptor)) + .with_schema(parameters.schema()) + .build(futures::stream::iter([Ok(parameters)])); + let mut flight = client + .client_with_headers(default_namespace_path, &request_id) + .await?; + let stream = tokio::time::timeout(read_timeout, flight.do_exchange(request)) + .await + .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? + .map_err(|err| sql_error(&request_id, err))?; + Ok(Self { + id, + client, + response: StdMutex::new(Some(ResultEndpointStream { + stream, + request_id, + read_timeout, + })), + lifecycle: StdMutex::new(Lifecycle::Running), + cancelled: Notify::new(), + terminal_at: OnceLock::new(), + last_accessed: StdMutex::new(Instant::now()), + }) + } + + pub(super) fn touch(&self) { + *self.last_accessed.lock().unwrap() = Instant::now(); + } + + pub(super) fn registry_expired(&self, abandoned: bool) -> bool { + if let Some(finished) = self.terminal_at.get() { + return finished.elapsed() >= TERMINAL_QUERY_RETENTION; + } + abandoned && self.last_accessed.lock().unwrap().elapsed() >= ABANDONED_QUERY_RETENTION + } + + pub(super) fn describe(&self) -> Result { + self.touch(); + let status = match &*self.lifecycle.lock().unwrap() { + Lifecycle::Running => QueryStatus::Running, + Lifecycle::Finished => QueryStatus::Finished, + Lifecycle::Cancelled => QueryStatus::Cancelled, + Lifecycle::Failed(message) => { + return Err(Error::Runtime { + message: message.clone(), + }); + } + }; + Ok(QueryDescription { + id: self.id, + status, + progress: (status == QueryStatus::Finished).then_some(1.0), + expires_at: None, + }) + } + + /// Stop the statement. Idempotent, and a no-op once it has ended. + pub(super) fn cancel(&self) { + self.touch(); + let settled = self.settle(Lifecycle::Cancelled); + // Unread rows are dropped here, which cancels the call; a query that + // already ended has none. A reader's stream is dropped by its task + // when it sees the cancellation. + self.response.lock().unwrap().take(); + if settled { + self.cancelled.notify_waiters(); + } + } + + /// Move a running statement to `state`. `true` if this call did. + fn settle(&self, state: Lifecycle) -> bool { + let mut lifecycle = self.lifecycle.lock().unwrap(); + if *lifecycle != Lifecycle::Running { + return false; + } + *lifecycle = state; + drop(lifecycle); + let _ = self.terminal_at.set(Instant::now()); + true + } + + fn is_cancelled(&self) -> bool { + *self.lifecycle.lock().unwrap() == Lifecycle::Cancelled + } + + fn cancelled_error(&self) -> Error { + Error::JobCancelled { + job_id: Some(self.id.to_string()), + } + } + + async fn wait_for_cancellation(&self) { + loop { + let cancelled = self.cancelled.notified(); + if self.is_cancelled() { + return; + } + cancelled.await; + } + } + + /// Take the rows and wait for the first batch, which carries the schema. + async fn prepare(&self) -> Result { + let Some(rows) = self.response.lock().unwrap().take() else { + // Taken by an earlier attempt that failed, or dropped by a cancel. + return Err(match &*self.lifecycle.lock().unwrap() { + Lifecycle::Failed(message) => Error::Runtime { + message: message.clone(), + }, + _ => self.cancelled_error(), + }); + }; + let mut rows = Borrowed { + query: self, + rows: Some(rows), + }; + let first = tokio::select! { + biased; + _ = self.wait_for_cancellation() => return Err(self.cancelled_error()), + result = rows.get().next_batch() => result, + }; + // Read or failed, the rows are this attempt's now: a failed read is + // not worth putting back. + let rows = rows.keep(); + let first = first?; + let schema = first + .as_ref() + .map(RecordBatch::schema) + .or_else(|| rows.stream.schema().cloned()) + .ok_or_else(|| Error::Runtime { + message: "SQL result stream did not include a schema".to_string(), + })?; + if first.is_none() { + self.settle(Lifecycle::Finished); + } + Ok(PreparedExchange { + schema, + rest: first.is_some().then_some(rows), + first, + }) + } + + async fn run( + self: Arc, + mut prepared: PreparedExchange, + sender: mpsc::Sender>, + ) -> Result<()> { + if let Some(batch) = prepared.first.take() + && !self.send(&sender, batch).await? + { + return Ok(()); + } + let Some(mut rows) = prepared.rest.take() else { + return Ok(()); + }; + loop { + let batch = tokio::select! { + biased; + // The reader was dropped: the rows go with this task, which + // ends the call. + _ = sender.closed() => { + self.settle(Lifecycle::Cancelled); + return Ok(()); + } + _ = self.wait_for_cancellation() => return Err(self.cancelled_error()), + result = rows.next_batch() => result, + }; + match batch { + Ok(Some(batch)) => { + if !self.send(&sender, batch).await? { + return Ok(()); + } + } + Ok(None) => { + self.settle(Lifecycle::Finished); + return Ok(()); + } + Err(error) => { + self.settle(failure(&error)); + return Err(error); + } + } + } + } + + async fn send( + &self, + sender: &mpsc::Sender>, + batch: RecordBatch, + ) -> Result { + tokio::select! { + biased; + _ = self.wait_for_cancellation() => Err(self.cancelled_error()), + result = sender.send(Ok(batch)) => { + if result.is_err() { + self.settle(Lifecycle::Cancelled); + } + Ok(result.is_ok()) + } + } + } +} + +pub(super) struct ExchangeQueryHandle { + query: Arc, + result_started: AtomicBool, +} + +impl ExchangeQueryHandle { + pub(super) fn new(query: Arc) -> Self { + Self { + query, + result_started: AtomicBool::new(false), + } + } +} + +impl Drop for ExchangeQueryHandle { + // Rows nobody took are rows nobody will read. Keeping the call open would + // hold the statement running on the server, stalled on flow control, + // until the registry forgot it. + fn drop(&mut self) { + if self.query.response.lock().unwrap().is_some() { + self.query.cancel(); + } + } +} + +#[async_trait::async_trait] +impl QueryHandle for ExchangeQueryHandle { + fn id(&self) -> Uuid { + self.query.touch(); + self.query.id + } + + async fn describe(&self) -> Result { + self.query.describe() + } + + async fn reader(&self) -> Result { + let timeout = self.query.client.overall_timeout()?; + if self + .result_started + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_err() + { + return Err(Error::Runtime { + message: "SQL query results can only be consumed once".to_string(), + }); + } + let result_start = ResultStartGuard::new(&self.result_started); + self.query.touch(); + let started = Instant::now(); + let prepared = + match with_overall_timeout(timeout, "SQL query result", self.query.prepare()).await { + Ok(prepared) => prepared, + Err(error) => { + // A failed read or the timeout ends the query, and with it + // the call, including rows a timed-out wait put back. + self.query.settle(failure(&error)); + self.query.response.lock().unwrap().take(); + return Err(error); + } + }; + let remaining_timeout = timeout.map(|timeout| timeout.saturating_sub(started.elapsed())); + let schema = prepared.schema.clone(); + let (sender, receiver) = mpsc::channel(2); + let error_sender = sender.clone(); + let query = self.query.clone(); + tokio::spawn(async move { + let result = with_overall_timeout( + remaining_timeout, + "SQL query result", + query.clone().run(prepared, sender), + ) + .await; + if let Err(error) = result { + query.settle(failure(&error)); + let _ = error_sender.send(Err(error)).await; + } + }); + let stream = futures::stream::unfold(receiver, |mut receiver| async move { + receiver.recv().await.map(|item| (item, receiver)) + }); + result_start.commit(); + Ok(Box::pin(SimpleRecordBatchStream::new(stream, schema))) + } + + async fn cancel(&self) -> Result<()> { + self.query.cancel(); + Ok(()) + } +} diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 0d7efc9c4..79ca339db 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -4,7 +4,8 @@ use std::sync::atomic::AtomicUsize; use arrow_array::builder::StringDictionaryBuilder; -use arrow_array::{Array, Int64Array, StringArray, types::Int32Type}; +use arrow_array::types::{Float32Type, Int32Type}; +use arrow_array::{Array, ArrayRef, FixedSizeListArray, Float32Array, Int64Array, StringArray}; use arrow_flight::encode::FlightDataEncoderBuilder; use arrow_flight::flight_service_server::{FlightService, FlightServiceServer}; use arrow_flight::sql::{Any, CommandStatementQuery}; @@ -15,6 +16,7 @@ use arrow_flight::{ use arrow_schema::{DataType, Field, Schema}; use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; +use tonic::metadata::MetadataMap; use tonic::{Request, Response, Status, Streaming}; use super::*; @@ -63,6 +65,9 @@ struct CapturedHeaders { database_prefix: String, } +/// Every exchange's statement and the parameter row it arrived with. +type ReceivedExchanges = Arc)>>>; + #[derive(Clone)] struct TestSqlService { query_count: Arc, @@ -77,11 +82,46 @@ struct TestSqlService { transient_poll_failures: Arc, headers: Arc>>, gate_refresh_release: Arc, + exchanges: ReceivedExchanges, + /// How many exchange responses have been dropped, finished or not. + exchange_drops: Arc, result: RecordBatch, large_result: RecordBatch, dictionary_result: RecordBatch, } +/// Counts a response stream being dropped, which is how a server sees a call +/// end -- by the client cancelling it or by the rows running out. +struct DropSignal(Arc); + +impl Drop for DropSignal { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } +} + +impl TestSqlService { + fn capture_headers(&self, metadata: &MetadataMap) { + let header = |name| { + metadata + .get(name) + .and_then(|value| value.to_str().ok()) + .unwrap() + .to_string() + }; + self.headers.lock().unwrap().push(CapturedHeaders { + database: header("database"), + namespace_path: header("namespace-path"), + request_id: header("x-request-id"), + api_key: header("x-api-key"), + database_prefix: metadata + .get("x-lancedb-database-prefix") + .map(|value| value.to_str().unwrap().to_string()) + .unwrap_or_default(), + }); + } +} + impl Default for TestSqlService { fn default() -> Self { let schema = Arc::new(Schema::new(vec![Field::new( @@ -126,6 +166,8 @@ impl Default for TestSqlService { transient_poll_failures: Arc::new(AtomicUsize::new(0)), headers: Arc::new(std::sync::Mutex::new(Vec::new())), gate_refresh_release: Arc::new(Notify::new()), + exchanges: Arc::new(std::sync::Mutex::new(Vec::new())), + exchange_drops: Arc::new(AtomicUsize::new(0)), result, large_result, dictionary_result, @@ -168,24 +210,7 @@ impl FlightService for TestSqlService { &self, request: Request, ) -> std::result::Result, Status> { - let metadata = request.metadata(); - let header = |name| { - metadata - .get(name) - .and_then(|value| value.to_str().ok()) - .unwrap() - .to_string() - }; - self.headers.lock().unwrap().push(CapturedHeaders { - database: header("database"), - namespace_path: header("namespace-path"), - request_id: header("x-request-id"), - api_key: header("x-api-key"), - database_prefix: metadata - .get("x-lancedb-database-prefix") - .map(|value| value.to_str().unwrap().to_string()) - .unwrap_or_default(), - }); + self.capture_headers(request.metadata()); let command = Any::decode(request.get_ref().cmd.as_ref()) .ok() @@ -394,11 +419,336 @@ impl FlightService for TestSqlService { Err(Status::unimplemented("list_actions")) } + /// Answers with the parameter row it was sent, so a test can compare what + /// came back with what went out. async fn do_exchange( &self, - _request: Request>, + request: Request>, ) -> std::result::Result, Status> { - Err(Status::unimplemented("do_exchange")) + self.capture_headers(request.metadata()); + let mut messages = request.into_inner(); + let first = messages + .message() + .await? + .ok_or_else(|| Status::invalid_argument("empty exchange"))?; + let query = first + .flight_descriptor + .as_ref() + .and_then(|descriptor| Any::decode(descriptor.cmd.as_ref()).ok()) + .and_then(|any| any.unpack::().ok().flatten()) + .map(|command| command.query) + .ok_or_else(|| Status::invalid_argument("exchange without a statement"))?; + let data = futures::stream::once(async move { Ok(first) }) + .chain(messages) + .map_err(|status| FlightError::Tonic(Box::new(status))); + let parameters = FlightRecordBatchStream::new_from_flight_data(data) + .try_next() + .await + .map_err(|error| Status::invalid_argument(error.to_string()))?; + self.exchanges + .lock() + .unwrap() + .push((query.clone(), parameters.clone())); + if query == "SELECT exchange refused" { + return Err(Status::invalid_argument( + "parameter 1 is not referenced by the statement", + )); + } + if query == "SELECT exchange stall" { + return Ok(Response::new(futures::stream::pending().boxed())); + } + let rows = + parameters.ok_or_else(|| Status::invalid_argument("exchange without parameters"))?; + let schema = rows.schema(); + let rest: BoxStream<'static, arrow_flight::error::Result> = + if query == "SELECT exchange slow" { + futures::stream::pending().boxed() + } else { + futures::stream::empty().boxed() + }; + let delay = if query == "SELECT exchange delayed" { + Duration::from_millis(300) + } else { + Duration::ZERO + }; + let dropped = DropSignal(self.exchange_drops.clone()); + let stream = FlightDataEncoderBuilder::new() + .with_schema(schema) + .build( + futures::stream::once(async move { + tokio::time::sleep(delay).await; + Ok(rows) + }) + .chain(rest), + ) + .map_err(Status::from) + .map(move |message| { + let _dropped = &dropped; + message + }); + Ok(Response::new(Box::pin(stream))) + } +} + +fn parameter_row() -> RecordBatch { + let vector = FixedSizeListArray::from_iter_primitive::( + [Some([0.1_f32, 1e-45, -0.0].map(Some))], + 3, + ); + RecordBatch::try_from_iter([ + ("vector", Arc::new(vector) as ArrayRef), + ("k", Arc::new(Int64Array::from(vec![10])) as ArrayRef), + ]) + .unwrap() +} + +fn parameterized(query: &str, parameters: RecordBatch) -> ExecuteQueryRequest { + ExecuteQueryRequest { + parameters: Some(parameters), + ..ExecuteQueryRequest::new(query) + } +} + +async fn start_sql_service( + service: TestSqlService, + client_config: ClientConfig, +) -> (SqlClient, tokio::sync::oneshot::Sender<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let incoming = futures::stream::try_unfold(listener, |listener| async { + let (socket, _) = listener.accept().await?; + Ok::<_, std::io::Error>(Some((socket, listener))) + }); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + tokio::spawn( + tonic::transport::Server::builder() + .add_service(FlightServiceServer::new(service)) + .serve_with_incoming_shutdown(incoming, async { + let _ = shutdown_rx.await; + }), + ); + let client = SqlClient::new( + "analytics".to_string(), + None, + "test-key".to_string(), + None, + Some(format!("grpc://{address}")), + client_config, + ); + (client, shutdown_tx) +} + +async fn wait_for(count: &AtomicUsize, expected: usize, what: &str) { + tokio::time::timeout(Duration::from_secs(5), async { + while count.load(Ordering::SeqCst) < expected { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{what}")); +} + +/// The statement and its row go out on one exchange -- no poll, no ticket -- +/// and the row arrives as exactly the bytes that were sent. +#[tokio::test] +async fn parameters_travel_on_one_exchange() { + let service = TestSqlService::default(); + let exchanges = service.exchanges.clone(); + let query_count = service.query_count.clone(); + let do_get_count = service.do_get_count.clone(); + let headers = service.headers.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + + let parameters = parameter_row(); + let query = client + .execute(ExecuteQueryRequest { + default_namespace_path: vec!["events".to_string(), "raw".to_string()], + ..parameterized("SELECT echo", parameters.clone()) + }) + .await + .unwrap(); + assert_eq!(query.describe().await.unwrap().status, QueryStatus::Running); + let rows = collect_result(&query).await.unwrap(); + assert_eq!(rows, vec![parameters.clone()]); + let echoed = rows[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .map(|value| value.to_bits()) + .collect::>(); + assert_eq!( + echoed, + [0.1_f32, 1e-45, -0.0].map(f32::to_bits), + "every float keeps its exact bits, the subnormal and the sign of zero included" + ); + + let finished = query.describe().await.unwrap(); + assert_eq!(finished.status, QueryStatus::Finished); + assert_eq!(finished.progress, Some(1.0)); + assert_eq!( + client.describe(query.id()).await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + + assert_eq!( + *exchanges.lock().unwrap(), + vec![("SELECT echo".to_string(), Some(parameters))] + ); + assert_eq!(query_count.load(Ordering::SeqCst), 0); + assert_eq!(do_get_count.load(Ordering::SeqCst), 0); + let headers = headers.lock().unwrap(); + assert_eq!(headers.len(), 1); + assert_eq!(headers[0].database, "analytics"); + assert_eq!(headers[0].namespace_path, "events$raw"); + assert_eq!(headers[0].api_key, "test-key"); +} + +/// The server plans before it answers, so a refusal is the submission's. +#[tokio::test] +async fn a_refused_statement_fails_its_submission() { + let (client, _shutdown) = + start_sql_service(TestSqlService::default(), ClientConfig::default()).await; + let refused = client + .execute(parameterized("SELECT exchange refused", parameter_row())) + .await + .unwrap_err(); + assert!( + refused + .to_string() + .contains("not referenced by the statement"), + "{refused}" + ); + + let two_rows = + RecordBatch::try_from_iter([("k", Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef)]) + .unwrap(); + assert!(matches!( + client.execute(parameterized("SELECT echo", two_rows)).await, + Err(Error::InvalidInput { .. }) + )); +} + +/// The call is the statement's lifetime: cancelling a read, or dropping a +/// handle nobody read, ends the call and so the work behind it. +#[tokio::test] +async fn cancelling_or_abandoning_an_exchange_ends_the_call() { + let service = TestSqlService::default(); + let drops = service.exchange_drops.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + + let reading = client + .execute(parameterized("SELECT exchange slow", parameter_row())) + .await + .unwrap(); + let mut rows = reading.reader().await.unwrap(); + assert_eq!(rows.try_next().await.unwrap().unwrap(), parameter_row()); + reading.cancel().await.unwrap(); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), rows.try_next()) + .await + .expect("cancellation must wake the reader"), + Err(Error::JobCancelled { .. }) + )); + assert_eq!( + reading.describe().await.unwrap().status, + QueryStatus::Cancelled + ); + wait_for(&drops, 1, "cancelling must end the call on the server").await; + reading.cancel().await.unwrap(); + + let unread = client + .execute(parameterized("SELECT exchange slow", parameter_row())) + .await + .unwrap(); + let id = unread.id(); + drop(unread); + wait_for(&drops, 2, "dropping an unread handle must end the call").await; + assert_eq!( + client.describe(id).await.unwrap().status, + QueryStatus::Cancelled + ); + + let dropped_reader = client + .execute(parameterized("SELECT exchange slow", parameter_row())) + .await + .unwrap(); + let rows = dropped_reader.reader().await.unwrap(); + drop(rows); + wait_for(&drops, 3, "dropping a reader must end the call").await; +} + +/// A reader abandoned before the first batch -- its task cancelled -- must +/// leave the rows for the next reader, as a polled query does, rather than +/// take the call down with it. +#[tokio::test] +async fn an_abandoned_reader_leaves_the_rows_for_the_next_one() { + let service = TestSqlService::default(); + let drops = service.exchange_drops.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + + let delayed = Arc::new( + client + .execute(parameterized("SELECT exchange delayed", parameter_row())) + .await + .unwrap(), + ); + let abandoned = { + let delayed = delayed.clone(); + tokio::spawn(async move { delayed.reader().await.map(|_| ()) }) + }; + tokio::time::sleep(Duration::from_millis(50)).await; + abandoned.abort(); + assert!(abandoned.await.unwrap_err().is_cancelled()); + assert_eq!( + delayed.describe().await.unwrap().status, + QueryStatus::Running + ); + assert_eq!(drops.load(Ordering::SeqCst), 0, "the call must survive"); + + assert_eq!( + collect_result(&delayed).await.unwrap(), + vec![parameter_row()] + ); + assert_eq!( + delayed.describe().await.unwrap().status, + QueryStatus::Finished + ); +} + +/// A reader that gives up before the first batch has already taken the rows, +/// so the query must say it failed rather than report `Running` for a call +/// that is gone. +#[tokio::test] +async fn a_reader_that_times_out_reports_the_failure() { + let mut client_config = ClientConfig::default(); + client_config.timeout_config.timeout = Some(Duration::from_millis(200)); + let (client, _shutdown) = start_sql_service(TestSqlService::default(), client_config).await; + + let stalled = client + .execute(parameterized("SELECT exchange stall", parameter_row())) + .await + .unwrap(); + assert_overall_timeout(stalled.reader().await, "result"); + let Err(retried) = stalled.reader().await else { + panic!("the rows are gone, so a second reader cannot succeed"); + }; + for error in [ + stalled.describe().await.unwrap_err(), + client.describe(stalled.id()).await.unwrap_err(), + retried, + ] { + assert!( + matches!(&error, Error::Runtime { message } if message == "SQL query result timed out"), + "{error}" + ); } } @@ -430,7 +780,10 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { let database = RemoteDatabase::for_catalog("https://catalog.example", Some(name), &options).unwrap(); let query = database - .execute_query_async("SELECT 42", &[]) + .execute_query_async(ExecuteQueryRequest { + default_namespace_path: vec![], + ..ExecuteQueryRequest::new("SELECT 42") + }) .await .unwrap(); assert_eq!( @@ -813,7 +1166,11 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { .submit("SELECT slow", &["public".to_string()]) .await .unwrap(); - let tracked_dropped_reader = client.queries.get(dropped_reader.id()).unwrap(); + let Some(RegisteredQuery::Polled(tracked_dropped_reader)) = + client.queries.get(dropped_reader.id()) + else { + panic!("a submitted statement is registered as a polled query"); + }; let continuation_count_before = first_continuation_count.load(Ordering::SeqCst); let dropped_result_stream = dropped_reader.reader().await.unwrap(); tokio::time::timeout(Duration::from_millis(100), async { @@ -1107,8 +1464,11 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { ) .unwrap(), ); - registry.insert(id, query.clone()); - assert!(Arc::ptr_eq(®istry.get(id).unwrap(), &query)); + registry.insert(id, RegisteredQuery::Polled(query.clone())); + let Some(RegisteredQuery::Polled(found)) = registry.get(id) else { + panic!("the registry returns the query it was given"); + }; + assert!(Arc::ptr_eq(&found, &query)); } let expired_id = Uuid::now_v7(); @@ -1125,7 +1485,7 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { ) .unwrap(), ); - registry.insert(expired_id, expired_query); + registry.insert(expired_id, RegisteredQuery::Polled(expired_query)); assert!(registry.get(expired_id).is_none()); let stale_id = Uuid::now_v7(); @@ -1142,7 +1502,7 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { .unwrap(), ); *stale_query.last_accessed.lock().unwrap() = Instant::now() - ABANDONED_QUERY_RETENTION; - registry.insert(stale_id, stale_query.clone()); + registry.insert(stale_id, RegisteredQuery::Polled(stale_query.clone())); drop(stale_query); assert!(registry.get(stale_id).is_none()); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index e4653fdca..119577dc5 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -47,6 +47,7 @@ use crate::utils::background_cache::BackgroundCache; use crate::utils::{ MaxBatchLengthStream, TimeoutStream, public_fts_field_path_by_id, resolve_arrow_field_path, resolve_arrow_fts_field_path, supported_btree_data_type, supported_vector_data_type, + validate_fts_field, }; use crate::{DistanceType, Error}; use crate::{ @@ -76,7 +77,6 @@ use lance::dataset::{ColumnAlteration, NewColumnTransform, Version}; use lance_datafusion::exec::{OneShotExec, execute_plan}; use reqwest::{RequestBuilder, Response}; use serde::{Deserialize, Serialize}; -use serde_json::Number; use std::collections::{HashMap, HashSet}; use std::io::Cursor; use std::pin::Pin; @@ -597,6 +597,7 @@ impl RemoteTable { Index::BloomFilter(p) => ("BLOOM_FILTER", Some(to_json(p)?)), Index::RTree(p) => ("RTREE", Some(to_json(p)?)), Index::FTS(p) => { + validate_fts_field(&field)?; let mut params = to_json(p)?; if p.get_document_granularity().is_list_element() { params["document_granularity"] = "list_element".into(); @@ -620,9 +621,9 @@ impl RemoteTable { }); } } - _ => { + Index::IvfHnswPq(_) => { return Err(Error::NotSupported { - message: "Index type not supported".into(), + message: "IVF_HNSW_PQ is not supported for remote tables. Please use IVF_HNSW_SQ instead.".into(), }); } }; @@ -1085,7 +1086,7 @@ impl RemoteTable { if let Some(full_text_search) = ¶ms.full_text_search { if full_text_search.wand_factor.is_some() { return Err(Error::NotSupported { - message: "Wand factor is not yet supported in LanceDB Cloud".into(), + message: "Wand factor is not supported for remote tables.".into(), }); } @@ -1145,17 +1146,14 @@ impl RemoteTable { if let Some(approx_mode) = query.approx_mode { body["approx_mode"] = serde_json::json!(approx_mode); } - // In 0.23.1 we migrated from `nprobes` to `minimum_nprobes` and `maximum_nprobes`. - // Old client / new server: since minimum_nprobes is missing, fallback to nprobes - // New client / old server: old server will only see nprobes, make sure to set both - // nprobes and minimum_nprobes - // New client / new server: since minimum_nprobes is present, server can ignore nprobes - body["nprobes"] = query.minimum_nprobes.into(); - body["minimum_nprobes"] = query.minimum_nprobes.into(); + if let Some(nprobes) = query.nprobes { + body["nprobes"] = nprobes.into(); + } + if let Some(minimum_nprobes) = query.minimum_nprobes { + body["minimum_nprobes"] = minimum_nprobes.into(); + } if let Some(maximum_nprobes) = query.maximum_nprobes { body["maximum_nprobes"] = maximum_nprobes.into(); - } else { - body["maximum_nprobes"] = serde_json::Value::Number(Number::from_u128(0).unwrap()) } body["lower_bound"] = query.lower_bound.into(); body["upper_bound"] = query.upper_bound.into(); @@ -2825,6 +2823,16 @@ impl BaseTable for RemoteTable { self.fetch_blobs_impl(column, row_ids).await } + async fn fetch_blob_ranges( + &self, + _column: &str, + _requests: &[crate::blob::BlobRangeRequest], + ) -> Result { + Err(Error::NotSupported { + message: "fetch_blob_ranges is not supported for remote tables.".into(), + }) + } + async fn fetch_blob_files( &self, column: &str, @@ -3139,7 +3147,7 @@ impl BaseTable for RemoteTable { async fn set_unenforced_primary_key(&self, _columns: &[&str]) -> Result<()> { Err(Error::NotSupported { - message: "set_unenforced_primary_key is not supported on LanceDB cloud.".into(), + message: "set_unenforced_primary_key is not supported for remote tables.".into(), }) } @@ -3331,7 +3339,7 @@ impl BaseTable for RemoteTable { async fn optimize(&self, _action: OptimizeAction) -> Result { self.check_mutable().await?; Err(Error::NotSupported { - message: "optimize is not supported on LanceDB cloud.".into(), + message: "optimize is not supported for remote tables.".into(), }) } async fn add_columns( @@ -3931,22 +3939,14 @@ impl TryFrom for MergeInsertRequest { let when_matched_update_all_filt = match value.when_matched_update_all_filt { Some(MergeFilter::Sql(sql)) => Some(sql), - Some(MergeFilter::Expr(_)) => { - return Err(Error::NotSupported { - message: "DataFusion expressions are not supported on remote tables".into(), - }); - } + Some(MergeFilter::Expr(expr)) => Some(expr_to_sql_string(&expr)?), None => None, }; let when_not_matched_by_source_delete_filt = match value.when_not_matched_by_source_delete_filt { Some(MergeFilter::Sql(sql)) => Some(sql), - Some(MergeFilter::Expr(_)) => { - return Err(Error::NotSupported { - message: "DataFusion expressions are not supported on remote tables".into(), - }); - } + Some(MergeFilter::Expr(expr)) => Some(expr_to_sql_string(&expr)?), None => None, }; @@ -4819,6 +4819,69 @@ mod tests { } } + #[rstest] + #[case::sql(false, false)] + #[case::expr_update(true, false)] + #[case::expr_delete(false, true)] + #[case::expr_both(true, true)] + #[tokio::test] + async fn test_merge_insert_filter_expressions( + #[case] update_expr: bool, + #[case] delete_expr: bool, + ) { + use datafusion_expr::{col, lit}; + + let batch = record_batch!(("id", Int32, [0, 1]), ("v", Int32, [100, 110])).unwrap(); + let data: Box = Box::new(RecordBatchIterator::new( + [Ok(batch.clone())], + batch.schema(), + )); + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); + let params = request.url().query_pairs().collect::>(); + assert_eq!(params["on"], "id"); + assert_eq!(params["when_matched_update_all"], "true"); + assert_eq!(params["when_not_matched_insert_all"], "false"); + assert_eq!(params["when_not_matched_by_source_delete"], "true"); + assert_eq!( + params["when_matched_update_all_filt"], + if update_expr { + "(`target`.v < `source`.v)" + } else { + "target.v < source.v" + } + ); + assert_eq!( + params["when_not_matched_by_source_delete_filt"], + if delete_expr { "(id > 3)" } else { "id > 3" } + ); + http::Response::builder() + .status(200) + .body(r#"{"version": 2, "num_updated_rows": 2, "num_inserted_rows": 0, "num_deleted_rows": 2}"#) + .unwrap() + }); + + let mut merge = table.merge_insert(&["id"]); + if update_expr { + merge.when_matched_update_all_expr(col("target.v").lt(col("source.v"))); + } else { + merge.when_matched_update_all(Some("target.v < source.v".into())); + } + if delete_expr { + merge.when_not_matched_by_source_delete_expr(col("id").gt(lit(3))); + } else { + merge.when_not_matched_by_source_delete(Some("id > 3".into())); + } + + let result = merge.execute(data).await.unwrap(); + assert_eq!(result.version, 2); + assert_eq!(result.num_updated_rows, 2); + assert_eq!(result.num_inserted_rows, 0); + assert_eq!(result.num_deleted_rows, 2); + } + #[rstest] #[case::no_batches_insert(false, false)] #[case::empty_batch_insert(true, false)] @@ -5693,6 +5756,72 @@ mod tests { } } + #[tokio::test] + async fn test_remote_unsupported_operations() { + let table = Table::new_with_handler("my_table", |_| -> http::Response { + panic!("unsupported operations must not send requests to the server"); + }); + let mut fts = FullTextSearchQuery::new("test".into()); + fts.wand_factor = Some(1.0); + let results = [ + ( + "set_unenforced_primary_key", + table.set_unenforced_primary_key(["id"]).await, + ), + ( + "optimize", + table.optimize(OptimizeAction::All).await.map(|_| ()), + ), + ( + "fetch_blob_ranges", + table.fetch_blob_ranges("image", []).await.map(|_| ()), + ), + ( + "Wand factor", + table + .query() + .full_text_search(fts) + .execute() + .await + .map(|_| ()), + ), + ]; + for (operation, result) in results { + match result.unwrap_err() { + Error::NotSupported { message } => assert_eq!( + message, + format!("{operation} is not supported for remote tables.") + ), + error => panic!("expected not-supported error, got {error}"), + } + } + } + + #[tokio::test] + async fn test_remote_unsupported_hnsw_pq() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/describe/"); + let schema = Schema::new(vec![Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + false, + )]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + }); + let error = table + .create_index(&["vector"], Index::IvfHnswPq(Default::default())) + .execute() + .await + .unwrap_err(); + assert_not_supported_error( + error, + "IVF_HNSW_PQ is not supported for remote tables. Please use IVF_HNSW_SQ instead.", + ); + } + #[tokio::test] async fn test_fetch_blobs_rejects_missing_column() { let batch = RecordBatch::try_new( @@ -5806,12 +5935,12 @@ mod tests { assert_not_supported_error( table.fetch_blobs("image", &[1]).await.unwrap_err(), - "fetch_blobs", + "fetch_blobs is not supported by this LanceDB server.", ); assert_not_supported_error( table.fetch_blob_files("image", &[1]).await.unwrap_err(), - "requires LanceDB Cloud server 0.5.0 or newer", + "requires LanceDB server 0.5.0 or newer", ); } @@ -5957,9 +6086,6 @@ mod tests { let body: serde_json::Value = serde_json::from_slice(body).unwrap(); let mut expected_body = serde_json::json!({ "prefilter": true, - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, "lower_bound": Option::::None, "upper_bound": Option::::None, "k": 10, @@ -5991,7 +6117,7 @@ mod tests { } #[tokio::test] - async fn test_query_vector_approx_mode_sent_when_set() { + async fn test_query_vector_minimum_only_and_approx_mode() { let expected_data = RecordBatch::try_new( Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], @@ -6011,9 +6137,7 @@ mod tests { let body: serde_json::Value = serde_json::from_slice(body).unwrap(); let mut expected_body = serde_json::json!({ "prefilter": true, - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, + "minimum_nprobes": 5, "approx_mode": "accurate", "lower_bound": Option::::None, "upper_bound": Option::::None, @@ -6037,6 +6161,8 @@ mod tests { .query() .nearest_to(vec![0.1, 0.2, 0.3]) .unwrap() + .minimum_nprobes(5) + .unwrap() .approx_mode(crate::ApproxMode::Accurate) .execute() .await; @@ -6045,6 +6171,39 @@ mod tests { assert_eq!(data[0].as_ref().unwrap(), &expected_data); } + #[tokio::test] + async fn test_query_vector_explicitly_unbounded_maximum() { + let table = Table::new_with_handler("my_table", |request| { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["nprobes"], 5); + assert_eq!(body["maximum_nprobes"], 0); + + let data = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) + .body(write_ipc_file(&data)) + .unwrap() + }); + + table + .query() + .nearest_to(vec![0.1, 0.2, 0.3]) + .unwrap() + .nprobes(5) + .unwrap() + .maximum_nprobes(None) + .unwrap() + .execute() + .await + .unwrap(); + } + #[tokio::test] async fn test_query_fts_default_values() { let expected_data = RecordBatch::try_new( @@ -6143,8 +6302,8 @@ mod tests { } ], "nprobes": 12, - "minimum_nprobes": 12, - "maximum_nprobes": 12, + "minimum_nprobes": 3, + "maximum_nprobes": 10, "lower_bound": Option::::None, "upper_bound": Option::::None, "ef": Option::::None, @@ -6184,6 +6343,10 @@ mod tests { .distance_type(crate::DistanceType::Cosine) .nprobes(12) .unwrap() + .minimum_nprobes(3) + .unwrap() + .maximum_nprobes(Some(10)) + .unwrap() .refine_factor(2) .bypass_vector_index() .execute() @@ -6214,9 +6377,6 @@ mod tests { "vector_column": "image.embedding", "prefilter": true, "k": 10, - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, "lower_bound": Option::::None, "upper_bound": Option::::None, "ef": Option::::None, @@ -6952,19 +7112,24 @@ mod tests { ]; for (index_type, expected_body, index) in cases { + let data_type = match &index { + Index::FTS(params) if params.get_document_granularity().is_list_element() => { + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))) + } + Index::FTS(_) => DataType::Utf8, + _ => DataType::Int32, + }; + let schema = Schema::new(vec![Field::new("a", data_type, false)]); let table = Table::new_with_handler_version( "my_table", semver::Version::new(0, 6, 0), move |request| { assert_eq!(request.method(), "POST"); match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), "/v1/table/my_table/create_index/" => { assert_eq!( request.headers().get("Content-Type").unwrap(), diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index ac37bf7ae..b8506adb3 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -393,7 +393,7 @@ impl RemoteTable { } if !self.server_version.support_blobs() { return Err(Error::NotSupported { - message: "fetch_blobs is not supported on this LanceDB Cloud server".into(), + message: "fetch_blobs is not supported by this LanceDB server.".into(), }); } let mut read_snapshot = self.snapshot_read_state().await; @@ -716,7 +716,7 @@ impl RemoteTable { Ok(()) } else { Err(Error::NotSupported { - message: "fetch_blob_files requires LanceDB Cloud server 0.5.0 or newer".into(), + message: "fetch_blob_files requires LanceDB server 0.5.0 or newer.".into(), }) } } diff --git a/rust/lancedb/src/sql.rs b/rust/lancedb/src/sql.rs index 7c040c51b..c2ef13ba5 100644 --- a/rust/lancedb/src/sql.rs +++ b/rust/lancedb/src/sql.rs @@ -17,6 +17,10 @@ pub enum QueryStatus { /// The server is still executing the query. Running, /// The server has made the complete result available. + /// + /// A query submitted with parameters streams its rows on the call that + /// runs it, so it reports `Finished` only once its reader has received + /// the last batch. Read it rather than waiting for this status. Finished, /// The server accepted cancellation but has not confirmed it yet. Cancelling, @@ -60,7 +64,9 @@ pub(crate) trait QueryHandle: Send + Sync { /// A handle to a submitted SQL query. /// /// The handle can be inspected, opened as an Arrow reader, or cancelled. -/// Dropping it does not cancel the server-side query. +/// Dropping it does not cancel the server-side query, unless the query was +/// submitted with parameters: those run on the call that returns their rows, +/// so dropping the handle before reading, or dropping the reader, stops them. /// Identifier lookup is scoped to the connection that submitted the query and /// is not a durable resume mechanism. pub struct Query { @@ -88,6 +94,11 @@ impl Query { } /// Get a point-in-time description of the query. + /// + /// For a query submitted with parameters the status follows its result + /// stream: it stays [`QueryStatus::Running`] until the reader has received + /// every row, so open [`Self::reader`] directly instead of polling for + /// [`QueryStatus::Finished`] first. pub async fn describe(&self) -> Result { self.handle.describe().await } diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 916b28ed1..3ac759bdd 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -29,8 +29,8 @@ use crate::index::Index; use crate::index::vector::{VectorIndex, suggested_num_sub_vectors}; use crate::utils::{ resolve_lance_fts_field_path, supported_bitmap_data_type, supported_btree_data_type, - supported_fm_data_type, supported_fts_data_type, supported_label_list_data_type, - supported_vector_data_type, supported_zonemap_data_type, + supported_fm_data_type, supported_label_list_data_type, supported_vector_data_type, + supported_zonemap_data_type, validate_fts_field, }; use super::NativeTable; @@ -292,7 +292,7 @@ impl NativeTable { )) } Index::FTS(fts_opts) => { - Self::validate_index_type(field, "FTS", supported_fts_data_type)?; + validate_fts_field(field)?; Ok(Box::new(fts_opts)) } Index::IvfFlat(index) => { diff --git a/rust/lancedb/src/table/datafusion.rs b/rust/lancedb/src/table/datafusion.rs index ea5222b15..26c423497 100644 --- a/rust/lancedb/src/table/datafusion.rs +++ b/rust/lancedb/src/table/datafusion.rs @@ -699,6 +699,44 @@ pub mod tests { ); } + #[tokio::test] + async fn test_topk_late_materialization_over_table_adapter() { + use lance::io::exec::topk_late_materialization::TopKLateMaterialization; + + let fixture = TestFixture::new().await; + let run = |rule: Option| { + let adapter = fixture.adapter2.clone(); + async move { + let mut state = + datafusion::execution::SessionStateBuilder::new().with_default_features(); + if let Some(rule) = rule { + state = state.with_physical_optimizer_rule(Arc::new(rule)); + } + let ctx = SessionContext::new_with_state(state.build()); + ctx.register_table("tbl2", adapter).unwrap(); + let df = ctx + .sql("SELECT ints, strings, jsons FROM tbl2 WHERE floats > 10 ORDER BY ints DESC LIMIT 5") + .await + .unwrap(); + let plan = df.clone().create_physical_plan().await.unwrap(); + let plan = datafusion_physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + (plan, df.collect().await.unwrap()) + } + }; + + let (plan, expected) = run(None).await; + assert!(!plan.contains("Take"), "{plan}"); + + // The sort lands below `MetadataEraserExec`, so the rule needs no + // certification for it. + let (plan, batches) = run(Some(TopKLateMaterialization::new())).await; + assert!(plan.contains("Take"), "{plan}"); + assert_eq!(batches, expected); + assert!(batches[0].schema().metadata().is_empty()); + } + #[tokio::test] async fn test_metadata_eraser_with_new_children_takes_child_schema() { let fixture = TestFixture::new().await; diff --git a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs index 000addbe1..543a59589 100644 --- a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs +++ b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs @@ -411,9 +411,7 @@ impl ExecutionPlan for VectorDedupExec { .map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?; let mut scanner = dataset.scan(); scanner - .with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter( - removed, - ))) + .with_row_id_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter(removed))) .blob_handling(BlobHandling::AllBinary) .batch_size(MATERIALIZE_BATCH_ROWS) .strict_batch_size(true); diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 354eef5a5..1ad084e22 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -1296,6 +1296,21 @@ mod lsm_tests { rows } + async fn collect_ids(stream: SendableRecordBatchStream) -> Vec { + let batches: Vec<_> = stream.try_collect().await.unwrap(); + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect() + } + /// Upsert `ids` (value = 0..n) through the LSM `merge_insert` path. async fn lsm_upsert(table: &Table, ids: Vec) { let mut builder = table.merge_insert(&[]); @@ -1772,7 +1787,6 @@ mod lsm_tests { use crate::index::Index; use crate::index::vector::IvfPqIndexBuilder; use arrow::array::{FixedSizeListBuilder, Float32Builder}; - use arrow::datatypes::Int64Type; const DIM: i32 = 8; const N: i64 = 256; @@ -1817,7 +1831,7 @@ mod lsm_tests { &["vec"], Index::IvfPq( IvfPqIndexBuilder::default() - .num_partitions(1) + .num_partitions(8) .num_sub_vectors(2), ), ) @@ -1851,21 +1865,32 @@ mod lsm_tests { .execute() .await .unwrap(); - let batches: Vec<_> = stream.try_collect().await.unwrap(); - let ids: Vec = batches - .iter() - .flat_map(|b| { - b.column_by_name("id") - .unwrap() - .as_primitive::() - .values() - .to_vec() - }) - .collect(); + let ids = collect_ids(stream).await; assert_eq!( ids, vec![9999], "LSM vector search must rank the memtable row first" ); + + let base_query = table + .query() + .nearest_to(&[0.0_f32; 8]) + .unwrap() + .only_if("id = 255") + .limit(1); + for query in [ + base_query.clone(), + base_query.clone().nprobes(8).unwrap(), + base_query.clone().minimum_nprobes(8).unwrap(), + base_query.clone().maximum_nprobes(Some(8)).unwrap(), + base_query + .nprobes(1) + .unwrap() + .maximum_nprobes(None) + .unwrap(), + ] { + let ids = collect_ids(query.execute().await.unwrap()).await; + assert_eq!(ids, vec![255]); + } } } diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index d611fe8a5..77c4ca1a5 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -106,19 +106,16 @@ async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> R } fn requires_local_namespace_execution(query: &AnyQuery) -> bool { - // The namespace QueryTable request has no approx_mode or use_lsm field yet, so - // pushing these down would silently ignore the user's setting. For use_lsm that - // is worse than a tuning miss: MemWAL read routing lives only in `create_plan`, - // so a pushed-down query would return stale base-only data with no error. + // The namespace QueryTable request cannot represent approx_mode or use_lsm. + // For use_lsm that is worse than a tuning miss: MemWAL read routing lives only + // in `create_plan`, so a pushed-down query would return stale base-only data. if query.base().use_lsm.is_some() || query.base().take_offsets.is_some() { return true; } matches!( query, - AnyQuery::VectorQuery(VectorQueryRequest { - approx_mode: Some(_), - .. - }) + AnyQuery::VectorQuery(vector_query) + if vector_query.approx_mode.is_some() ) } @@ -296,8 +293,17 @@ pub async fn create_plan( scanner.approx_mode(approx_mode.into()); } - scanner.minimum_nprobes(query.minimum_nprobes); - if let Some(maximum_nprobes) = query.maximum_nprobes { + if let Some(nprobes) = query.nprobes { + if query.maximum_nprobes == Some(0) { + scanner.minimum_nprobes(nprobes); + } else { + scanner.nprobes(nprobes); + } + } + if let Some(minimum_nprobes) = query.minimum_nprobes { + scanner.minimum_nprobes(minimum_nprobes); + } + if let Some(maximum_nprobes) = query.maximum_nprobes.filter(|nprobes| *nprobes != 0) { scanner.maximum_nprobes(maximum_nprobes); } } @@ -630,7 +636,9 @@ fn convert_to_namespace_query(query: &AnyQuery) -> Result { columns, offset: vq.base.offset.map(|o| o as i32), distance_type: vq.distance_type.map(|dt| dt.to_string()), - nprobes: Some(vq.minimum_nprobes as i32), + nprobes: vq.nprobes.map(|nprobes| nprobes as i32), + minimum_nprobes: vq.minimum_nprobes.map(|nprobes| nprobes as i32), + maximum_nprobes: vq.maximum_nprobes.map(|nprobes| nprobes as i32), ef: vq.ef.map(|e| e as i32), refine_factor: vq.refine_factor.map(|r| r as i32), lower_bound: vq.lower_bound, @@ -919,7 +927,6 @@ mod tests { column: Some("vector".to_string()), // We cast here to satisfy the struct definition query_vector: vec![query_vector as Arc], - minimum_nprobes: 20, distance_type: Some(crate::DistanceType::L2), ..Default::default() }; @@ -940,6 +947,7 @@ mod tests { ); assert_eq!(ns_request.vector_column, Some("vector".to_string())); assert_eq!(ns_request.distance_type, Some("l2".to_string())); + assert_eq!(ns_request.nprobes, None); // Verify the vector data was extracted correctly assert!(ns_request.vector.single_vector.is_some()); @@ -949,6 +957,66 @@ mod tests { ); } + #[test] + fn test_namespace_query_probe_pushdown_compatibility() { + let default_query = AnyQuery::VectorQuery(VectorQueryRequest::default()); + assert!(!requires_local_namespace_execution(&default_query)); + + let nprobes_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + nprobes: Some(20), + ..Default::default() + }); + assert!(!requires_local_namespace_execution(&nprobes_query)); + assert_eq!( + convert_to_namespace_query(&nprobes_query).unwrap().nprobes, + Some(20) + ); + + let minimum_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + minimum_nprobes: Some(5), + ..Default::default() + }); + assert!(!requires_local_namespace_execution(&minimum_query)); + let request = convert_to_namespace_query(&minimum_query).unwrap(); + assert_eq!(request.minimum_nprobes, Some(5)); + assert_eq!(request.maximum_nprobes, None); + + let maximum_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + maximum_nprobes: Some(10), + ..Default::default() + }); + assert!(!requires_local_namespace_execution(&maximum_query)); + let request = convert_to_namespace_query(&maximum_query).unwrap(); + assert_eq!(request.minimum_nprobes, None); + assert_eq!(request.maximum_nprobes, Some(10)); + + let exact_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + nprobes: Some(8), + minimum_nprobes: Some(20), + maximum_nprobes: Some(20), + ..Default::default() + }); + assert!(!requires_local_namespace_execution(&exact_query)); + let request = convert_to_namespace_query(&exact_query).unwrap(); + assert_eq!(request.nprobes, Some(8)); + assert_eq!(request.minimum_nprobes, Some(20)); + assert_eq!(request.maximum_nprobes, Some(20)); + + let explicitly_unbounded_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + nprobes: Some(8), + maximum_nprobes: Some(0), + ..Default::default() + }); + let request = convert_to_namespace_query(&explicitly_unbounded_query).unwrap(); + assert_eq!(request.nprobes, Some(8)); + assert_eq!(request.maximum_nprobes, Some(0)); + } + #[test] fn test_convert_to_namespace_query_plain_query() { let q = QueryRequest { diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs index f3e40fc7b..5ec053483 100644 --- a/rust/lancedb/src/table/query/lsm.rs +++ b/rust/lancedb/src/table/query/lsm.rs @@ -170,9 +170,8 @@ fn reject_unsupported(query: &VectorQueryRequest) -> Result<()> { // Vector-only knobs the LSM scanner cannot honor. Both change results rather // than just recall, so error instead of silently ignoring them: distance_range // would return rows outside the bound, and use_index(false) asks for a - // brute-force search the index-only base arm can't do. (ef / approx_mode / - // maximum_nprobes are recall/speed knobs and are left to no-op — and - // maximum_nprobes defaults to Some, so it cannot be rejected on presence.) + // brute-force search the index-only base arm can't do. (ef / approx_mode are + // recall/speed knobs and are left to no-op.) if !query.query_vector.is_empty() { if query.lower_bound.is_some() || query.upper_bound.is_some() { return unsupported("distance_range on vector search"); @@ -740,8 +739,20 @@ async fn vector_plan( let mut scanner = base_scanner(dataset, query, pk_columns, snapshots, in_memory)? .with_overfetch_factor(LSM_OVERFETCH_FACTOR) .nearest(&column, query_vector.as_ref(), k)? - .nprobes(query.minimum_nprobes) .distance_metric(distance_type.into()); + if let Some(nprobes) = query.nprobes { + scanner = if query.maximum_nprobes == Some(0) { + scanner.minimum_nprobes(nprobes) + } else { + scanner.nprobes(nprobes) + }; + } + if let Some(minimum_nprobes) = query.minimum_nprobes { + scanner = scanner.minimum_nprobes(minimum_nprobes); + } + if let Some(maximum_nprobes) = query.maximum_nprobes.filter(|nprobes| *nprobes != 0) { + scanner = scanner.maximum_nprobes(maximum_nprobes); + } if let Some(refine_factor) = query.refine_factor { scanner = scanner.refine(refine_factor); } diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs index 52b7bd5e8..20f209f2e 100644 --- a/rust/lancedb/src/utils/mod.rs +++ b/rust/lancedb/src/utils/mod.rs @@ -19,15 +19,27 @@ use std::pin::Pin; use crate::error::{Error, Result}; use datafusion_physical_plan::SendableRecordBatchStream; -/// The characters any object name may contain: a table, a namespace segment, a -/// Secret, a materialized view. -/// -/// No positional rule on top of it -- a name may begin with `_`, `-` or `.`, -/// as LanceDB namespaces already do. `.` and `..` are excluded separately, by -/// [`reject_relative_segment`]: that is a property of where a name sits in a -/// URL, not of the name. Length is the service's to bound. -static OBJECT_NAME_REGEX: std::sync::LazyLock = - std::sync::LazyLock::new(|| regex::Regex::new(r"^[A-Za-z0-9_.\-]+$").unwrap()); +/// The maximum length of a table name. Table names are used in URLs and in the keys used with the +/// object store. +const MAX_TABLE_NAME_CHARS: usize = 128; + +/// The maximum length of a component of a namespace. +const MAX_NAMESPACE_NAME_CHARS: usize = 64; + +/// The maximum number of namespace components that are allowed. +const MAX_NAMESPACE_COMPONENTS: usize = 10; + +/// The maximum length of a view name. +const MAX_VIEW_NAME_CHARS: usize = 128; + +/// The maximum length of a secret name. +const MAX_SECRET_NAME_CHARS: usize = 128; + +/// The maximum length of a function name. +const MAX_FUNCTION_NAME_CHARS: usize = 128; + +/// The maximum length of a database name. +const MAX_DB_NAME_CHARS: usize = 64; pub trait PatchStoreParam { fn patch_with_store_wrapper( @@ -86,129 +98,190 @@ impl PatchReadParam for ReadParams { } } -/// The reason `.` and `..` are refused wherever a name becomes a path segment. -const RELATIVE_SEGMENT_REASON: &str = - "'.' and '..' are read as relative path segments and cannot address an object"; - -/// Whether URL parsing would resolve this component away rather than keep it. -/// -/// Exactly `.` and `..`, and their percent-encoded spellings -- resolution -/// happens after decoding, so `%2E%2E` collapses as surely as `..` does, and -/// `drop_table("..")` would reach `/v1/drop/`. No wider than that: `...` is an -/// ordinary segment that addresses fine. -fn is_relative_segment(value: &str) -> bool { - let decoded = value.replace("%2e", ".").replace("%2E", "."); - decoded == "." || decoded == ".." +/// Translate a string into something that can be safely logged. This involves removing control +/// characters that could be used in an exploit like CVE-2003-0063. +fn sanitize(input: &str) -> String { + input.chars().filter(|c| !c.is_control()).collect() } -/// Refuse a path component that URL parsing resolves as a relative segment. -/// -/// Reachable on its own for an identifier with no other validator: a Function -/// name has no client-side grammar, so this is the only rule that applies. -pub(crate) fn reject_relative_segment(what: &str, value: &str) -> Result<()> { - if is_relative_segment(value) { - return Err(Error::InvalidInput { - message: format!("invalid {what} '{value}': {RELATIVE_SEGMENT_REASON}"), +pub fn validate_table_name(name: &str) -> Result<()> { + if name.is_empty() || name == "." || name == ".." { + return Err(Error::InvalidTableName { + name: name.to_string(), + reason: "Illegal table name.".to_string(), + }); + } + if name.len() > MAX_TABLE_NAME_CHARS { + return Err(Error::InvalidTableName { + name: sanitize(name), + reason: format!( + "Table name cannot be more than {MAX_TABLE_NAME_CHARS} characters long." + ), + }); + } + if !name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.') + { + return Err(Error::InvalidTableName { + name: sanitize(name), + reason: "Only alphanumeric characters, dashes, underscores, and\ + periods are allowed in table names." + .to_string(), }); } Ok(()) } -/// Every rule an object name obeys: non-empty, inside [`OBJECT_NAME_REGEX`], -/// and addressable as a path segment. -/// -/// Returns the reason rather than an [`Error`], because the error type is each -/// API's own -- a table reports [`Error::InvalidTableName`], the rest -/// [`Error::InvalidInput`]. Sharing the rules but not the error keeps a table, -/// a namespace segment and a Secret from drifting apart. -fn check_object_name(name: &str) -> std::result::Result<(), &'static str> { - if name.is_empty() { - return Err("it must not be empty"); +/// Validate that an input string contains alphanumerics, or dashes, underscores, and periods. +fn validate_view_or_namespace_or_secret(input: &str, kind: &str, max_len: usize) -> Result<()> { + if input.is_empty() || input == "." || input == ".." { + return Err(Error::InvalidInput { + message: format!("Illegal {} '{}'.", kind, input), + }); } - if !OBJECT_NAME_REGEX.is_match(name) { - return Err( - "it may contain only alphanumeric characters, underscores, hyphens and periods", - ); + if input.len() > max_len { + return Err(Error::InvalidInput { + message: format!( + "Invalid {} '{}' is more than {} characters long.", + kind, + sanitize(input), + max_len + ), + }); } - if is_relative_segment(name) { - return Err(RELATIVE_SEGMENT_REASON); + if !input + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.') + { + return Err(Error::InvalidInput { + message: format!( + "Invalid {} '{}': only alphanumeric characters, \ + dashes, underscores, and periods are allowed.", + kind, + sanitize(input) + ), + }); } Ok(()) } -/// Validate a table name. -pub fn validate_table_name(name: &str) -> Result<()> { - check_object_name(name).map_err(|reason| Error::InvalidTableName { - name: name.to_string(), - reason: reason.to_string(), - }) -} - /// Validate one component of a namespace path -- a single segment, not the /// whole path. [`validate_namespace`] covers a path. -pub fn validate_namespace_name(name: &str) -> Result<()> { - check_object_name(name).map_err(|reason| Error::InvalidInput { - message: format!("invalid namespace name '{name}': {reason}"), - }) +pub fn validate_namespace_component(name: &str) -> Result<()> { + validate_view_or_namespace_or_secret(name, "namespace component", MAX_NAMESPACE_NAME_CHARS) } -/// Validate one component of a Secret identifier: a Secret name, or one segment -/// of the namespace path holding it. -/// -/// The join decides identity, and the service only sees what the split -/// produced. `"a$b"` is not a name the service accepts, but joined and split it -/// reads as the namespace `a` and the name `b` -- a different Secret that may -/// already exist. This is not a second opinion on the name; it is what lets the -/// service have one. -pub fn validate_secret_component(what: &str, value: &str) -> Result<()> { - check_object_name(value).map_err(|reason| Error::InvalidInput { - message: format!("invalid {what} '{value}': {reason}"), - }) +/// Validate all components of a namespace +pub fn validate_namespace(namespace: &[String]) -> Result<()> { + if namespace.len() > 1 && namespace.first().map(String::as_str) == Some("public") { + return Err(Error::InvalidInput { + message: "Namespaces that begin with public cannot have multiple namespace components." + .to_string(), + }); + } + if namespace.len() > MAX_NAMESPACE_COMPONENTS { + return Err(Error::InvalidInput { + message: format!( + "Cannot have a namespace with more than {MAX_NAMESPACE_COMPONENTS} components." + ), + }); + } + for component in namespace { + validate_namespace_component(component)?; + } + Ok(()) +} + +/// Validate a secret name. +pub fn validate_secret_name(name: &str) -> Result<()> { + validate_view_or_namespace_or_secret(name, "secret name", MAX_SECRET_NAME_CHARS) } /// Validate a Secret name and every segment of the namespace path holding it. pub fn validate_secret_reference(name: &str, namespace_path: &[String]) -> Result<()> { - for segment in namespace_path { - validate_secret_component("Secret namespace path segment", segment)?; - } - validate_secret_component("Secret name", name) + validate_namespace(namespace_path)?; + validate_secret_name(name) +} + +pub fn validate_view_name(name: &str) -> Result<()> { + validate_view_or_namespace_or_secret(name, "view name", MAX_VIEW_NAME_CHARS) } /// Validate a view name and every segment of the namespace path holding it. -/// -/// Worded for a view rather than deferring to [`validate_table_name`]: a view -/// is not a table, and a caller who mistypes one should not be told their -/// table name is invalid. The rule is the same one every object name follows, -/// which is what lets the `$`-joined identifier split back apart. pub fn validate_view_reference(name: &str, namespace_path: &[String]) -> Result<()> { - for segment in namespace_path { - validate_view_component("view namespace path segment", segment)?; + validate_namespace(namespace_path)?; + validate_view_name(name) +} + +/// Validate a function name. These do not allow dashes or periods. +pub fn validate_function_name(name: &str) -> Result<()> { + if name.is_empty() { + return Err(Error::InvalidInput { + message: format!("Invalid function name '{}'.", name), + }); } - validate_view_component("view name", name) + if name.len() > MAX_FUNCTION_NAME_CHARS { + return Err(Error::InvalidInput { + message: format!( + "Invalid function name '{}' is more than {} characters long.", + sanitize(name), + MAX_FUNCTION_NAME_CHARS + ), + }); + } + if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { + return Err(Error::InvalidInput { + message: format!( + "Invalid function name '{}': only alphanumeric \ + characters and underscores are allowed.", + sanitize(name) + ), + }); + } + Ok(()) } -/// Validate one component of a view identifier: a view name, or one segment of -/// the namespace path holding it. -pub fn validate_view_component(what: &str, value: &str) -> Result<()> { - check_object_name(value).map_err(|reason| Error::InvalidInput { - message: format!("invalid {what} '{value}': {reason}"), - }) -} - -/// Validate all components of a namespace -/// -/// Iterates through all namespace components and validates each one. -/// Returns an error if any component is invalid. -/// -/// # Arguments -/// * `namespace` - The namespace components to validate -/// -/// # Returns -/// * `Ok(())` if all namespace components are valid -/// * `Err(Error)` if any component is invalid -pub fn validate_namespace(namespace: &[String]) -> Result<()> { - for component in namespace { - validate_namespace_name(component)?; +/// Validate a database name, including slash-separated components. +pub fn validate_database_name(name: &str) -> Result<()> { + if name.len() > MAX_DB_NAME_CHARS { + return Err(Error::InvalidInput { + message: format!( + "Invalid database name '{}' is more than {} characters long.", + sanitize(name), + MAX_DB_NAME_CHARS + ), + }); + } + if !name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' || c == '/') + { + return Err(Error::InvalidInput { + message: format!( + "Invalid database name '{}': only alphanumeric characters, dashes, \ + underscores, periods, and slashes are allowed.", + sanitize(name) + ), + }); + } + for component in name.split('/') { + if component.is_empty() { + return Err(Error::InvalidInput { + message: format!( + "Invalid database name '{}': invalid empty path component.", + sanitize(name) + ), + }); + } + if component == "." || component == ".." { + return Err(Error::InvalidInput { + message: format!( + "Invalid database name '{}': illegal path component.", + sanitize(name) + ), + }); + } } Ok(()) } @@ -494,6 +567,26 @@ pub fn supported_fts_data_type(dtype: &DataType) -> bool { supported_fts_data_type_impl(dtype, false) } +/// Validate FTS input without losing the logical type in field metadata. +pub(crate) fn validate_fts_field(field: &Field) -> Result<()> { + if lance_arrow::json::is_json_field(field) || supported_fts_data_type(field.data_type()) { + return Ok(()); + } + let logical_type = field + .metadata() + .get("ARROW:extension:name") + .map(|name| format!(" (logical type {name})")) + .unwrap_or_default(); + Err(Error::Schema { + message: format!( + "A FTS index cannot be created on the field `{}` which has data type {}{}. FTS supports strings, lists of strings, and lance.json fields stored as LargeBinary; raw binary is not supported", + field.name(), + field.data_type(), + logical_type, + ), + }) +} + fn supported_fts_data_type_impl(dtype: &DataType, in_list: bool) -> bool { match (dtype, in_list) { (DataType::Utf8 | DataType::LargeUtf8, _) => true, @@ -707,6 +800,43 @@ mod tests { use super::*; + #[test] + fn test_fts_field_logical_type_validation() { + let json = lance_arrow::json::json_field("doc", true); + assert!(validate_fts_field(&json).is_ok()); + + for dtype in [DataType::Binary, DataType::LargeBinary] { + let raw = Field::new("doc", dtype.clone(), true); + let err = validate_fts_field(&raw).unwrap_err().to_string(); + assert!(err.contains("raw binary is not supported"), "{err}"); + assert!(err.contains("lance.json"), "{err}"); + + let other_extension = raw.with_metadata(std::collections::HashMap::from([( + "ARROW:extension:name".into(), + "other.json".into(), + )])); + let err = validate_fts_field(&other_extension) + .unwrap_err() + .to_string(); + assert!(err.contains("logical type other.json"), "{err}"); + } + + // A JSON marker on the wrong binary storage type is insufficient. + let wrong_storage = + Field::new("doc", DataType::Binary, true).with_metadata(json.metadata().clone()); + assert!(validate_fts_field(&wrong_storage).is_err()); + for dtype in [DataType::Utf8, DataType::LargeUtf8] { + let text = Field::new("text", dtype, true); + assert!(validate_fts_field(&text).is_ok()); + for list_type in [ + DataType::List(Arc::new(text.clone())), + DataType::LargeList(Arc::new(text.clone())), + ] { + assert!(validate_fts_field(&Field::new("texts", list_type, true)).is_ok()); + } + } + } + #[test] fn test_public_fts_field_path_prefers_exact_case() { let text_list = || { @@ -870,6 +1000,37 @@ mod tests { assert!(err.contains("text.embedding")); } + #[test] + fn test_validate_database_name() { + for name in [ + "db", + "team/search", + "Team_1/db-2.3", + &"a".repeat(MAX_DB_NAME_CHARS), + ] { + assert!(validate_database_name(name).is_ok(), "rejected {name:?}"); + } + for name in [ + "", + "/", + "//", + ".", + "..", + "/db", + "db/", + "a//b", + "a/./b", + "a/../b", + "a$b", + "a b", + "a?b", + "café", + &"a".repeat(MAX_DB_NAME_CHARS + 1), + ] { + assert!(validate_database_name(name).is_err(), "accepted {name:?}"); + } + } + #[test] fn test_validate_table_name() { assert!(validate_table_name("my_table").is_ok()); @@ -890,27 +1051,27 @@ mod tests { } #[test] - fn test_validate_namespace_name() { + fn test_validate_namespace_component() { // Valid namespace names - assert!(validate_namespace_name("ns1").is_ok()); - assert!(validate_namespace_name("namespace_123").is_ok()); - assert!(validate_namespace_name("my-namespace").is_ok()); - assert!(validate_namespace_name("my.namespace").is_ok()); - assert!(validate_namespace_name("NS_1.2.3").is_ok()); - assert!(validate_namespace_name("a").is_ok()); - assert!(validate_namespace_name("123").is_ok()); - assert!(validate_namespace_name("_underscore").is_ok()); - assert!(validate_namespace_name("-hyphen").is_ok()); - assert!(validate_namespace_name(".period").is_ok()); + assert!(validate_namespace_component("ns1").is_ok()); + assert!(validate_namespace_component("namespace_123").is_ok()); + assert!(validate_namespace_component("my-namespace").is_ok()); + assert!(validate_namespace_component("my.namespace").is_ok()); + assert!(validate_namespace_component("NS_1.2.3").is_ok()); + assert!(validate_namespace_component("a").is_ok()); + assert!(validate_namespace_component("123").is_ok()); + assert!(validate_namespace_component("_underscore").is_ok()); + assert!(validate_namespace_component("-hyphen").is_ok()); + assert!(validate_namespace_component(".period").is_ok()); // Invalid namespace names - assert!(validate_namespace_name("").is_err()); - assert!(validate_namespace_name("namespace with spaces").is_err()); - assert!(validate_namespace_name("namespace/with/slashes").is_err()); - assert!(validate_namespace_name("namespace\\with\\backslashes").is_err()); - assert!(validate_namespace_name("namespace$with$delimiter").is_err()); - assert!(validate_namespace_name("namespace@special").is_err()); - assert!(validate_namespace_name("namespace#hash").is_err()); + assert!(validate_namespace_component("").is_err()); + assert!(validate_namespace_component("namespace with spaces").is_err()); + assert!(validate_namespace_component("namespace/with/slashes").is_err()); + assert!(validate_namespace_component("namespace\\with\\backslashes").is_err()); + assert!(validate_namespace_component("namespace$with$delimiter").is_err()); + assert!(validate_namespace_component("namespace@special").is_err()); + assert!(validate_namespace_component("namespace#hash").is_err()); } #[test] @@ -926,6 +1087,9 @@ mod tests { // Empty namespace (root) is valid assert!(validate_namespace(&[]).is_ok()); + // Valid public namespace (which is equivalent to the root namespace.) + assert!(validate_namespace(&["public".to_string()]).is_ok()); + // Invalid: contains empty component assert!(validate_namespace(&["ns1".to_string(), "".to_string()]).is_err()); @@ -937,6 +1101,9 @@ mod tests { assert!(validate_namespace(&["ns1".to_string(), "ns/2".to_string()]).is_err()); assert!(validate_namespace(&["ns1".to_string(), "ns$2".to_string()]).is_err()); + // Invalid: starts with public but contains other components. + assert!(validate_namespace(&["public".to_string(), "foo".to_string()]).is_err()); + // Valid: underscores, hyphens, and periods are allowed assert!( validate_namespace(&["ns_1".to_string(), "ns-2".to_string(), "ns.3".to_string()]) diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json index c6bbd83e3..e07c29f98 100644 --- a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json @@ -23,7 +23,7 @@ "nullable": false } }, - "created_at": "2026-08-21T00:00:00Z", + "created_at_millis": 1787270400000, "image": { "manifest_digest": "sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6", "descriptor": { diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json index 63f3c824e..44a84b17c 100644 --- a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json @@ -1 +1 @@ -{"created_at":"2026-08-21T00:00:00Z","disabled":false,"image":{"descriptor":{"behavior":{"result_stability":"input_and_context","side_effects":"non_idempotent"},"entrypoint":"app.function:create","format_version":1,"interface":{"type":"lance.scalar","version":1},"python":{"abi_tag":"cp312","executable":"/usr/local/bin/python3","implementation":"cpython","import_paths":["/opt/function/code"],"version":"3.12.14"},"python_api":1,"requires":{"capabilities":[],"kernel_min":"4.18.0"},"schemas":{"initialization":"/opt/function/schemas/initialization.arrow","input":"/opt/function/schemas/input.arrow","output":"/opt/function/schemas/output.arrow"}},"manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6","source":false},"location":"memory:///fixture","metadata":{},"name":"embed","object_id":"fixture","secret_bindings":[{"kind":"env","secret_ref":{"name":"hf-prod"},"variable":"HF_TOKEN"}],"signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list","kind":"scalar","nullable":false}},"version":"1"} +{"created_at_millis":1787270400000,"disabled":false,"image":{"descriptor":{"behavior":{"result_stability":"input_and_context","side_effects":"non_idempotent"},"entrypoint":"app.function:create","format_version":1,"interface":{"type":"lance.scalar","version":1},"python":{"abi_tag":"cp312","executable":"/usr/local/bin/python3","implementation":"cpython","import_paths":["/opt/function/code"],"version":"3.12.14"},"python_api":1,"requires":{"capabilities":[],"kernel_min":"4.18.0"},"schemas":{"initialization":"/opt/function/schemas/initialization.arrow","input":"/opt/function/schemas/input.arrow","output":"/opt/function/schemas/output.arrow"}},"manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6","source":false},"location":"memory:///fixture","metadata":{},"name":"embed","object_id":"fixture","secret_bindings":[{"kind":"env","secret_ref":{"name":"hf-prod"},"variable":"HF_TOKEN"}],"signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list","kind":"scalar","nullable":false}},"version":"1"}