mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-03 20:18:54 +00:00
Compare commits
18 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8d3185dec6 | |||
| bb14667c81 | |||
| 983ad2c011 | |||
| bcff95b109 | |||
| dfbb5f5ade | |||
| 12ee5ca626 | |||
| fcef32bfd3 | |||
| 3863189b16 | |||
| c858e25494 | |||
| cf9a8bf91d | |||
| 7357d63e87 | |||
| 624a75edf7 | |||
| c7ea91f3ea | |||
| 8e24dd3828 | |||
| f79dc017c4 | |||
| e6ae93f52a | |||
| 3dd9c598e9 | |||
| 9e26bf3fba |
Generated
+259
-297
File diff suppressed because it is too large
Load Diff
+14
-14
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=10.1.0-beta.2", default-features = false, "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=10.1.0-beta.2", default-features = false, "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=10.1.0-beta.2", default-features = false, "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
|
||||
+1
-1
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>10.1.0-beta.2</lance-core.version>
|
||||
<lance-core.version>10.1.0-beta.1</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.version>
|
||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
# Keep Node's Linux x64 addons compatible with pre-Haswell CPUs.
|
||||
# lance-linalg dispatches hot vector kernels to newer SIMD tiers at runtime.
|
||||
[env]
|
||||
LANCEDB_NODE_ENFORCE_X86_64_V2 = "1"
|
||||
|
||||
[target.x86_64-unknown-linux-gnu]
|
||||
rustflags = [
|
||||
"-C",
|
||||
"target-cpu=x86-64-v2",
|
||||
"-C",
|
||||
"target-feature=-avx,-avx2,-fma,-f16c",
|
||||
]
|
||||
|
||||
# Preserve the workspace's dynamic C runtime configuration for musl.
|
||||
[target.x86_64-unknown-linux-musl]
|
||||
rustflags = [
|
||||
"-C",
|
||||
"target-cpu=x86-64-v2",
|
||||
"-C",
|
||||
"target-feature=-crt-static,-avx,-avx2,-fma,-f16c",
|
||||
]
|
||||
@@ -12,6 +12,10 @@ categories.workspace = true
|
||||
[lib]
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[[test]]
|
||||
name = "x86_64_v2_build_flags"
|
||||
path = "build_support/x86_64_v2.rs"
|
||||
|
||||
[dependencies]
|
||||
async-trait.workspace = true
|
||||
arrow-ipc.workspace = true
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as tmp from "tmp";
|
||||
|
||||
import { connect } from "../lancedb";
|
||||
import {
|
||||
Field,
|
||||
FixedSizeList,
|
||||
Float32,
|
||||
Int32,
|
||||
Schema,
|
||||
makeArrowTable,
|
||||
} from "../lancedb/arrow";
|
||||
|
||||
test("cosine vector search runs on the pre-Haswell build baseline", async () => {
|
||||
const tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
|
||||
try {
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int32(), false),
|
||||
new Field(
|
||||
"vector",
|
||||
new FixedSizeList(3, new Field("item", new Float32(), false)),
|
||||
false,
|
||||
),
|
||||
]);
|
||||
const data = makeArrowTable(
|
||||
[
|
||||
{ id: 1, vector: [1, 0, 0] },
|
||||
{ id: 2, vector: [0, 1, 0] },
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
const table = await db.createTable("vectors", data);
|
||||
|
||||
const results = await table
|
||||
.vectorSearch([1, 0, 0])
|
||||
.distanceType("cosine")
|
||||
.limit(1)
|
||||
.toArray();
|
||||
|
||||
expect(results).toHaveLength(1);
|
||||
expect(results[0].id).toBe(1);
|
||||
expect(results[0]._distance).toBeCloseTo(0);
|
||||
} finally {
|
||||
tmpDir.removeCallback();
|
||||
}
|
||||
}, 60_000);
|
||||
@@ -3,6 +3,29 @@
|
||||
|
||||
extern crate napi_build;
|
||||
|
||||
use std::env;
|
||||
|
||||
#[path = "build_support/x86_64_v2.rs"]
|
||||
mod x86_64_v2;
|
||||
|
||||
const ENFORCE_BASELINE: &str = "LANCEDB_NODE_ENFORCE_X86_64_V2";
|
||||
|
||||
fn main() {
|
||||
napi_build::setup();
|
||||
println!("cargo:rerun-if-env-changed={ENFORCE_BASELINE}");
|
||||
println!("cargo:rerun-if-env-changed=CARGO_ENCODED_RUSTFLAGS");
|
||||
|
||||
let is_linux_x64 = env::var("CARGO_CFG_TARGET_ARCH").as_deref() == Ok("x86_64")
|
||||
&& env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("linux");
|
||||
let is_release = env::var("PROFILE").as_deref() == Ok("release");
|
||||
let is_node_build = env::var(ENFORCE_BASELINE).as_deref() == Ok("1");
|
||||
|
||||
if !is_linux_x64 || (!is_release && !is_node_build) {
|
||||
return;
|
||||
}
|
||||
|
||||
let encoded_rustflags = env::var("CARGO_ENCODED_RUSTFLAGS").unwrap_or_default();
|
||||
x86_64_v2::validate_encoded_rustflags(&encoded_rustflags).unwrap_or_else(|error| {
|
||||
panic!("Linux x64 Node addons must use the x86-64-v2 baseline; {error}")
|
||||
});
|
||||
}
|
||||
|
||||
@@ -0,0 +1,430 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
const SAFE_TARGET_CPUS: [&str; 2] = ["x86-64", "x86-64-v2"];
|
||||
const BASELINE_FEATURES: [&str; 9] = [
|
||||
"cmpxchg16b",
|
||||
"fxsr",
|
||||
"popcnt",
|
||||
"sse",
|
||||
"sse2",
|
||||
"sse3",
|
||||
"sse4.1",
|
||||
"sse4.2",
|
||||
"ssse3",
|
||||
];
|
||||
|
||||
pub(crate) fn validate_encoded_rustflags(encoded: &str) -> Result<(), String> {
|
||||
let mut target_cpu = None;
|
||||
let mut feature_states = BTreeMap::new();
|
||||
let mut required_disables = BTreeSet::new();
|
||||
let mut unsupported_features = BTreeSet::new();
|
||||
|
||||
for option in codegen_options(encoded)? {
|
||||
let Some((name, value)) = option.split_once('=') else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let name = name.replace('_', "-");
|
||||
match name.as_str() {
|
||||
"target-cpu" => target_cpu = Some(value),
|
||||
"target-feature" => {
|
||||
for toggle in value.split(',').filter(|toggle| !toggle.is_empty()) {
|
||||
let (enabled, feature) = match toggle.as_bytes()[0] {
|
||||
b'+' => (true, &toggle[1..]),
|
||||
b'-' => (false, &toggle[1..]),
|
||||
_ => return Err(format!("invalid target feature flag: {toggle}")),
|
||||
};
|
||||
|
||||
feature_states.insert(feature, enabled);
|
||||
if enabled && !BASELINE_FEATURES.contains(&feature) {
|
||||
match feature {
|
||||
// These are inherited from the workspace configuration and
|
||||
// explicitly canceled by the Node configuration. Account for
|
||||
// their implied AVX prerequisite as well as the named feature.
|
||||
"avx" => {
|
||||
required_disables.insert("avx");
|
||||
}
|
||||
"avx2" => {
|
||||
required_disables.extend(["avx", "avx2"]);
|
||||
}
|
||||
"f16c" => {
|
||||
required_disables.extend(["avx", "f16c"]);
|
||||
}
|
||||
"fma" => {
|
||||
required_disables.extend(["avx", "fma"]);
|
||||
}
|
||||
_ => {
|
||||
unsupported_features.insert(feature);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// LLVM arguments can independently alter the target feature set and
|
||||
// cannot be proven safe by inspecting rustc's target options.
|
||||
"llvm-args" => return Err("LLVM arguments can override the CPU baseline".to_owned()),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(cpu) = target_cpu.filter(|cpu| !SAFE_TARGET_CPUS.contains(cpu)) {
|
||||
return Err(format!(
|
||||
"effective target CPU is {}, expected x86-64-v2 or lower",
|
||||
cpu
|
||||
));
|
||||
}
|
||||
|
||||
if !unsupported_features.is_empty() {
|
||||
return Err(format!(
|
||||
"features above v2: {}",
|
||||
unsupported_features
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
));
|
||||
}
|
||||
|
||||
let not_disabled = required_disables
|
||||
.into_iter()
|
||||
.filter(|feature| feature_states.get(feature) != Some(&false))
|
||||
.collect::<Vec<_>>();
|
||||
if !not_disabled.is_empty() {
|
||||
return Err(format!(
|
||||
"inherited features not fully disabled: {}",
|
||||
not_disabled.join(", ")
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn codegen_options(encoded: &str) -> Result<Vec<&str>, String> {
|
||||
let arguments = encoded.split('\u{1f}').collect::<Vec<_>>();
|
||||
if arguments.iter().any(|argument| argument.starts_with('@')) {
|
||||
return Err("rustc response-file arguments cannot be validated".to_owned());
|
||||
}
|
||||
|
||||
let mut options = Vec::new();
|
||||
let mut index = 0;
|
||||
|
||||
while index < arguments.len() {
|
||||
let argument = arguments[index];
|
||||
if argument == "--cfg" {
|
||||
index += 1;
|
||||
let cfg = arguments
|
||||
.get(index)
|
||||
.copied()
|
||||
.ok_or_else(|| "missing value after --cfg".to_owned())?;
|
||||
reject_builtin_target_feature_cfg(cfg)?;
|
||||
} else if let Some(cfg) = argument.strip_prefix("--cfg=") {
|
||||
reject_builtin_target_feature_cfg(cfg)?;
|
||||
} else if argument == "-C" || argument == "--codegen" {
|
||||
index += 1;
|
||||
let option = arguments
|
||||
.get(index)
|
||||
.copied()
|
||||
.ok_or_else(|| format!("missing value after {argument}"))?;
|
||||
options.push(option.trim_start_matches('='));
|
||||
} else if let Some(option) = argument.strip_prefix("-C") {
|
||||
if !option.is_empty() {
|
||||
options.push(option.trim_start_matches('='));
|
||||
}
|
||||
} else if let Some(option) = argument.strip_prefix("--codegen=") {
|
||||
options.push(option);
|
||||
}
|
||||
index += 1;
|
||||
}
|
||||
|
||||
Ok(options)
|
||||
}
|
||||
|
||||
fn reject_builtin_target_feature_cfg(cfg: &str) -> Result<(), String> {
|
||||
let key = cfg.split_once('=').map_or(cfg, |(key, _)| key);
|
||||
if key.contains("/*") || key.contains("//") {
|
||||
return Err("comment-bearing cfgs cannot be validated".to_owned());
|
||||
}
|
||||
if !key.is_ascii() {
|
||||
return Err("non-ASCII cfg keys cannot be validated".to_owned());
|
||||
}
|
||||
|
||||
let name = key.trim();
|
||||
let name = name.strip_prefix("r#").unwrap_or(name);
|
||||
if name == "target_feature" {
|
||||
return Err("built-in target_feature cfgs can override runtime CPU detection".to_owned());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn encoded(arguments: &[&str]) -> String {
|
||||
arguments.join("\u{1f}")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_merged_workspace_and_node_flags() {
|
||||
let flags = encoded(&[
|
||||
"-C",
|
||||
"target-cpu=haswell",
|
||||
"-C",
|
||||
"target-feature=+avx2,+fma,+f16c",
|
||||
"-C",
|
||||
"target-cpu=x86-64-v2",
|
||||
"-C",
|
||||
"target-feature=-avx,-avx2,-fma,-f16c",
|
||||
]);
|
||||
|
||||
assert_eq!(validate_encoded_rustflags(&flags), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_non_merging_v2_boundary() {
|
||||
let flags = encoded(&["-Ctarget-cpu=x86-64-v2"]);
|
||||
|
||||
assert_eq!(validate_encoded_rustflags(&flags), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_default_cpu_with_non_codegen_flags() {
|
||||
let flags = encoded(&["-D", "warnings"]);
|
||||
|
||||
assert_eq!(validate_encoded_rustflags(&flags), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_unrelated_custom_cfg() {
|
||||
let flags = encoded(&["--cfg=tokio_unstable"]);
|
||||
|
||||
assert_eq!(validate_encoded_rustflags(&flags), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_comment_like_syntax_in_cfg_value() {
|
||||
let flags = encoded(&[r#"--cfg=endpoint="https://example.com/*""#]);
|
||||
|
||||
assert_eq!(validate_encoded_rustflags(&flags), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_explicit_x86_64_v1_cpu() {
|
||||
let flags = encoded(&["-Ctarget-cpu=x86-64"]);
|
||||
|
||||
assert_eq!(validate_encoded_rustflags(&flags), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_musl_dynamic_crt_configuration() {
|
||||
let flags = encoded(&[
|
||||
"-C",
|
||||
"target-cpu=haswell",
|
||||
"-C",
|
||||
"target-feature=-crt-static,+avx2,+fma,+f16c",
|
||||
"-C",
|
||||
"target-cpu=x86-64-v2",
|
||||
"-C",
|
||||
"target-feature=-crt-static,-avx,-avx2,-fma,-f16c",
|
||||
]);
|
||||
|
||||
assert_eq!(validate_encoded_rustflags(&flags), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_feature_omitted_from_target_cfg() {
|
||||
let flags = encoded(&["-Ctarget-cpu=x86-64-v2", "-Ctarget-feature=+apxf"]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("features above v2: apxf".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_underscore_spelling_above_baseline_feature() {
|
||||
let flags = encoded(&["-Ctarget_cpu=x86-64-v2", "-Ctarget_feature=+apxf"]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("features above v2: apxf".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unexpected_above_baseline_feature() {
|
||||
let flags = encoded(&[
|
||||
"--codegen=target-cpu=x86-64-v2",
|
||||
"--codegen",
|
||||
"target-feature=+bmi2",
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("features above v2: bmi2".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_cpu_above_v2() {
|
||||
let flags = encoded(&["-Ctarget-cpu=haswell"]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("effective target CPU is haswell, expected x86-64-v2 or lower".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_incompletely_disabled_feature_implications() {
|
||||
let flags = encoded(&["-Ctarget-cpu=x86-64-v2", "-Ctarget-feature=+avx2,-avx2"]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("inherited features not fully disabled: avx".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_llvm_feature_overrides() {
|
||||
let flags = encoded(&["-Ctarget-cpu=x86-64-v2", "-Cllvm-args=-mattr=+apxf"]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("LLVM arguments can override the CPU baseline".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_underscore_spelling_llvm_feature_overrides() {
|
||||
let flags = encoded(&["-Ctarget_cpu=x86-64-v2", "-Cllvm_args=-mattr=+apxf"]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("LLVM arguments can override the CPU baseline".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_response_file_arguments() {
|
||||
let flags = encoded(&["-Ctarget-cpu=x86-64-v2", "@flags.rsp"]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("rustc response-file arguments cannot be validated".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_split_builtin_target_feature_cfg() {
|
||||
let flags = encoded(&[
|
||||
"-Ctarget-cpu=x86-64-v2",
|
||||
"--cfg",
|
||||
r#"target_feature="avx2""#,
|
||||
"-Aexplicit_builtin_cfgs_in_flags",
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("built-in target_feature cfgs can override runtime CPU detection".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_equals_builtin_target_feature_cfg() {
|
||||
let flags = encoded(&[
|
||||
"-Ctarget-cpu=x86-64-v2",
|
||||
r#"--cfg=target_feature="avx2""#,
|
||||
"-Aexplicit_builtin_cfgs_in_flags",
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("built-in target_feature cfgs can override runtime CPU detection".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_raw_identifier_builtin_target_feature_cfg() {
|
||||
let flags = encoded(&[
|
||||
"-Ctarget-cpu=x86-64-v2",
|
||||
r#"--cfg=r#target_feature="avx2""#,
|
||||
"-Aexplicit_builtin_cfgs_in_flags",
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("built-in target_feature cfgs can override runtime CPU detection".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_block_comment_after_cfg_name() {
|
||||
let flags = encoded(&[
|
||||
"--cfg",
|
||||
r#"target_feature/*gate*/="avx2""#,
|
||||
"-Aexplicit_builtin_cfgs_in_flags",
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("comment-bearing cfgs cannot be validated".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_block_comment_before_cfg_name() {
|
||||
let flags = encoded(&[
|
||||
r#"--cfg=/*gate*/target_feature="avx2""#,
|
||||
"-Aexplicit_builtin_cfgs_in_flags",
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("comment-bearing cfgs cannot be validated".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_line_comment_cfg_trivia() {
|
||||
let flags = encoded(&[
|
||||
"--cfg",
|
||||
"target_feature// gate\n=\"avx2\"",
|
||||
"-Aexplicit_builtin_cfgs_in_flags",
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("comment-bearing cfgs cannot be validated".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_leading_bom_in_cfg_key() {
|
||||
let flags = encoded(&[
|
||||
"--cfg",
|
||||
"\u{feff}target_feature=\"avx2\"",
|
||||
"-Aexplicit_builtin_cfgs_in_flags",
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("non-ASCII cfg keys cannot be validated".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_ascii_pattern_whitespace_in_cfg_key() {
|
||||
for whitespace in ['\u{200e}', '\u{200f}'] {
|
||||
let cfg = format!("{whitespace}target_feature=\"avx2\"");
|
||||
let flags = encoded(&["--cfg", &cfg, "-Aexplicit_builtin_cfgs_in_flags"]);
|
||||
|
||||
assert_eq!(
|
||||
validate_encoded_rustflags(&flags),
|
||||
Err("non-ASCII cfg keys cannot be validated".to_owned())
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
const assert = require("node:assert/strict");
|
||||
const fs = require("node:fs");
|
||||
const os = require("node:os");
|
||||
const path = require("node:path");
|
||||
|
||||
const { connect } = require("../dist");
|
||||
|
||||
async function main() {
|
||||
const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "lancedb-cpu-"));
|
||||
|
||||
try {
|
||||
const db = await connect(tmpDir);
|
||||
const table = await db.createTable("vectors", [
|
||||
{ id: 1, vector: [1, 0, 0] },
|
||||
{ id: 2, vector: [0, 1, 0] },
|
||||
]);
|
||||
const results = await table
|
||||
.vectorSearch([1, 0, 0])
|
||||
.distanceType("cosine")
|
||||
.limit(1)
|
||||
.toArray();
|
||||
|
||||
assert.equal(results.length, 1);
|
||||
assert.equal(results[0].id, 1);
|
||||
assert.ok(Math.abs(results[0]._distance) < 1e-6);
|
||||
} finally {
|
||||
fs.rmSync(tmpDir, { force: true, recursive: true });
|
||||
}
|
||||
}
|
||||
|
||||
main().catch((error) => {
|
||||
console.error(error);
|
||||
process.exitCode = 1;
|
||||
});
|
||||
@@ -88,6 +88,7 @@
|
||||
"lint-fix": "biome check --write . && biome format --write .",
|
||||
"prepublishOnly": "napi prepublish -t npm",
|
||||
"test": "jest --verbose",
|
||||
"test:pre-haswell": "node ci/pre_haswell_smoke.js",
|
||||
"integration": "S3_TEST=1 pnpm test",
|
||||
"universal": "napi universalize",
|
||||
"version": "napi version"
|
||||
|
||||
+6
-2
@@ -339,7 +339,9 @@ impl Table {
|
||||
let transforms = NewColumnTransform::SqlExpressions(transforms);
|
||||
let res = self
|
||||
.inner_ref()?
|
||||
.add_columns(transforms, None)
|
||||
.add_columns()
|
||||
.transform(transforms)
|
||||
.execute()
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(res.into())
|
||||
@@ -356,7 +358,9 @@ impl Table {
|
||||
let transforms = NewColumnTransform::AllNulls(schema);
|
||||
let res = self
|
||||
.inner_ref()?
|
||||
.add_columns(transforms, None)
|
||||
.add_columns()
|
||||
.transform(transforms)
|
||||
.execute()
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(res.into())
|
||||
|
||||
@@ -707,6 +707,9 @@ class LanceDBConnection(DBConnection):
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
if _inner is not None:
|
||||
self._conn = _inner
|
||||
# Native-derived wrappers resolve this in their async reconstruction
|
||||
# path so construction never synchronously re-enters LOOP.
|
||||
self._read_consistency_interval = read_consistency_interval
|
||||
self._cached_namespace_client = None
|
||||
return
|
||||
|
||||
@@ -756,11 +759,14 @@ class LanceDBConnection(DBConnection):
|
||||
# storage_options. Also, this class really shouldn't be holding any state
|
||||
# beyond _conn.
|
||||
self._conn = AsyncConnection(LOOP.run(do_connect()))
|
||||
# Keep property access synchronous so debugger introspection cannot wait on
|
||||
# the background loop while that thread is suspended at a breakpoint.
|
||||
self._read_consistency_interval = read_consistency_interval
|
||||
self._cached_namespace_client: Optional[LanceNamespace] = None
|
||||
|
||||
@property
|
||||
def read_consistency_interval(self) -> Optional[timedelta]:
|
||||
return LOOP.run(self._conn.get_read_consistency_interval())
|
||||
return self._read_consistency_interval
|
||||
|
||||
@property
|
||||
def session(self) -> Optional[Session]:
|
||||
@@ -771,8 +777,16 @@ class LanceDBConnection(DBConnection):
|
||||
return self._conn.uri
|
||||
|
||||
@classmethod
|
||||
def from_inner(cls, inner: LanceDbConnection):
|
||||
return cls(None, _inner=inner)
|
||||
def from_inner(
|
||||
cls,
|
||||
inner: LanceDbConnection,
|
||||
read_consistency_interval: Optional[timedelta],
|
||||
):
|
||||
return cls(
|
||||
None,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
_inner=inner,
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
|
||||
|
||||
@@ -92,8 +92,10 @@ class LanceMergeInsertBuilder(object):
|
||||
self._when_not_matched_by_source_delete = True
|
||||
if isinstance(condition, Expr):
|
||||
self._when_not_matched_by_source_condition_expr = condition._inner
|
||||
elif condition is not None:
|
||||
self._when_not_matched_by_source_condition = None
|
||||
else:
|
||||
self._when_not_matched_by_source_condition = condition
|
||||
self._when_not_matched_by_source_condition_expr = None
|
||||
return self
|
||||
|
||||
def use_index(self, use_index: bool) -> LanceMergeInsertBuilder:
|
||||
|
||||
@@ -226,7 +226,7 @@ class PermutationBuilder:
|
||||
|
||||
async def do_execute():
|
||||
inner_tbl = await self._async.execute()
|
||||
return LanceTable.from_inner(inner_tbl)
|
||||
return await LanceTable.from_inner(inner_tbl)
|
||||
|
||||
return LOOP.run(do_execute())
|
||||
|
||||
|
||||
@@ -2697,7 +2697,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._fts_query.phrase_query(True)
|
||||
if self._distance_type:
|
||||
self._vector_query.metric(self._distance_type)
|
||||
if self._minimum_nprobes:
|
||||
if self._minimum_nprobes is not None:
|
||||
self._vector_query.minimum_nprobes(self._minimum_nprobes)
|
||||
if self._maximum_nprobes is not None:
|
||||
self._vector_query.maximum_nprobes(self._maximum_nprobes)
|
||||
|
||||
@@ -1595,8 +1595,8 @@ class Table(ABC):
|
||||
Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is
|
||||
a ``list[int]`` or a query ``pyarrow.Table`` carrying row identity via
|
||||
``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null
|
||||
rows are ``None``. Unsupported on LanceDB Cloud, where
|
||||
:meth:`fetch_blobs` returns full bytes instead.
|
||||
rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or
|
||||
newer.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -2182,11 +2182,15 @@ class LanceTable(Table):
|
||||
return self.name
|
||||
|
||||
@classmethod
|
||||
def from_inner(cls, tbl: LanceDBTable):
|
||||
from .db import LanceDBConnection
|
||||
async def from_inner(cls, tbl: LanceDBTable):
|
||||
from .db import AsyncConnection, LanceDBConnection
|
||||
|
||||
async_tbl = AsyncTable(tbl)
|
||||
conn = LanceDBConnection.from_inner(tbl.database())
|
||||
inner_conn = tbl.database()
|
||||
read_consistency_interval = await AsyncConnection(
|
||||
inner_conn
|
||||
).get_read_consistency_interval()
|
||||
conn = LanceDBConnection.from_inner(inner_conn, read_consistency_interval)
|
||||
return cls(
|
||||
conn,
|
||||
async_tbl.name,
|
||||
|
||||
@@ -77,6 +77,23 @@ def test_sync_repr_does_not_use_background_loop(tmp_path, monkeypatch):
|
||||
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
|
||||
|
||||
|
||||
def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkeypatch):
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.db import LanceDBConnection
|
||||
|
||||
consistency_interval = timedelta(seconds=5)
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
|
||||
db_from_inner = LanceDBConnection.from_inner(db._inner, consistency_interval)
|
||||
|
||||
def fail_run(*args, **kwargs):
|
||||
raise AssertionError("properties should not use the Python background loop")
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", fail_run)
|
||||
|
||||
assert db.read_consistency_interval == consistency_interval
|
||||
assert db_from_inner.read_consistency_interval == consistency_interval
|
||||
|
||||
|
||||
def test_ingest_pd(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
|
||||
@@ -123,6 +123,19 @@ async def test_async_hybrid_query_default_limit(table: AsyncTable):
|
||||
assert texts.count("a") == 1
|
||||
|
||||
|
||||
def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
|
||||
# minimum_nprobes(0) must raise the same validation error a plain vector
|
||||
# query raises, not silently no-op because 0 is falsy.
|
||||
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
|
||||
(
|
||||
sync_table.search(query_type="hybrid")
|
||||
.vector([0.0, 0.4])
|
||||
.text("dog")
|
||||
.minimum_nprobes(0)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
|
||||
def test_hybrid_query_distance_range(sync_table: Table):
|
||||
reranker = RRFReranker(return_score="all")
|
||||
result = (
|
||||
|
||||
@@ -6,6 +6,7 @@ import math
|
||||
import pytest
|
||||
|
||||
from lancedb import DBConnection, Table, connect
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.permutation import Permutation, Permutations, permutation_builder
|
||||
|
||||
|
||||
@@ -31,6 +32,25 @@ def test_split_random_ratios(mem_db):
|
||||
assert 65 <= split_1_count <= 75 # ~70% ± tolerance
|
||||
|
||||
|
||||
def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
|
||||
import threading
|
||||
|
||||
db = connect(tmp_path)
|
||||
tbl = db.create_table("test_table", pa.table({"x": range(10)}))
|
||||
original_run = LOOP.run
|
||||
|
||||
def fail_on_reentry(future):
|
||||
assert threading.current_thread() is not LOOP.thread
|
||||
return original_run(future)
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", fail_on_reentry)
|
||||
|
||||
permutation_tbl = permutation_builder(tbl).execute()
|
||||
|
||||
assert permutation_tbl.count_rows() == 10
|
||||
assert permutation_tbl._conn.read_consistency_interval is None
|
||||
|
||||
|
||||
def test_split_random_counts(mem_db):
|
||||
"""Test random splitting with absolute counts."""
|
||||
tbl = mem_db.create_table(
|
||||
|
||||
@@ -2055,6 +2055,24 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
|
||||
request.send_header("phalanx-version", str(server_version))
|
||||
request.end_headers()
|
||||
request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
|
||||
elif request.path.startswith("/v1/table/test/blob/image/"):
|
||||
path = request.path.partition("?")[0]
|
||||
row_id = int(path.split("/")[-2])
|
||||
payload = {10: b"alpha", 20: None, 30: b"gamma"}[row_id]
|
||||
if payload is None:
|
||||
request.send_response(204)
|
||||
request.end_headers()
|
||||
return
|
||||
byte_range = request.headers["Range"].removeprefix("bytes=")
|
||||
start_text, end_text = byte_range.split("-", maxsplit=1)
|
||||
start = int(start_text)
|
||||
end = int(end_text) if end_text else len(payload) - 1
|
||||
chunk = payload[start : end + 1]
|
||||
request.send_response(206)
|
||||
request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}")
|
||||
request.send_header("Content-Length", str(len(chunk)))
|
||||
request.end_headers()
|
||||
request.wfile.write(chunk)
|
||||
elif request.path == "/v1/table/test/query/":
|
||||
content_len = int(request.headers.get("Content-Length", 0))
|
||||
body = json.loads(request.rfile.read(content_len))
|
||||
@@ -2092,8 +2110,21 @@ def test_remote_blob_columns_and_fetch():
|
||||
assert table.blob_columns() == ["image"]
|
||||
blobs = table.fetch_blobs("image", [10, 20, 30])
|
||||
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
|
||||
with pytest.raises(NotImplementedError, match="Use fetch_blobs for full bytes"):
|
||||
table.fetch_blob_files("image", [10, 20, 30])
|
||||
|
||||
|
||||
def test_remote_blob_files_are_lazy_seekable_handles():
|
||||
with blob_remote_table() as table:
|
||||
files = table.fetch_blob_files("image", [10, 20, 30])
|
||||
|
||||
assert len(files) == 3
|
||||
alpha, null_row, gamma = files
|
||||
assert null_row is None
|
||||
assert alpha is not None
|
||||
assert gamma is not None
|
||||
assert alpha.size() == 5
|
||||
assert alpha.read_range(1, 3) == b"lph"
|
||||
gamma.seek(2)
|
||||
assert gamma.read() == b"mma"
|
||||
|
||||
|
||||
def test_remote_blob_fetch_accepts_query_table():
|
||||
|
||||
@@ -6,6 +6,7 @@ import os
|
||||
import sys
|
||||
import threading
|
||||
import warnings
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import date, datetime, timedelta
|
||||
from time import sleep
|
||||
from typing import List
|
||||
@@ -2124,6 +2125,27 @@ def test_delete(mem_db: DBConnection):
|
||||
assert table.to_arrow()["id"].to_pylist() == [1]
|
||||
|
||||
|
||||
def test_concurrent_deletes_are_thread_safe(mem_db: DBConnection):
|
||||
num_workers = 8
|
||||
table = mem_db.create_table(
|
||||
"my_table", data=[{"id": row_id} for row_id in range(num_workers)]
|
||||
)
|
||||
barrier = threading.Barrier(num_workers)
|
||||
|
||||
def delete(row_id: int):
|
||||
barrier.wait()
|
||||
return table.delete(f"id = {row_id}")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=num_workers) as pool:
|
||||
results = list(pool.map(delete, range(num_workers)))
|
||||
|
||||
assert all(result.num_deleted_rows == 1 for result in results)
|
||||
assert sorted(result.version for result in results) == list(
|
||||
range(2, num_workers + 2)
|
||||
)
|
||||
assert table.count_rows() == 0
|
||||
|
||||
|
||||
def test_delete_expr(mem_db: DBConnection):
|
||||
table = mem_db.create_table(
|
||||
"my_table",
|
||||
@@ -2364,6 +2386,29 @@ def test_merge_insert_by_source_delete_expr(mem_db: DBConnection):
|
||||
assert table.to_arrow().sort_by("a") == expected
|
||||
|
||||
|
||||
def test_merge_insert_by_source_delete_reconfigure(mem_db: DBConnection):
|
||||
# Calling when_not_matched_by_source_delete() again with no condition must
|
||||
# widen the delete to unconditional, not keep the earlier condition around.
|
||||
table = mem_db.create_table(
|
||||
"my_table",
|
||||
data=pa.table({"a": [1, 2, 3], "b": ["a", "b", "c"]}),
|
||||
)
|
||||
new_data = pa.table({"a": [2, 4], "b": ["x", "z"]})
|
||||
|
||||
merge_insert_res = (
|
||||
table.merge_insert("a")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.when_not_matched_by_source_delete("a > 2")
|
||||
.when_not_matched_by_source_delete()
|
||||
.execute(new_data)
|
||||
)
|
||||
assert merge_insert_res.num_deleted_rows == 2
|
||||
|
||||
expected = pa.table({"a": [2, 4], "b": ["x", "z"]})
|
||||
assert table.to_arrow().sort_by("a") == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_insert_by_source_delete_expr_async(
|
||||
mem_db_async: AsyncConnection,
|
||||
|
||||
+40
-13
@@ -426,9 +426,11 @@ pub struct PyBlobFile {
|
||||
impl PyBlobFile {
|
||||
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read().await })
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read().await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
pub fn read(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
@@ -444,24 +446,32 @@ impl PyBlobFile {
|
||||
|
||||
fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.close().await })
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.close().await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
|
||||
}
|
||||
|
||||
fn is_closed(self_: PyRef<'_, Self>) -> bool {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.is_closed().await })
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.is_closed().await }))
|
||||
}
|
||||
|
||||
fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.seek(position).await })
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.seek(position).await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
|
||||
}
|
||||
|
||||
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.tell().await })
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.tell().await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
|
||||
}
|
||||
|
||||
@@ -475,16 +485,20 @@ impl PyBlobFile {
|
||||
.checked_add(length as u64)
|
||||
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read_range(offset..end).await })
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_range(offset..end).await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read_up_to(length).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_up_to(length).await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?;
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -731,6 +745,9 @@ impl Table {
|
||||
|
||||
#[allow(private_interfaces)]
|
||||
pub fn delete(self_: PyRef<'_, Self>, condition: PredicateArg) -> PyResult<Bound<'_, PyAny>> {
|
||||
// Do not hold the Python borrow across the await. The cloned Rust table
|
||||
// handle is thread-safe and allows deletes on the same Python table to
|
||||
// run concurrently without PyO3 reporting "Already borrowed".
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = match &condition {
|
||||
@@ -1361,7 +1378,12 @@ impl Table {
|
||||
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner.add_columns(definitions, None).await.infer_error()?;
|
||||
let result = inner
|
||||
.add_columns()
|
||||
.transform(definitions)
|
||||
.execute()
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(AddColumnsResult::from(result))
|
||||
})
|
||||
}
|
||||
@@ -1375,7 +1397,12 @@ impl Table {
|
||||
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner.add_columns(transform, None).await.infer_error()?;
|
||||
let result = inner
|
||||
.add_columns()
|
||||
.transform(transform)
|
||||
.execute()
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(AddColumnsResult::from(result))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -49,8 +49,8 @@ lance-namespace = { workspace = true }
|
||||
lance-namespace-impls = { workspace = true }
|
||||
metrics = { workspace = true, optional = true }
|
||||
metrics-util = { workspace = true, optional = true }
|
||||
# Keep the direct dependency aligned with the version required by OpenDAL.
|
||||
goosefs-sdk = { version = "=0.1.8", optional = true }
|
||||
# Pin the transitive GooseFS SDK until the 0.1.6 compile break is fixed upstream.
|
||||
goosefs-sdk = { version = "=0.1.5", optional = true }
|
||||
moka = { workspace = true }
|
||||
pin-project = { workspace = true }
|
||||
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
||||
|
||||
+198
-2
@@ -9,6 +9,7 @@
|
||||
//!
|
||||
//! Blob tables require Lance file format >= 2.2 and stable row ids at create.
|
||||
|
||||
use std::ops::Range;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::LargeBinaryArray;
|
||||
@@ -17,10 +18,202 @@ use arrow_schema::{DataType, Field, Schema};
|
||||
use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams};
|
||||
use lance_arrow::FieldExt;
|
||||
use lance_encoding::version::LanceFileVersion;
|
||||
use lance_io::object_store::ObjectStore;
|
||||
use object_store::path::Path;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
pub use lance::dataset::BlobFile;
|
||||
/// Seekable handle for one blob value, backed by local storage or a remote
|
||||
/// HTTP byte-range endpoint.
|
||||
#[derive(Debug)]
|
||||
pub struct BlobFile {
|
||||
inner: BlobFileInner,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum BlobFileInner {
|
||||
Native(lance::dataset::BlobFile),
|
||||
#[cfg(feature = "remote")]
|
||||
Remote(Box<crate::remote::table::blobs::RemoteBlobFile>),
|
||||
}
|
||||
|
||||
impl From<lance::dataset::BlobFile> for BlobFile {
|
||||
fn from(value: lance::dataset::BlobFile) -> Self {
|
||||
Self {
|
||||
inner: BlobFileInner::Native(value),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "remote")]
|
||||
impl From<crate::remote::table::blobs::RemoteBlobFile> for BlobFile {
|
||||
fn from(value: crate::remote::table::blobs::RemoteBlobFile) -> Self {
|
||||
Self {
|
||||
inner: BlobFileInner::Remote(Box::new(value)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BlobFile {
|
||||
/// Inline reader over a data-file slice.
|
||||
pub fn new_inline(
|
||||
object_store: Arc<ObjectStore>,
|
||||
path: Path,
|
||||
position: u64,
|
||||
size: u64,
|
||||
) -> Self {
|
||||
lance::dataset::BlobFile::new_inline(object_store, path, position, size).into()
|
||||
}
|
||||
|
||||
/// Dedicated sidecar-file reader.
|
||||
pub fn new_dedicated(object_store: Arc<ObjectStore>, path: Path, size: u64) -> Self {
|
||||
lance::dataset::BlobFile::new_dedicated(object_store, path, size).into()
|
||||
}
|
||||
|
||||
/// Packed reader for a slice in a shared sidecar.
|
||||
pub fn new_packed(
|
||||
object_store: Arc<ObjectStore>,
|
||||
path: Path,
|
||||
position: u64,
|
||||
size: u64,
|
||||
) -> Self {
|
||||
lance::dataset::BlobFile::new_packed(object_store, path, position, size).into()
|
||||
}
|
||||
|
||||
/// External reader at a resolved object location.
|
||||
pub fn new_external(
|
||||
object_store: Arc<ObjectStore>,
|
||||
path: Path,
|
||||
uri: String,
|
||||
position: u64,
|
||||
size: u64,
|
||||
) -> Self {
|
||||
lance::dataset::BlobFile::new_external(object_store, path, uri, position, size).into()
|
||||
}
|
||||
|
||||
/// Close the handle.
|
||||
pub async fn close(&self) -> lance_core::Result<()> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.close().await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.close().await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the handle is closed.
|
||||
pub async fn is_closed(&self) -> bool {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.is_closed().await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.is_closed(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Read a range without moving the cursor.
|
||||
pub async fn read_range(&self, range: Range<u64>) -> lance_core::Result<bytes::Bytes> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.read_range(range).await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.read_range(range).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read ranges without moving the cursor.
|
||||
pub async fn read_ranges(
|
||||
&self,
|
||||
ranges: &[Range<u64>],
|
||||
) -> lance_core::Result<Vec<bytes::Bytes>> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.read_ranges(ranges).await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.read_ranges(ranges).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read from the cursor to the end.
|
||||
pub async fn read(&self) -> lance_core::Result<bytes::Bytes> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.read().await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.read().await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read up to `len` bytes and advance the cursor.
|
||||
pub async fn read_up_to(&self, len: usize) -> lance_core::Result<bytes::Bytes> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.read_up_to(len).await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.read_up_to(len).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Move the cursor to `new_cursor`.
|
||||
pub async fn seek(&self, new_cursor: u64) -> lance_core::Result<()> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.seek(new_cursor).await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.seek(new_cursor).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Current cursor position.
|
||||
pub async fn tell(&self) -> lance_core::Result<u64> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.tell().await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.tell().await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Blob length in bytes.
|
||||
pub fn size(&self) -> u64 {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.size(),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.size(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Physical byte offset in the data file. `None` on remote handles. The
|
||||
/// Cloud byte-range route does not expose storage layout.
|
||||
pub fn position(&self) -> Option<u64> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => Some(file.position()),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Path of the data file holding the blob. `None` on remote handles. The
|
||||
/// Cloud byte-range route does not expose storage layout.
|
||||
pub fn data_path(&self) -> Option<&Path> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => Some(file.data_path()),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Native storage layout. `None` on remote handles. The Cloud byte-range
|
||||
/// route does not expose layout.
|
||||
pub fn kind(&self) -> Option<lance_core::datatypes::BlobKind> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => Some(file.kind()),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// External URI for native handles. Remote handles do not expose storage URIs.
|
||||
pub fn uri(&self) -> Option<&str> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.uri(),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One row-specific blob range read request.
|
||||
///
|
||||
@@ -264,7 +457,10 @@ pub(crate) async fn take_blob_files_aligned(
|
||||
|
||||
let handles = dataset.take_blobs(row_ids, column).await?;
|
||||
ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?;
|
||||
Ok(handles)
|
||||
Ok(handles
|
||||
.into_iter()
|
||||
.map(|handle| handle.map(Into::into))
|
||||
.collect())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
mod blobs;
|
||||
pub mod blobs;
|
||||
pub mod insert;
|
||||
|
||||
use self::insert::{RemoteWriteExec, WriteOp};
|
||||
@@ -2939,7 +2939,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
}
|
||||
|
||||
#[derive(Serialize, Clone, Debug)]
|
||||
pub struct MergeInsertRequest {
|
||||
pub(crate) struct MergeInsertRequest {
|
||||
on: String,
|
||||
when_matched_update_all: bool,
|
||||
when_matched_update_all_filt: Option<String>,
|
||||
@@ -3089,10 +3089,12 @@ mod tests {
|
||||
Box::pin(table.delete("false").map_ok(|_| ())),
|
||||
Box::pin(
|
||||
table
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("x".into(), "y".into())]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"x".into(),
|
||||
"y".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.map_ok(|_| ()),
|
||||
),
|
||||
Box::pin(async {
|
||||
@@ -4300,32 +4302,9 @@ mod tests {
|
||||
"fetch_blobs",
|
||||
);
|
||||
|
||||
let message = table
|
||||
.fetch_blob_files("image", &[1])
|
||||
.await
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(
|
||||
message.contains("fetch_blob_files is not supported on LanceDB Cloud"),
|
||||
"got: {message}"
|
||||
);
|
||||
assert!(
|
||||
!message.contains("Use fetch_blobs"),
|
||||
"old server must not be told to use fetch_blobs, got: {message}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_blob_files_point_at_fetch_blobs_on_a_blob_capable_server() {
|
||||
let table = Table::new_with_handler_version(
|
||||
"my_table",
|
||||
semver::Version::new(0, 5, 0),
|
||||
|_| -> http::Response<String> { panic!("fetch_blob_files must not reach the server") },
|
||||
);
|
||||
|
||||
assert_not_supported_error(
|
||||
table.fetch_blob_files("image", &[1]).await.unwrap_err(),
|
||||
"Use fetch_blobs for full bytes",
|
||||
"requires LanceDB Cloud server 0.5.0 or newer",
|
||||
);
|
||||
}
|
||||
|
||||
@@ -6411,13 +6390,12 @@ mod tests {
|
||||
});
|
||||
|
||||
let result = table
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![
|
||||
("b".into(), "a + 1".into()),
|
||||
("x".into(), "cast(NULL as int32)".into()),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![
|
||||
("b".into(), "a + 1".into()),
|
||||
("x".into(), "cast(NULL as int32)".into()),
|
||||
]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -7142,10 +7120,12 @@ mod tests {
|
||||
}
|
||||
"add_columns" => {
|
||||
let _ = table
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + 1".into())]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"c".into(),
|
||||
"a + 1".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.await;
|
||||
}
|
||||
"drop_columns" => {
|
||||
@@ -9903,10 +9883,12 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
branch
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"b".into(),
|
||||
"a + 1".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
branch
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -33,7 +33,7 @@ use crate::table::{AddResult, MergeResult};
|
||||
/// same Arrow-IPC streaming body and error side-channel; only the target
|
||||
/// endpoint, query parameters, and parsed result type differ.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum WriteOp {
|
||||
pub(crate) enum WriteOp {
|
||||
/// `add`: stream to `/v1/table/{id}/insert/`, optionally overwriting.
|
||||
Insert { overwrite: bool },
|
||||
/// `merge_insert`: stream to `/v1/table/{id}/merge_insert/` with the merge
|
||||
@@ -49,7 +49,7 @@ pub enum WriteOp {
|
||||
/// The parsed server response for a completed write, discriminated by the
|
||||
/// operation that produced it.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum WriteResult {
|
||||
pub(crate) enum WriteResult {
|
||||
Add(AddResult),
|
||||
Merge(MergeResult),
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
//! LanceDB Table APIs
|
||||
|
||||
use crate::blob::BlobFile;
|
||||
use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
|
||||
use arrow_schema::{Schema, SchemaRef};
|
||||
use async_trait::async_trait;
|
||||
@@ -12,7 +13,6 @@ use datafusion_physical_plan::ExecutionPlan;
|
||||
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
||||
use futures::StreamExt;
|
||||
use futures::stream::FuturesUnordered;
|
||||
use lance::dataset::BlobFile;
|
||||
pub use lance::dataset::ColumnAlteration;
|
||||
pub use lance::dataset::NewColumnTransform;
|
||||
pub use lance::dataset::ReadParams;
|
||||
@@ -65,6 +65,7 @@ use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
|
||||
use self::dataset::DatasetConsistencyWrapper;
|
||||
use self::merge::MergeInsertBuilder;
|
||||
|
||||
pub mod add_columns;
|
||||
mod add_data;
|
||||
pub mod branch_merge;
|
||||
mod create_index;
|
||||
@@ -79,6 +80,7 @@ pub mod schema_evolution;
|
||||
pub mod update;
|
||||
pub mod write_progress;
|
||||
use crate::index::waiter::wait_for_index;
|
||||
pub use add_columns::AddColumnsBuilder;
|
||||
#[cfg(feature = "remote")]
|
||||
pub(crate) use add_data::PreprocessingOutput;
|
||||
pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior};
|
||||
@@ -1620,12 +1622,8 @@ impl Table {
|
||||
}
|
||||
|
||||
/// Add new columns to the table, providing values to fill in.
|
||||
pub async fn add_columns(
|
||||
&self,
|
||||
transforms: NewColumnTransform,
|
||||
read_columns: Option<Vec<String>>,
|
||||
) -> Result<AddColumnsResult> {
|
||||
self.inner.add_columns(transforms, read_columns).await
|
||||
pub fn add_columns(&self) -> AddColumnsBuilder {
|
||||
AddColumnsBuilder::new(self.inner.clone())
|
||||
}
|
||||
|
||||
/// Change a column's name or nullability.
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Builder for adding columns to a table.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use lance::dataset::NewColumnTransform;
|
||||
|
||||
use super::BaseTable;
|
||||
use super::schema_evolution::AddColumnsResult;
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Adds columns to a table. See [`Table::add_columns`](super::Table::add_columns).
|
||||
pub struct AddColumnsBuilder {
|
||||
parent: Arc<dyn BaseTable>,
|
||||
transform: Option<NewColumnTransform>,
|
||||
read_columns: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AddColumnsBuilder {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AddColumnsBuilder")
|
||||
.field("parent", &self.parent)
|
||||
.field("has_transform", &self.transform.is_some())
|
||||
.field("read_columns", &self.read_columns)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl AddColumnsBuilder {
|
||||
pub(crate) fn new(parent: Arc<dyn BaseTable>) -> Self {
|
||||
Self {
|
||||
parent,
|
||||
transform: None,
|
||||
read_columns: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set how the new columns' values are produced. Required.
|
||||
pub fn transform(mut self, transform: NewColumnTransform) -> Self {
|
||||
self.transform = Some(transform);
|
||||
self
|
||||
}
|
||||
|
||||
/// Limit which existing columns a [`NewColumnTransform::BatchUDF`] mapper
|
||||
/// receives. Every other transform determines what it reads, so setting
|
||||
/// this alongside one is an error rather than a silent no-op.
|
||||
pub fn read_columns(mut self, columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
|
||||
self.read_columns = Some(columns.into_iter().map(Into::into).collect());
|
||||
self
|
||||
}
|
||||
|
||||
/// Add the columns.
|
||||
pub async fn execute(self) -> Result<AddColumnsResult> {
|
||||
let Self {
|
||||
parent,
|
||||
transform,
|
||||
read_columns,
|
||||
} = self;
|
||||
|
||||
let Some(transform) = transform else {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "add_columns requires a transform".into(),
|
||||
});
|
||||
};
|
||||
|
||||
if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "read_columns applies only to a BatchUDF transform; \
|
||||
every other transform determines what it reads"
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
|
||||
parent.add_columns(transform, read_columns).await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{Int32Array, RecordBatch, record_batch};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use lance::dataset::{BatchUDF, NewColumnTransform};
|
||||
|
||||
use crate::Table;
|
||||
use crate::connect;
|
||||
|
||||
async fn table_with_two_columns(name: &str) -> Table {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let batch = record_batch!(("x", Int32, [1, 2, 3]), ("y", Int32, [10, 20, 30])).unwrap();
|
||||
conn.create_table(name, batch).execute().await.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_requires_a_transform() {
|
||||
let table = table_with_two_columns("no_transform").await;
|
||||
let err = table.add_columns().execute().await.unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("requires a transform"),
|
||||
"got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_columns_with_sql_expressions_is_rejected() {
|
||||
let table = table_with_two_columns("read_cols_sql").await;
|
||||
let err = table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"doubled".into(),
|
||||
"x * 2".into(),
|
||||
)]))
|
||||
.read_columns(["x"])
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("BatchUDF"), "got: {err}");
|
||||
|
||||
let schema = table.schema().await.unwrap();
|
||||
assert!(
|
||||
schema.field_with_name("doubled").is_err(),
|
||||
"a rejected call must not commit"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_columns_limits_what_a_batch_udf_sees() {
|
||||
let table = table_with_two_columns("read_cols_udf").await;
|
||||
|
||||
let output_schema = Arc::new(Schema::new(vec![Field::new("sum", DataType::Int32, true)]));
|
||||
let mapper_schema = output_schema.clone();
|
||||
let udf = BatchUDF {
|
||||
mapper: Box::new(move |batch: &RecordBatch| {
|
||||
assert!(batch.column_by_name("x").is_some());
|
||||
assert!(batch.column_by_name("y").is_none(), "y was not requested");
|
||||
let x = batch["x"].as_any().downcast_ref::<Int32Array>().unwrap();
|
||||
let doubled: Int32Array = x.iter().map(|v| v.map(|v| v * 2)).collect();
|
||||
Ok(RecordBatch::try_new(
|
||||
mapper_schema.clone(),
|
||||
vec![Arc::new(doubled)],
|
||||
)?)
|
||||
}),
|
||||
output_schema,
|
||||
result_checkpoint: None,
|
||||
};
|
||||
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::BatchUDF(udf))
|
||||
.read_columns(["x"])
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let schema = table.schema().await.unwrap();
|
||||
assert!(schema.field_with_name("sum").is_ok());
|
||||
}
|
||||
}
|
||||
@@ -576,10 +576,12 @@ mod tests {
|
||||
|
||||
// Add a new physical column AFTER the embedding column.
|
||||
table
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("score".into(), "42.0".into())]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"score".into(),
|
||||
"42.0".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -683,7 +685,9 @@ mod tests {
|
||||
true,
|
||||
)]));
|
||||
table
|
||||
.add_columns(NewColumnTransform::AllNulls(nested_schema), None)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::AllNulls(nested_schema))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
||||
@@ -18,7 +18,6 @@
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use arrow_array::cast::AsArray;
|
||||
|
||||
@@ -84,8 +84,9 @@ pub(super) async fn create_lsm_plan(
|
||||
let pk_columns = pk_columns(&ds_ref)?;
|
||||
// The base index an indexed arm relies on may lag compaction; resolve it so the
|
||||
// snapshot retains SSTables the index has not yet caught up to.
|
||||
let arm_indexes = arm_maintained_index_names(&ds_ref, &query, &details).await?;
|
||||
let (snapshots, in_memory) = build_read_context(table, &ds_ref, &details, &arm_indexes).await?;
|
||||
let arm_index = arm_maintained_index_name(&ds_ref, &query, &details).await?;
|
||||
let (snapshots, in_memory) =
|
||||
build_read_context(table, &ds_ref, &details, arm_index.as_deref()).await?;
|
||||
|
||||
let limit = query.base.limit;
|
||||
let offset = query.base.offset;
|
||||
@@ -231,40 +232,28 @@ fn pk_columns(dataset: &Dataset) -> Result<Vec<String>> {
|
||||
Ok(pk)
|
||||
}
|
||||
|
||||
/// Per-shard SSTable exclusion watermark: the generation at or below which
|
||||
/// SSTables are safe to drop for this query.
|
||||
///
|
||||
/// A generation is droppable only once it is compacted into the base table AND
|
||||
/// covered by the catch-up of every index the query relies on, so the watermark
|
||||
/// is the minimum across `index_names`. Gating on fewer than all of them would
|
||||
/// drop SSTables holding rows an uncounted index has not yet indexed, and that
|
||||
/// arm would silently return fewer rows.
|
||||
///
|
||||
/// See [`arm_maintained_index_names`] for which indexes are collected today: a
|
||||
/// vector search with a scalar prefilter is not yet among them.
|
||||
///
|
||||
/// An empty `index_names` (a plain scan) uses the compaction watermark alone.
|
||||
/// First occurrence per shard mirrors Lance's `compacted_generation_for_shard`.
|
||||
/// Per-shard SSTable exclusion watermark: the generation at or below which SSTables
|
||||
/// are safe to drop for this arm. A generation is droppable only once it is
|
||||
/// compacted into the base table AND covered by `index_name`'s catch-up (for an
|
||||
/// indexed arm); a plain scan (`index_name == None`) uses the compaction watermark
|
||||
/// alone. Capping at the index catch-up keeps rows the base index has not yet
|
||||
/// indexed visible through their SSTable. First occurrence per shard mirrors Lance's
|
||||
/// `compacted_generation_for_shard`.
|
||||
fn exclusion_watermarks(
|
||||
details: &MemWalIndexDetails,
|
||||
index_names: &[String],
|
||||
index_name: Option<&str>,
|
||||
) -> HashMap<Uuid, u64> {
|
||||
let mut exclude: HashMap<Uuid, u64> = HashMap::new();
|
||||
for entry in &details.compacted_sstables {
|
||||
let mut watermark = entry.generation;
|
||||
for name in index_names {
|
||||
match details
|
||||
if let Some(name) = index_name
|
||||
&& let Some(caught_up) = details
|
||||
.index_catchup
|
||||
.iter()
|
||||
.find(|icp| icp.index_name == *name)
|
||||
.find(|icp| icp.index_name == name)
|
||||
.and_then(|icp| icp.caught_up_generation_for_shard(&entry.shard_id))
|
||||
{
|
||||
Some(caught_up) => watermark = watermark.min(caught_up),
|
||||
// No entry means the index is *not* known to hold these rows,
|
||||
// and the base arm is index-only -- so every generation stays
|
||||
// readable from its SSTable.
|
||||
None => watermark = 0,
|
||||
}
|
||||
{
|
||||
watermark = watermark.min(caught_up);
|
||||
}
|
||||
exclude.entry(entry.shard_id).or_insert(watermark);
|
||||
}
|
||||
@@ -282,9 +271,9 @@ async fn build_read_context(
|
||||
table: &NativeTable,
|
||||
dataset: &Dataset,
|
||||
details: &MemWalIndexDetails,
|
||||
index_names: &[String],
|
||||
index_name: Option<&str>,
|
||||
) -> Result<(Vec<ShardSnapshot>, HashMap<Uuid, InMemoryMemTables>)> {
|
||||
let exclude = exclusion_watermarks(details, index_names);
|
||||
let exclude = exclusion_watermarks(details, index_name);
|
||||
|
||||
let shard_ids = dataset.list_mem_wal_latest_shard_ids().await?;
|
||||
// Use the dataset's own object store (not `ObjectStore::from_uri`, which
|
||||
@@ -498,33 +487,19 @@ async fn index_maintained(
|
||||
}))
|
||||
}
|
||||
|
||||
/// Every maintained base index this query relies on, used to gate SSTable
|
||||
/// exclusion by index catch-up.
|
||||
///
|
||||
/// Returns a list because the watermark must be the lowest across every index a
|
||||
/// query relies on. Today it never holds more than one: `reject_unsupported`
|
||||
/// refuses hybrid search, so the vector and full-text arms are mutually
|
||||
/// exclusive.
|
||||
///
|
||||
/// The case that is genuinely multi-index -- a vector search with a scalar or
|
||||
/// bitmap prefilter -- is **not collected yet**. Identifying those needs the
|
||||
/// planner's chosen indexes, not the columns the filter names, and no Lance API
|
||||
/// exposes them. Until it does, such a query is gated on its vector index alone.
|
||||
///
|
||||
/// Empty for a plain scan, or when no maintained index covers the searched
|
||||
/// column.
|
||||
async fn arm_maintained_index_names(
|
||||
/// The maintained base index the query's arm relies on (vector index for ANN, FTS
|
||||
/// index for full-text), used to gate SSTable compaction exclusion by index catch-up.
|
||||
/// `None` for a plain scan or when no maintained index covers the searched column.
|
||||
async fn arm_maintained_index_name(
|
||||
dataset: &Dataset,
|
||||
query: &VectorQueryRequest,
|
||||
details: &MemWalIndexDetails,
|
||||
) -> Result<Vec<String>> {
|
||||
) -> Result<Option<String>> {
|
||||
use lance::index::DatasetIndexExt;
|
||||
|
||||
// Each arm's searched column, the index-detail type it relies on, and a
|
||||
// Resolve the arm's searched column, the index-detail type it relies on, and a
|
||||
// label for diagnostics — catch-up is taken from the vector/FTS index
|
||||
// specifically, not a BTree on the same column.
|
||||
let mut arms: Vec<(String, &str, &str)> = Vec::new();
|
||||
if !query.query_vector.is_empty() {
|
||||
let (column, type_url_suffix, arm) = if !query.query_vector.is_empty() {
|
||||
let arrow_schema = ArrowSchema::from(dataset.schema());
|
||||
let column = match &query.column {
|
||||
Some(column) => column.clone(),
|
||||
@@ -533,43 +508,31 @@ async fn arm_maintained_index_names(
|
||||
default_vector_column(&arrow_schema, dim)?
|
||||
}
|
||||
};
|
||||
arms.push((column, "VectorIndexDetails", "vector"));
|
||||
}
|
||||
if let Some(fts) = &query.base.full_text_search
|
||||
&& let Some(column) = fts.columns().into_iter().next()
|
||||
{
|
||||
arms.push((column, "InvertedIndexDetails", "full-text"));
|
||||
}
|
||||
if arms.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let indices = dataset.load_indices().await?;
|
||||
let mut names = Vec::with_capacity(arms.len());
|
||||
for (column, type_url_suffix, arm) in arms {
|
||||
let Some(field) = dataset.schema().field(&column) else {
|
||||
continue;
|
||||
};
|
||||
let segment_names: Vec<String> = indices
|
||||
.iter()
|
||||
.filter(|idx| {
|
||||
idx.fields.contains(&field.id)
|
||||
&& idx
|
||||
.index_details
|
||||
.as_ref()
|
||||
.is_some_and(|d| d.type_url.ends_with(type_url_suffix))
|
||||
})
|
||||
.map(|idx| idx.name.clone())
|
||||
.collect();
|
||||
if let Some(name) =
|
||||
resolve_single_index(segment_names, &details.maintained_indexes, arm, &column)?
|
||||
{
|
||||
names.push(name);
|
||||
(column, "VectorIndexDetails", "vector")
|
||||
} else if let Some(fts) = &query.base.full_text_search {
|
||||
match fts.columns().into_iter().next() {
|
||||
Some(column) => (column, "InvertedIndexDetails", "full-text"),
|
||||
None => return Ok(None),
|
||||
}
|
||||
}
|
||||
names.sort();
|
||||
names.dedup();
|
||||
Ok(names)
|
||||
} else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(field) = dataset.schema().field(&column) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let indices = dataset.load_indices().await?;
|
||||
let segment_names: Vec<String> = indices
|
||||
.iter()
|
||||
.filter(|idx| {
|
||||
idx.fields.contains(&field.id)
|
||||
&& idx
|
||||
.index_details
|
||||
.as_ref()
|
||||
.is_some_and(|d| d.type_url.ends_with(type_url_suffix))
|
||||
})
|
||||
.map(|idx| idx.name.clone())
|
||||
.collect();
|
||||
resolve_single_index(segment_names, &details.maintained_indexes, arm, &column)
|
||||
}
|
||||
|
||||
/// Resolve the single logical index from the names of its matching physical
|
||||
@@ -771,99 +734,24 @@ mod tests {
|
||||
};
|
||||
|
||||
// Plain scan: drop every compacted generation (through 5).
|
||||
assert_eq!(exclusion_watermarks(&details, &[]).get(&shard), Some(&5));
|
||||
assert_eq!(exclusion_watermarks(&details, None).get(&shard), Some(&5));
|
||||
|
||||
// FTS arm with a lagging index: exclusion is capped at the index catch-up
|
||||
// (2), so SSTable generations 3..=5 are retained until the index covers
|
||||
// them — otherwise those documents would silently vanish from FTS results.
|
||||
assert_eq!(
|
||||
exclusion_watermarks(&details, &["fts_idx".to_string()]).get(&shard),
|
||||
exclusion_watermarks(&details, Some("fts_idx")).get(&shard),
|
||||
Some(&2)
|
||||
);
|
||||
|
||||
// An index with no entry has not recorded that it holds these rows, so
|
||||
// nothing is excluded. This is the case a table written before catch-up
|
||||
// was maintained lands in, and it errs toward reading the SSTables.
|
||||
// A caught-up index — or one untracked in index_catchup — falls back to the
|
||||
// compaction watermark.
|
||||
assert_eq!(
|
||||
exclusion_watermarks(&details, &["untracked_idx".to_string()]).get(&shard),
|
||||
Some(&0)
|
||||
);
|
||||
|
||||
// An index recorded as covering the compaction watermark excludes up to it.
|
||||
let caught_up = MemWalIndexDetails {
|
||||
index_catchup: vec![IndexCatchupProgress::new(
|
||||
"caught_up_idx".to_string(),
|
||||
vec![CompactedSsTable::new(shard, 5)],
|
||||
)],
|
||||
..details.clone()
|
||||
};
|
||||
assert_eq!(
|
||||
exclusion_watermarks(&caught_up, &["caught_up_idx".to_string()]).get(&shard),
|
||||
exclusion_watermarks(&details, Some("caught_up_idx")).get(&shard),
|
||||
Some(&5)
|
||||
);
|
||||
}
|
||||
|
||||
/// A hybrid search reads a vector and a full-text index, and either may lag.
|
||||
/// Retaining to the lower of the two is what keeps both arms complete;
|
||||
/// gating on one alone would drop SSTables the other has not indexed.
|
||||
#[test]
|
||||
fn exclusion_watermark_takes_the_minimum_across_every_index_used() {
|
||||
let shard = Uuid::from_u128(1);
|
||||
let details = MemWalIndexDetails {
|
||||
compacted_sstables: vec![CompactedSsTable::new(shard, 9)],
|
||||
index_catchup: vec![
|
||||
IndexCatchupProgress::new(
|
||||
"vec_idx".to_string(),
|
||||
vec![CompactedSsTable::new(shard, 7)],
|
||||
),
|
||||
IndexCatchupProgress::new(
|
||||
"fts_idx".to_string(),
|
||||
vec![CompactedSsTable::new(shard, 4)],
|
||||
),
|
||||
],
|
||||
maintained_indexes: vec!["vec_idx".to_string(), "fts_idx".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Each index alone stops at its own catch-up.
|
||||
assert_eq!(
|
||||
exclusion_watermarks(&details, &["vec_idx".to_string()]).get(&shard),
|
||||
Some(&7)
|
||||
);
|
||||
assert_eq!(
|
||||
exclusion_watermarks(&details, &["fts_idx".to_string()]).get(&shard),
|
||||
Some(&4)
|
||||
);
|
||||
|
||||
// Used together, the lower one governs regardless of order.
|
||||
let both = ["vec_idx".to_string(), "fts_idx".to_string()];
|
||||
assert_eq!(exclusion_watermarks(&details, &both).get(&shard), Some(&4));
|
||||
let reversed = ["fts_idx".to_string(), "vec_idx".to_string()];
|
||||
assert_eq!(
|
||||
exclusion_watermarks(&details, &reversed).get(&shard),
|
||||
Some(&4)
|
||||
);
|
||||
}
|
||||
|
||||
/// An index with no catch-up entry is not known to hold anything, so it
|
||||
/// governs over a lagging sibling rather than the other way round.
|
||||
#[test]
|
||||
fn an_untracked_index_retains_everything() {
|
||||
let shard = Uuid::from_u128(1);
|
||||
let details = MemWalIndexDetails {
|
||||
compacted_sstables: vec![CompactedSsTable::new(shard, 9)],
|
||||
index_catchup: vec![IndexCatchupProgress::new(
|
||||
"fts_idx".to_string(),
|
||||
vec![CompactedSsTable::new(shard, 4)],
|
||||
)],
|
||||
maintained_indexes: vec!["fts_idx".to_string(), "untracked_idx".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let both = ["fts_idx".to_string(), "untracked_idx".to_string()];
|
||||
assert_eq!(exclusion_watermarks(&details, &both).get(&shard), Some(&0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_single_index_dedupes_segments() {
|
||||
let maintained = vec!["fts_idx".to_string()];
|
||||
|
||||
@@ -193,10 +193,12 @@ mod tests {
|
||||
|
||||
// Add a computed column
|
||||
let result = table
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("doubled".into(), "id * 2".into())]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"doubled".into(),
|
||||
"id * 2".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -251,13 +253,12 @@ mod tests {
|
||||
|
||||
// Add multiple columns at once
|
||||
table
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![
|
||||
("y".into(), "x + 1".into()),
|
||||
("z".into(), "x * x".into()),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![
|
||||
("y".into(), "x + 1".into()),
|
||||
("z".into(), "x * x".into()),
|
||||
]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -283,10 +284,12 @@ mod tests {
|
||||
|
||||
// Add a column with a constant value
|
||||
table
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("constant".into(), "42".into())]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"constant".into(),
|
||||
"42".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -659,10 +662,12 @@ mod tests {
|
||||
|
||||
// Add column increments version
|
||||
let add_result = table
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + b".into())]),
|
||||
None,
|
||||
)
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"c".into(),
|
||||
"a + b".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(add_result.version > v1);
|
||||
|
||||
@@ -9,14 +9,17 @@ use arrow_array::{
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||
use futures::TryStreamExt;
|
||||
use lance::Dataset;
|
||||
use lance_encoding::version::LanceFileVersion;
|
||||
use lancedb::{
|
||||
Connection, Error, Result, Table,
|
||||
blob::{BlobRangeRequest, blob},
|
||||
connect, connect_namespace,
|
||||
database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||
database::listing::{
|
||||
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||
},
|
||||
query::{ExecutableQuery, QueryBase},
|
||||
table::{AddDataMode, CompactionOptions, OptimizeAction},
|
||||
table::{AddDataMode, CompactionOptions, OptimizeAction, OptimizeStats},
|
||||
};
|
||||
use tempfile::tempdir;
|
||||
|
||||
@@ -1075,3 +1078,252 @@ async fn fetch_blob_files_aligns_across_fragments_with_nulls_and_dups() -> Resul
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Rows exercising the null/empty interleavings from
|
||||
/// <https://github.com/lancedb/lancedb/issues/3744>: a payload, a null, a valid
|
||||
/// empty value, then payloads whose descriptors a fragment rewrite used to zero.
|
||||
fn null_empty_input_batch() -> RecordBatch {
|
||||
let owned = [
|
||||
Some(dedicated_blob_bytes(1)),
|
||||
None,
|
||||
Some(Vec::new()),
|
||||
Some(dedicated_blob_bytes(4)),
|
||||
Some(dedicated_blob_bytes(5)),
|
||||
Some(dedicated_blob_bytes(6)),
|
||||
];
|
||||
let payloads: Vec<Option<&[u8]>> = owned.iter().map(|payload| payload.as_deref()).collect();
|
||||
binary_input_batch(&[1, 2, 3, 4, 5, 6], &payloads)
|
||||
}
|
||||
|
||||
/// One `(id, Some((payload length, first byte)))` per live row, or `(id, None)`
|
||||
/// for a null blob. Comparing lengths and first bytes keeps failure output
|
||||
/// readable where comparing whole payloads would not.
|
||||
type BlobSummary = Vec<(i64, Option<(usize, Option<u8>)>)>;
|
||||
|
||||
/// The rows [`null_empty_input_batch`] leaves behind after `id IN (1, 4)` is
|
||||
/// deleted: a null, a valid empty value, and the two payloads that follow them.
|
||||
fn expected_null_empty_survivors() -> BlobSummary {
|
||||
vec![
|
||||
(2, None),
|
||||
(3, Some((0, None))),
|
||||
(5, Some((DEDICATED_BLOB_LEN, Some(5)))),
|
||||
(6, Some((DEDICATED_BLOB_LEN, Some(6)))),
|
||||
]
|
||||
}
|
||||
|
||||
/// `optimize()` only rewrites a fragment when lance's compaction planner selects
|
||||
/// it — here because the delete pushes the fragment past
|
||||
/// `materialize_deletions_threshold` (0.1 by default; these tests delete 2 of 6
|
||||
/// rows). Without this check, a planner or threshold change upstream would leave
|
||||
/// both regression tests green while no rewrite happened at all.
|
||||
fn assert_compacted(stats: &OptimizeStats) {
|
||||
let metrics = stats
|
||||
.compaction
|
||||
.as_ref()
|
||||
.expect("OptimizeAction::All runs compaction");
|
||||
assert!(
|
||||
metrics.fragments_removed >= 1,
|
||||
"optimize() rewrote no fragment, so this test proves nothing: {metrics:?}"
|
||||
);
|
||||
}
|
||||
|
||||
fn summarize(rows: &[(i64, Option<Vec<u8>>)]) -> BlobSummary {
|
||||
rows.iter()
|
||||
.map(|(id, payload)| {
|
||||
(
|
||||
*id,
|
||||
payload
|
||||
.as_ref()
|
||||
.map(|bytes| (bytes.len(), bytes.first().copied())),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn sorted_id_rowid(table: &Table) -> Result<Vec<(i64, u64)>> {
|
||||
let mut pairs = collect_id_rowid(table).await?;
|
||||
pairs.sort_by_key(|(id, _)| *id);
|
||||
Ok(pairs)
|
||||
}
|
||||
|
||||
/// `{position, size}` descriptors of a legacy v1 blob column, keyed by `id`.
|
||||
async fn v1_blob_descriptors(table: &Table) -> Result<Vec<(i64, Option<(u64, u64)>)>> {
|
||||
let batches = table
|
||||
.query()
|
||||
.execute()
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
let batch = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap();
|
||||
let ids = batch
|
||||
.column_by_name("id")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.unwrap();
|
||||
let descriptors = batch
|
||||
.column_by_name("image")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<StructArray>()
|
||||
.expect("v1 blob column reads back as a descriptor struct");
|
||||
let position = descriptors
|
||||
.column_by_name("position")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap();
|
||||
let size = descriptors
|
||||
.column_by_name("size")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap();
|
||||
let mut rows: Vec<(i64, Option<(u64, u64)>)> = (0..batch.num_rows())
|
||||
.map(|row| {
|
||||
let descriptor =
|
||||
(!descriptors.is_null(row)).then(|| (position.value(row), size.value(row)));
|
||||
(ids.value(row), descriptor)
|
||||
})
|
||||
.collect();
|
||||
rows.sort_by_key(|(id, _)| *id);
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Payload bytes of every live row of a legacy v1 blob column, keyed by `id`.
|
||||
/// [`Table::fetch_blobs`] rejects v1 columns, so read them through lance.
|
||||
async fn v1_blob_payloads(dataset_uri: &str, table: &Table) -> Result<Vec<(i64, Option<Vec<u8>>)>> {
|
||||
let pairs = sorted_id_rowid(table).await?;
|
||||
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
|
||||
let dataset = Arc::new(Dataset::open(dataset_uri).await?);
|
||||
let files = dataset.take_blobs(&row_ids, "image").await?;
|
||||
assert_eq!(
|
||||
files.len(),
|
||||
pairs.len(),
|
||||
"take_blobs returned {} handles for {} live rows",
|
||||
files.len(),
|
||||
pairs.len()
|
||||
);
|
||||
let mut rows = Vec::with_capacity(pairs.len());
|
||||
for ((id, _), file) in pairs.iter().zip(files) {
|
||||
let payload = match file {
|
||||
Some(file) => Some(file.read().await?.to_vec()),
|
||||
None => None,
|
||||
};
|
||||
rows.push((*id, payload));
|
||||
}
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Length and first byte of every live blob v2 value, keyed by `id`.
|
||||
async fn blob_v2_values(table: &Table) -> Result<BlobSummary> {
|
||||
let pairs = sorted_id_rowid(table).await?;
|
||||
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
|
||||
let bytes = table.fetch_blobs("image", &row_ids).await?;
|
||||
Ok(pairs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(slot, (id, _))| {
|
||||
let value = (!bytes.is_null(slot))
|
||||
.then(|| (bytes.value(slot).len(), bytes.value(slot).first().copied()));
|
||||
(*id, value)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Regression test for [#3744]: on storage 2.0 (legacy v1 descriptors),
|
||||
/// compaction rewrote every payload following a null or empty value in the same
|
||||
/// fragment as `{position: 0, size: 0}`, so the payload bytes read back as `b""`
|
||||
/// and the new fragment no longer referenced them at all.
|
||||
///
|
||||
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
|
||||
#[tokio::test]
|
||||
async fn optimize_preserves_v1_blob_payloads_with_null_and_empty() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db_uri = tmp.path().to_str().unwrap().to_string();
|
||||
let db = connect(&db_uri)
|
||||
.database_options(&ListingDatabaseOptions {
|
||||
new_table_config: NewTableConfig {
|
||||
data_storage_version: Some(LanceFileVersion::V2_0),
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
})
|
||||
.execute()
|
||||
.await?;
|
||||
let legacy = Field::new("image", DataType::LargeBinary, true).with_metadata(
|
||||
std::collections::HashMap::from([("lance-encoding:blob".to_string(), "true".to_string())]),
|
||||
);
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
legacy,
|
||||
]));
|
||||
let table = db.create_empty_table("t", schema).execute().await?;
|
||||
table.add(null_empty_input_batch()).execute().await?;
|
||||
assert_eq!(
|
||||
storage_format_version(&table).await,
|
||||
LanceFileVersion::V2_0.resolve(),
|
||||
"v1 blob descriptors only exist below storage 2.2"
|
||||
);
|
||||
let dataset_uri = table.uri().await?;
|
||||
|
||||
// Any rewrite triggers it; deleting rows is the shape from the issue.
|
||||
table.delete("id IN (1, 4)").await?;
|
||||
let descriptors_before = v1_blob_descriptors(&table).await?;
|
||||
let before = v1_blob_payloads(&dataset_uri, &table).await?;
|
||||
assert_eq!(
|
||||
summarize(&before),
|
||||
expected_null_empty_survivors(),
|
||||
"test setup no longer produces the null/empty/payload mix"
|
||||
);
|
||||
|
||||
let stats = table.optimize(OptimizeAction::All).await?;
|
||||
assert_compacted(&stats);
|
||||
|
||||
let descriptors_after = v1_blob_descriptors(&table).await?;
|
||||
let after = v1_blob_payloads(&dataset_uri, &table).await?;
|
||||
assert_eq!(
|
||||
summarize(&after),
|
||||
summarize(&before),
|
||||
"optimize() lost blob payloads; descriptors before={descriptors_before:?} after={descriptors_after:?}"
|
||||
);
|
||||
assert!(after == before, "optimize() changed blob payload bytes");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Regression test for the blob v2 half of [#3744]: compaction rewrote a valid
|
||||
/// empty value as null, destroying the null-vs-empty distinction.
|
||||
///
|
||||
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
|
||||
#[tokio::test]
|
||||
async fn optimize_preserves_blob_v2_null_and_empty_distinction() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
|
||||
let table = db
|
||||
.create_empty_table("t", blob_table_schema())
|
||||
.execute()
|
||||
.await?;
|
||||
table.add(null_empty_input_batch()).execute().await?;
|
||||
assert!(
|
||||
storage_format_version(&table).await >= LanceFileVersion::V2_2,
|
||||
"blob v2 columns require storage >= 2.2"
|
||||
);
|
||||
|
||||
table.delete("id IN (1, 4)").await?;
|
||||
let before = blob_v2_values(&table).await?;
|
||||
assert_eq!(
|
||||
before,
|
||||
expected_null_empty_survivors(),
|
||||
"test setup no longer produces the null/empty/payload mix"
|
||||
);
|
||||
|
||||
let stats = table.optimize(OptimizeAction::All).await?;
|
||||
assert_compacted(&stats);
|
||||
|
||||
assert_eq!(
|
||||
blob_v2_values(&table).await?,
|
||||
before,
|
||||
"optimize() changed blob v2 values"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user