use crate::{
RetryPolicy,
config::{self, CredsSet},
handlers::NamespaceIdent,
sessions::{self},
};
use anyhow::{Context, Result, anyhow, bail};
use lance::Dataset;
use lance::dataset::builder::DatasetBuilder;
use lance::dataset::index::DatasetIndexRemapperOptions;
use lance::dataset::optimize::{
CompactionMode, CompactionOptions, commit_compaction, plan_compaction,
};
pub use lance::dataset::write::merge_insert::MergeStats;
use lance::dataset::write::merge_insert::SourceDedupeBehavior;
use lance::dataset::{InsertBuilder, MergeInsertBuilder, WhenMatched, WhenNotMatched, WriteMode};
pub use lance::dataset::{WriteParams, WriteStats};
use lance::deps::arrow_array::{Array, RecordBatch, RecordBatchIterator, StringArray};
use lance::deps::datafusion::physical_plan::SendableRecordBatchStream;
use lance::index::DatasetIndexExt;
use lance::index::DatasetIndexInternalExt;
use lance::index::vector::VectorIndexParams;
use lance::session::Session;
use lance_index::IndexType;
use lance_index::optimize::OptimizeOptions;
use lance_index::scalar::{BuiltinIndexType, InvertedIndexParams, ScalarIndexParams};
use lance_index::vector::ivf::IvfBuildParams;
use lance_index::vector::sq::builder::SQBuildParams;
use lance_io::object_store::{
ChainedWrappingObjectStore, ObjectStore, ObjectStoreParams, ObjectStoreRegistry,
StorageOptionsAccessor, WrappingObjectStore, uri_to_url,
};
use lance_linalg::distance::MetricType;
use lance_namespace::LanceNamespace;
use lance_namespace::error::{ErrorCode, NamespaceError};
use lance_namespace::models::DescribeTableRequest;
use lance_namespace_impls::ConnectBuilder;
use std::{
collections::{BTreeMap, HashMap},
path::PathBuf,
sync::Arc,
time::{Duration, Instant},
};
use tokio::sync::{Mutex, OnceCell};
use tokio_stream::StreamExt;
use url::Url;
pub const VECTOR_INDEX_ACTIVATION_ROWS: usize = 100_000;
pub const DELTA_MERGE_THRESHOLD: usize = 4;
#[derive(Debug, Clone, PartialEq)]
pub struct StorageUrl {
canonical: Url,
lance: Url,
scheme_options: Vec<(&'static str, String)>,
query_options: Vec<(&'static str, String)>,
creds_pointer: Option<String>,
endpoint: Option<S3Endpoint>,
}
#[derive(Debug, Clone, PartialEq)]
struct S3Endpoint {
scheme: &'static str,
authority: String,
bucket: String,
}
const RECOGNIZED_QUERY_PARAMS: [&str; 3] = ["creds", "region", "virtual_hosted_style_request"];
fn network_store_rejected(input: &str) -> anyhow::Error {
anyhow!(
"storage path {input:?} is a Windows network path; the storage stack drops the host, so \
the store would silently open elsewhere. Use a path on a local drive, or a remote scheme \
(s3://, s3+https://, gs://, az://). Mapped network drives have the same fault and cannot \
be detected here."
)
}
impl StorageUrl {
pub fn parse(input: &str) -> Result<Self> {
let trimmed = input.trim();
if trimmed.is_empty() {
bail!("storage path is empty");
}
if !trimmed.contains("://") || trimmed.starts_with("file://") {
if trimmed
.get(..8)
.is_some_and(|prefix| prefix.eq_ignore_ascii_case(r"\\?\UNC\"))
{
return Err(network_store_rejected(trimmed));
}
if trimmed.starts_with(r"\\?\") {
bail!(
"storage path {trimmed:?} uses the extended-length `\\\\?\\` prefix, which \
pond does not accept: it is not understood consistently below pond, and it \
leaks into every path handed to another program. Use the ordinary form \
(`C:\\srv\\pond`)."
);
}
if trimmed.starts_with(r"\\") || (cfg!(windows) && trimmed.starts_with("//")) {
return Err(network_store_rejected(trimmed));
}
let url =
uri_to_url(trimmed).with_context(|| format!("invalid storage path {trimmed:?}"))?;
if url
.host_str()
.is_some_and(|host| !host.is_empty() && !host.eq_ignore_ascii_case("localhost"))
{
return Err(network_store_rejected(trimmed));
}
if url.query().is_some() {
bail!("storage URL {trimmed:?} carries query params; local URLs take none");
}
return Ok(Self::plain(url));
}
let url =
Url::parse(trimmed).with_context(|| format!("invalid storage URL {trimmed:?}"))?;
if !url.username().is_empty() || url.password().is_some() {
bail!(
"storage URL {trimmed:?} embeds credentials; put them in [creds.*] (or POND_CREDS_*) instead"
);
}
match url.scheme() {
"memory" | "shared-memory" => {
if url.query().is_some() {
bail!(
"storage URL {trimmed:?} carries query params; {}:// URLs take none",
url.scheme(),
);
}
Ok(Self::plain(url))
}
"s3" | "gs" => {
let (canonical, query_options, creds_pointer) = strip_query(url)?;
let mut lance = canonical.clone();
lance.set_query(None);
Ok(Self {
canonical,
lance,
scheme_options: Vec::new(),
query_options,
creds_pointer,
endpoint: None,
})
}
"s3+https" | "s3+http" => {
let (mut canonical, query_options, creds_pointer) = strip_query(url)?;
let tls = canonical.scheme() == "s3+https";
if canonical.port() == Some(if tls { 443 } else { 80 }) {
let _ = canonical.set_port(None);
}
let host = canonical
.host_str()
.ok_or_else(|| anyhow!("storage URL {trimmed:?} has no endpoint host"))?;
let endpoint_authority = match canonical.port() {
Some(port) => format!("{host}:{port}"),
None => host.to_owned(),
};
let mut segments = canonical.path().trim_start_matches('/').splitn(2, '/');
let bucket = segments.next().unwrap_or_default().to_owned();
let prefix = segments.next().unwrap_or_default().to_owned();
if bucket.is_empty() {
bail!(
"storage URL {trimmed:?} is missing the bucket: the form is {}://host/bucket/prefix",
canonical.scheme(),
);
}
let lance = Url::parse(&format!("s3://{bucket}/{prefix}")).with_context(|| {
format!("storage URL {trimmed:?}: bucket/prefix do not form a valid s3:// URL")
})?;
let scheme = if tls { "https" } else { "http" };
let virtual_hosted = host.parse::<std::net::IpAddr>().is_err()
&& !matches!(canonical.host(), Some(url::Host::Ipv6(_)));
let scheme_options = vec![
("allow_http", (!tls).to_string()),
("virtual_hosted_style_request", virtual_hosted.to_string()),
("region", "us-east-1".to_owned()),
];
Ok(Self {
canonical,
lance,
scheme_options,
query_options,
creds_pointer,
endpoint: Some(S3Endpoint {
scheme,
authority: endpoint_authority,
bucket,
}),
})
}
"az" => {
let (canonical, query_options, creds_pointer) = strip_query(url)?;
let account = canonical
.host_str()
.ok_or_else(|| anyhow!("storage URL {trimmed:?} has no account: the form is az://account/container/prefix"))?
.to_owned();
let mut segments = canonical.path().trim_start_matches('/').splitn(2, '/');
let container = segments.next().unwrap_or_default();
if container.is_empty() {
bail!(
"storage URL {trimmed:?} is missing the container: the form is az://account/container/prefix"
);
}
let prefix = segments.next().unwrap_or_default();
let lance = Url::parse(&format!("az://{container}/{prefix}"))
.with_context(|| format!("storage URL {trimmed:?}: container/prefix do not form a valid az:// URL"))?;
Ok(Self {
canonical,
lance,
scheme_options: vec![("account_name", account)],
query_options,
creds_pointer,
endpoint: None,
})
}
other => bail!(
"storage URL scheme {other:?} not recognized; use a local path, s3://, s3+https://, s3+http://, gs://, or az://"
),
}
}
fn plain(url: Url) -> Self {
Self {
canonical: url.clone(),
lance: url,
scheme_options: Vec::new(),
query_options: Vec::new(),
creds_pointer: None,
endpoint: None,
}
}
pub fn lance_url(&self) -> &Url {
&self.lance
}
pub fn canonical(&self) -> &Url {
&self.canonical
}
pub fn is_local(&self) -> bool {
config::is_local(&self.canonical)
}
pub fn display(&self) -> String {
config::display(&self.canonical)
}
fn takes_credentials(&self) -> bool {
!matches!(
self.canonical.scheme(),
"file" | "file+uring" | "memory" | "shared-memory"
)
}
pub fn resolve(&self, creds: &BTreeMap<String, CredsSet>) -> Result<ResolvedStorage> {
if !self.takes_credentials() {
return Ok(ResolvedStorage {
storage: self.clone(),
options: HashMap::new(),
binding: CredsBinding::NotApplicable,
});
}
let matched: Option<(&String, &CredsSet, BindVia)> = match &self.creds_pointer {
Some(name) => {
let set = creds.get(name).ok_or_else(|| {
anyhow!(
"URL names ?creds={name} but no [creds.{name}] set is configured; define it or drop the pointer"
)
})?;
Some((name, set, BindVia::Pointer))
}
None => {
let mut best: Option<(&String, &CredsSet, String)> = None;
for (name, set) in creds {
let Some(scope) = &set.scope else { continue };
let scope_url = parse_scope(scope).with_context(|| {
format!("[creds.{name}] scope {scope:?} is not a valid URL prefix")
})?;
if scope_matches(&scope_url, &self.canonical)
&& best
.as_ref()
.is_none_or(|(_, _, len)| scope_url.as_str().len() > len.len())
{
best = Some((name, set, scope_url.as_str().to_owned()));
}
}
match best {
Some((name, set, _)) => Some((name, set, BindVia::Scope)),
None => creds
.iter()
.find(|(_, set)| set.scope.is_none())
.map(|(name, set)| (name, set, BindVia::CatchAll)),
}
}
};
let mut options: HashMap<String, String> = self
.scheme_options
.iter()
.map(|(key, value)| ((*key).to_owned(), value.clone()))
.collect();
let binding = match matched {
None => CredsBinding::Ambient,
Some((name, set, via)) => {
if let Some(region) = &set.region {
options.insert("region".to_owned(), region.clone());
}
if let Some(virtual_hosted) = set.virtual_hosted_style_request {
options.insert(
"virtual_hosted_style_request".to_owned(),
virtual_hosted.to_string(),
);
}
for (key, value) in &set.extra {
options.insert(key.clone(), value.clone());
}
if let Some(value) = materialize_secret(
name,
"access_key_id",
set.access_key_id.as_deref(),
set.access_key_id_file.as_deref(),
None,
)? {
options.insert("access_key_id".to_owned(), value);
}
if let Some(value) = materialize_secret(
name,
"secret_access_key",
set.secret_access_key.as_deref(),
set.secret_access_key_file.as_deref(),
set.secret_access_key_command.as_deref(),
)? {
options.insert("secret_access_key".to_owned(), value);
}
CredsBinding::Set {
name: name.clone(),
via,
}
}
};
for (key, value) in &self.query_options {
options.insert((*key).to_owned(), value.clone());
}
if let Some(endpoint) = &self.endpoint
&& !options.keys().any(|key| {
key.eq_ignore_ascii_case("endpoint") || key.eq_ignore_ascii_case("aws_endpoint")
})
{
let virtual_hosted = options
.get("virtual_hosted_style_request")
.is_some_and(|value| value == "true");
let url = if virtual_hosted {
format!(
"{}://{}.{}",
endpoint.scheme, endpoint.bucket, endpoint.authority
)
} else {
format!("{}://{}", endpoint.scheme, endpoint.authority)
};
options.insert("endpoint".to_owned(), url);
}
Ok(ResolvedStorage {
storage: self.clone(),
options,
binding,
})
}
}
type StrippedQuery = (Url, Vec<(&'static str, String)>, Option<String>);
fn strip_query(url: Url) -> Result<StrippedQuery> {
let mut query_options = Vec::new();
let mut creds_pointer = None;
for (key, value) in url.query_pairs() {
match RECOGNIZED_QUERY_PARAMS
.iter()
.find(|known| **known == key.as_ref())
{
Some(&"creds") => creds_pointer = Some(value.into_owned()),
Some(known) => query_options.push((*known, value.into_owned())),
None => bail!(
"storage URL query param {key:?} not recognized (known: {})",
RECOGNIZED_QUERY_PARAMS.join(", "),
),
}
}
let mut canonical = url;
canonical.set_query(None);
Ok((canonical, query_options, creds_pointer))
}
pub(crate) fn parse_scope(scope: &str) -> Result<Url> {
let mut url = Url::parse(scope.trim())?;
if !url.username().is_empty() || url.password().is_some() {
bail!("scope embeds credentials");
}
if url.query().is_some() {
bail!("scope carries query params; scopes are plain URL prefixes");
}
match (url.scheme(), url.port()) {
("s3+https", Some(443)) | ("s3+http", Some(80)) => {
let _ = url.set_port(None);
}
_ => {}
}
Ok(url)
}
fn scope_matches(scope: &Url, address: &Url) -> bool {
if scope.scheme() != address.scheme()
|| scope.host_str() != address.host_str()
|| scope.port() != address.port()
{
return false;
}
let scope_path = scope.path().trim_end_matches('/');
let address_path = address.path().trim_end_matches('/');
address_path == scope_path
|| address_path
.strip_prefix(scope_path)
.is_some_and(|rest| rest.starts_with('/'))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BindVia {
Pointer,
Scope,
CatchAll,
}
#[derive(Debug, Clone, PartialEq)]
pub enum CredsBinding {
Set { name: String, via: BindVia },
Ambient,
NotApplicable,
}
impl CredsBinding {
pub fn describe(&self) -> String {
match self {
Self::Set { name, via } => {
let via = match via {
BindVia::Pointer => "?creds",
BindVia::Scope => "scope match",
BindVia::CatchAll => "catch-all",
};
format!("creds {name} ({via})")
}
Self::Ambient => "ambient chain".to_owned(),
Self::NotApplicable => "local (no credentials)".to_owned(),
}
}
}
#[derive(Debug, Clone)]
pub struct ResolvedStorage {
storage: StorageUrl,
pub options: HashMap<String, String>,
pub binding: CredsBinding,
}
impl ResolvedStorage {
pub fn lance_url(&self) -> &Url {
self.storage.lance_url()
}
pub fn display(&self) -> String {
self.storage.display()
}
}
pub fn unmatched_creds_sets<'c>(
resolved: &[&ResolvedStorage],
creds: &'c BTreeMap<String, CredsSet>,
) -> Vec<&'c str> {
if resolved
.iter()
.all(|entry| matches!(entry.binding, CredsBinding::NotApplicable))
{
return Vec::new();
}
creds
.keys()
.filter(|name| {
!resolved.iter().any(|entry| {
matches!(&entry.binding, CredsBinding::Set { name: bound, .. } if bound == *name)
})
})
.map(String::as_str)
.collect()
}
fn materialize_secret(
set: &str,
field: &str,
inline: Option<&str>,
file: Option<&std::path::Path>,
command: Option<&str>,
) -> Result<Option<String>> {
if let Some(value) = inline {
return Ok(Some(value.to_owned()));
}
if let Some(path) = file {
let text = std::fs::read_to_string(path).with_context(|| {
format!(
"[creds.{set}] {field}_file: failed to read {}",
path.display()
)
})?;
return Ok(Some(strip_one_newline(text)));
}
if let Some(command) = command {
return Ok(Some(run_secret_command(set, field, command)?));
}
Ok(None)
}
fn run_secret_command(set: &str, field: &str, command: &str) -> Result<String> {
static CACHE: std::sync::OnceLock<std::sync::Mutex<HashMap<String, String>>> =
std::sync::OnceLock::new();
let cache = CACHE.get_or_init(Default::default);
if let Some(hit) = cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(command)
{
return Ok(hit.clone());
}
#[cfg(windows)]
let output = {
use std::os::windows::process::CommandExt as _;
let shell = std::env::var_os("COMSPEC").unwrap_or_else(|| "cmd".into());
std::process::Command::new(shell)
.raw_arg(format!("/C {command}"))
.output()
.with_context(|| format!("[creds.{set}] {field}_command failed to spawn: {command}"))?
};
#[cfg(not(windows))]
let output = std::process::Command::new("sh")
.args(["-c", command])
.output()
.with_context(|| format!("[creds.{set}] {field}_command failed to spawn: {command}"))?;
if !output.status.success() {
bail!(
"[creds.{set}] {field}_command exited {}: {command}\n{}",
output.status,
String::from_utf8_lossy(&output.stderr).trim_end(),
);
}
let value = strip_one_newline(
String::from_utf8(output.stdout)
.with_context(|| format!("[creds.{set}] {field}_command output is not UTF-8"))?,
);
cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(command.to_owned(), value.clone());
Ok(value)
}
fn strip_one_newline(mut text: String) -> String {
if text.ends_with('\n') {
text.pop();
if text.ends_with('\r') {
text.pop();
}
}
text
}
#[derive(Debug, thiserror::Error)]
pub enum CheckFailure {
#[error(
"authentication failed and no creds set matched this URL; add one with `pond creds add` (or set POND_CREDS_*), or provide ambient AWS_* credentials"
)]
NoCreds { source: anyhow::Error },
#[error("authentication failed using creds set {set:?}; check its keys and scope")]
Auth { set: String, source: anyhow::Error },
#[error(
"backend does not enforce conditional writes (If-None-Match); concurrent pond writers would corrupt each other - {detail}"
)]
OccUnsupported { detail: String },
#[error("storage probe failed")]
Io { source: anyhow::Error },
}
impl CheckFailure {
pub fn concise_cause(&self) -> Option<String> {
let source = match self {
Self::NoCreds { source } | Self::Auth { source, .. } | Self::Io { source } => source,
Self::OccUnsupported { .. } => return None,
};
Some(condense_error_chain(source))
}
}
fn condense_error_chain(error: &anyhow::Error) -> String {
let mut text = error
.chain()
.last()
.map(ToString::to_string)
.unwrap_or_else(|| format!("{error:#}"));
if let Some(pos) = text.find(", <WORKSPACE>") {
text.truncate(pos);
}
text = text.replace(
"Encountered internal error. Please file a bug report at https://github.com/lance-format/lance/issues. ",
"",
);
let line = text.split_whitespace().collect::<Vec<_>>().join(" ");
const HEAD: usize = 120;
const TAIL: usize = 120;
let chars: Vec<char> = line.chars().collect();
if chars.len() > HEAD + TAIL + 5 {
let head: String = chars[..HEAD].iter().collect();
let tail: String = chars[chars.len() - TAIL..].iter().collect();
format!("{head} ... {tail}")
} else {
line
}
}
pub async fn storage_check(resolved: &ResolvedStorage) -> std::result::Result<(), CheckFailure> {
use object_store::{Error as OsError, ObjectStoreExt, PutMode, PutOptions, PutPayload};
let classify =
|error: OsError, step: &str| classify_check_error(error, &resolved.binding, step);
let probe_uri = format!(
"{}/_config-check/{}",
resolved.lance_url().as_str().trim_end_matches('/'),
uuid::Uuid::now_v7(),
);
let params = ObjectStoreParams {
storage_options_accessor: (!resolved.options.is_empty()).then(|| {
Arc::new(StorageOptionsAccessor::with_static_options(
resolved.options.clone(),
))
}),
..Default::default()
};
let registry = Arc::new(ObjectStoreRegistry::default());
let (store, path) = ObjectStore::from_uri_and_params(registry, &probe_uri, ¶ms)
.await
.map_err(|error| CheckFailure::Io {
source: anyhow!(error).context(format!("failed to open object store for {probe_uri}")),
})?;
let body: &[u8] = b"pond storage check";
let create = PutOptions::from(PutMode::Create);
store
.inner
.put_opts(&path, PutPayload::from_static(body), create.clone())
.await
.map_err(|error| classify(error, "initial conditional put"))?;
let outcome = async {
match store
.inner
.put_opts(&path, PutPayload::from_static(body), create)
.await
{
Err(OsError::AlreadyExists { .. }) => {}
Ok(_) => {
return Err(CheckFailure::OccUnsupported {
detail: "a second create over an existing key succeeded".to_owned(),
});
}
Err(OsError::NotImplemented { .. }) => {
return Err(CheckFailure::OccUnsupported {
detail: "the backend rejects conditional puts as unimplemented".to_owned(),
});
}
Err(error) => return Err(classify(error, "conditional-put probe")),
}
let read_back = store
.inner
.get(&path)
.await
.map_err(|error| classify(error, "read-back"))?
.bytes()
.await
.map_err(|error| classify(error, "read-back body"))?;
if read_back.as_ref() != body {
return Err(CheckFailure::Io {
source: anyhow!("read-back returned different bytes than written"),
});
}
Ok(())
}
.await;
let cleanup = store.inner.delete(&path).await;
outcome?;
cleanup.map_err(|error| classify(error, "cleanup delete"))?;
Ok(())
}
fn classify_check_error(
error: object_store::Error,
binding: &CredsBinding,
step: &str,
) -> CheckFailure {
use object_store::Error as OsError;
let auth_class = matches!(
error,
OsError::Unauthenticated { .. } | OsError::PermissionDenied { .. }
) || {
let rendered = error.to_string();
rendered.contains("CredentialsNotLoaded")
|| rendered.contains("no providers in chain provided credentials")
};
match (auth_class, binding) {
(true, CredsBinding::Set { name, .. }) => CheckFailure::Auth {
set: name.clone(),
source: anyhow!(error).context(step.to_owned()),
},
(true, _) => CheckFailure::NoCreds {
source: anyhow!(error).context(step.to_owned()),
},
(false, _) => CheckFailure::Io {
source: anyhow!(error).context(step.to_owned()),
},
}
}
pub const DEFAULT_COMPACTION_FRAGMENT_CAP: usize = 64;
pub const TARGET_FRAGMENT_BYTES: u64 = 256 * 1024 * 1024;
const MAX_TARGET_ROWS_PER_FRAGMENT: u64 = 1024 * 1024;
pub const COMPACTION_ABSORB_FACTOR: u64 = 4;
pub fn default_cleanup_older_than() -> chrono::Duration {
chrono::Duration::hours(1)
}
pub const DEFAULT_SYNC_CLEANUP_INTERVAL: u64 = 16;
pub const DEFAULT_SYNC_SCALAR_FOLD_ROWS: usize = 50_000;
pub const DEFAULT_SYNC_INDEX_FOLD_ROWS: usize = 5_000;
#[derive(Debug, Clone, Copy)]
pub struct MaintenancePolicy {
pub compaction_fragment_cap: usize,
pub cleanup_older_than: chrono::Duration,
pub cleanup_interval: u64,
pub scalar_fold_row_threshold: usize,
pub index_fold_row_threshold: usize,
}
impl MaintenancePolicy {
pub fn always_compact() -> Self {
Self {
compaction_fragment_cap: 0,
cleanup_older_than: default_cleanup_older_than(),
cleanup_interval: 1,
scalar_fold_row_threshold: 0,
index_fold_row_threshold: 0,
}
}
#[must_use]
pub fn with_cleanup_interval(mut self, interval: u64) -> Self {
self.cleanup_interval = interval.max(1);
self
}
#[must_use]
pub fn with_scalar_fold_row_threshold(mut self, threshold: usize) -> Self {
self.scalar_fold_row_threshold = threshold;
self
}
#[must_use]
pub fn with_index_fold_row_threshold(mut self, threshold: usize) -> Self {
self.index_fold_row_threshold = threshold;
self
}
fn fold_thresholds(&self) -> FoldThresholds {
FoldThresholds {
scalar: self.scalar_fold_row_threshold,
index: self.index_fold_row_threshold,
}
}
}
#[derive(Debug, Clone, Copy)]
struct FoldThresholds {
scalar: usize,
index: usize,
}
struct FragmentStat {
bytes: Option<u64>,
rows: u64,
deleted_rows: u64,
}
fn fragment_bytes(fragment: &lance::table::format::Fragment) -> Option<u64> {
fragment.files.iter().try_fold(0u64, |total, file| {
Some(total + file.file_size_bytes.get()?.get())
})
}
fn fragment_stat(fragment: &lance::table::format::Fragment) -> FragmentStat {
FragmentStat {
bytes: fragment_bytes(fragment),
rows: fragment.physical_rows.unwrap_or(0) as u64,
deleted_rows: fragment
.deletion_file
.as_ref()
.and_then(|deletions| deletions.num_deleted_rows)
.unwrap_or(0) as u64,
}
}
fn derived_target_rows(stats: &[FragmentStat]) -> usize {
let (mut bytes, mut rows) = (0u64, 0u64);
for stat in stats {
if let Some(fragment_bytes) = stat.bytes
&& stat.rows > 0
{
bytes += fragment_bytes;
rows += stat.rows;
}
}
if bytes == 0 || rows == 0 {
return MAX_TARGET_ROWS_PER_FRAGMENT as usize;
}
((u128::from(TARGET_FRAGMENT_BYTES / 2) * u128::from(rows) / u128::from(bytes))
.clamp(1, u128::from(MAX_TARGET_ROWS_PER_FRAGMENT))) as usize
}
fn task_veto_reason(
stats: &[FragmentStat],
cap: usize,
deletion_threshold: f32,
target_rows_per_fragment: usize,
max_bytes_per_file: u64,
) -> Option<&'static str> {
if stats.iter().any(|stat| {
stat.rows > 0 && (stat.deleted_rows as f32 / stat.rows as f32) > deletion_threshold
}) {
return None;
}
let budget = u128::from(max_bytes_per_file);
if budget == 0 {
return Some("invalid_byte_budget");
}
let (mut total_bytes, mut largest) = (0u128, 0u128);
for stat in stats {
let Some(bytes) = stat.bytes.map(u128::from) else {
return Some("missing_sizes");
};
total_bytes += bytes;
largest = largest.max(bytes);
}
let minimum_outputs = total_bytes.div_ceil(budget).max(1);
if u128::try_from(stats.len()).unwrap_or(u128::MAX) <= minimum_outputs {
return Some("cannot_shrink");
}
if stats.len() >= cap {
return None;
}
if total_bytes > budget {
for stat in stats {
let bytes = u128::from(stat.bytes.unwrap_or(0));
let rows = u128::from(stat.rows);
if rows == 0 || bytes * target_rows_per_fragment as u128 * 2 > rows * budget {
return Some("row_target_unattainable");
}
}
}
if (total_bytes - largest) * u128::from(COMPACTION_ABSORB_FACTOR) < largest {
return Some("absorb_veto");
}
None
}
#[derive(Debug, Clone)]
pub struct IndexIntent {
pub name: &'static str,
pub column: &'static str,
pub trigger: IndexTrigger,
pub params: IndexParamsKind,
}
impl IndexIntent {
fn presence_column(&self) -> &'static str {
match self.trigger {
IndexTrigger::OnNonNullCount { column, .. } => column,
IndexTrigger::OnAnyRows => self.column,
}
}
}
#[derive(Debug, Clone)]
pub enum IndexTrigger {
OnAnyRows,
OnNonNullCount {
column: &'static str,
threshold: usize,
},
}
#[derive(Debug, Clone)]
pub enum IndexParamsKind {
Scalar(BuiltinIndexType),
InvertedFtsWord,
IvfSqCosine { num_bits: u16, max_iters: usize },
}
impl IndexTrigger {
async fn should_create(&self, dataset: &Dataset) -> Result<bool> {
match self {
Self::OnAnyRows => Ok(dataset.count_rows(None).await? > 0),
Self::OnNonNullCount { column, threshold } => {
let count = dataset
.count_rows(Some(format!("{column} IS NOT NULL")))
.await?;
Ok(count >= *threshold)
}
}
}
}
impl IndexParamsKind {
fn index_type(&self) -> IndexType {
match self {
Self::Scalar(BuiltinIndexType::Bitmap) => IndexType::Bitmap,
Self::Scalar(BuiltinIndexType::ZoneMap) => IndexType::ZoneMap,
Self::Scalar(_) => IndexType::BTree,
Self::InvertedFtsWord => IndexType::Inverted,
Self::IvfSqCosine { .. } => IndexType::Vector,
}
}
async fn build(&self, dataset: &Dataset) -> Result<Box<dyn lance::index::IndexParams>> {
match self {
Self::Scalar(kind) => Ok(Box::new(ScalarIndexParams::for_builtin(kind.clone()))),
Self::InvertedFtsWord => Ok(Box::new(
InvertedIndexParams::default()
.base_tokenizer("simple".to_owned())
.stem(true)
.remove_stop_words(false),
)),
Self::IvfSqCosine {
num_bits,
max_iters,
} => {
let count = dataset
.count_rows(Some("vector IS NOT NULL".to_owned()))
.await?;
let partitions = count.checked_div(4096).unwrap_or(0).max(1);
let mut ivf = IvfBuildParams::new(partitions);
ivf.max_iters = *max_iters;
let sq = SQBuildParams {
num_bits: *num_bits,
..Default::default()
};
Ok(Box::new(VectorIndexParams::with_ivf_sq_params(
MetricType::Cosine,
ivf,
sq,
)))
}
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IndexStatus {
pub table: Table,
pub intent_name: String,
pub fragments_covered: usize,
pub unindexed_fragments: usize,
pub unindexed_rows: usize,
pub exists: bool,
}
#[derive(Debug, Clone, Copy)]
pub struct ConflictExhausted {
pub attempts: u8,
}
impl std::fmt::Display for ConflictExhausted {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"commit conflict exhausted after {} attempt(s)",
self.attempts
)
}
}
impl std::error::Error for ConflictExhausted {}
#[derive(Debug)]
pub enum PhaseOutcome {
Ok,
Noop,
SkippedConflict,
Failed(anyhow::Error),
NotAttempted,
}
impl PhaseOutcome {
pub fn is_failed(&self) -> bool {
matches!(self, Self::Failed(_))
}
}
#[derive(Debug)]
pub struct TableOptimizeOutcome {
pub table: Table,
pub indices: PhaseOutcome,
pub compaction: PhaseOutcome,
}
#[derive(Debug, Clone)]
pub enum OptimizeEvent {
PhaseStart {
table: Table,
phase: OptimizePhase,
detail: Option<String>,
},
PhaseDone {
table: Table,
phase: OptimizePhase,
elapsed_ms: u64,
},
IndexStage {
table: Table,
index: String,
stage: String,
completed: u64,
total: Option<u64>,
unit: String,
},
}
#[derive(Debug, Clone, Copy)]
pub enum OptimizePhase {
Compact,
Cleanup,
IndexCreate,
IndexRebuild,
IndexAppend,
}
impl OptimizePhase {
pub fn label(self) -> &'static str {
match self {
Self::Compact => "compact",
Self::Cleanup => "cleanup",
Self::IndexCreate => "index-create",
Self::IndexRebuild => "index-rebuild",
Self::IndexAppend => "index-append",
}
}
}
pub type OptimizeProgressFn = Arc<dyn Fn(OptimizeEvent) + Send + Sync>;
fn emit(progress: Option<&OptimizeProgressFn>, event: OptimizeEvent) {
if let Some(callback) = progress {
callback(event);
}
}
struct PondIndexProgress {
callback: OptimizeProgressFn,
table: Table,
index: String,
state: std::sync::Mutex<PondIndexStageState>,
}
impl std::fmt::Debug for PondIndexProgress {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PondIndexProgress")
.field("table", &self.table)
.field("index", &self.index)
.finish_non_exhaustive()
}
}
#[derive(Debug, Default)]
struct PondIndexStageState {
total: Option<u64>,
unit: String,
last_emit: Option<Instant>,
}
impl PondIndexProgress {
fn new(callback: OptimizeProgressFn, table: Table, index: String) -> Arc<Self> {
Arc::new(Self {
callback,
table,
index,
state: std::sync::Mutex::new(PondIndexStageState::default()),
})
}
}
#[async_trait::async_trait]
impl lance_index::progress::IndexBuildProgress for PondIndexProgress {
async fn stage_start(&self, stage: &str, total: Option<u64>, unit: &str) -> lance::Result<()> {
if let Ok(mut state) = self.state.lock() {
state.total = total;
state.unit = unit.to_owned();
state.last_emit = Some(Instant::now());
}
(self.callback)(OptimizeEvent::IndexStage {
table: self.table,
index: self.index.clone(),
stage: stage.to_owned(),
completed: 0,
total,
unit: unit.to_owned(),
});
Ok(())
}
async fn stage_progress(&self, stage: &str, completed: u64) -> lance::Result<()> {
let (total, unit) = {
let Ok(mut state) = self.state.lock() else {
return Ok(());
};
let now = Instant::now();
if let Some(prev) = state.last_emit
&& now.duration_since(prev) < Duration::from_millis(100)
{
return Ok(());
}
state.last_emit = Some(now);
(state.total, state.unit.clone())
};
(self.callback)(OptimizeEvent::IndexStage {
table: self.table,
index: self.index.clone(),
stage: stage.to_owned(),
completed,
total,
unit,
});
Ok(())
}
async fn stage_complete(&self, stage: &str) -> lance::Result<()> {
let (total, unit) = {
let Ok(state) = self.state.lock() else {
return Ok(());
};
(state.total, state.unit.clone())
};
(self.callback)(OptimizeEvent::IndexStage {
table: self.table,
index: self.index.clone(),
stage: stage.to_owned(),
completed: total.unwrap_or(0),
total,
unit,
});
Ok(())
}
}
fn lance_progress(
progress: Option<&OptimizeProgressFn>,
table: Table,
index: &str,
) -> Arc<dyn lance_index::progress::IndexBuildProgress> {
match progress {
Some(callback) => PondIndexProgress::new(callback.clone(), table, index.to_owned()),
None => Arc::new(lance_index::progress::NoopIndexBuildProgress),
}
}
pub fn is_commit_conflict(error: &anyhow::Error) -> bool {
#[cfg(windows)]
if is_transient_sharing_violation(error) {
return true;
}
error.downcast_ref::<lance::Error>().is_some_and(|err| {
matches!(
err,
lance::Error::CommitConflict { .. }
| lance::Error::RetryableCommitConflict { .. }
| lance::Error::TooMuchWriteContention { .. }
)
})
}
#[cfg(any(windows, test))]
fn is_transient_sharing_violation(error: &anyhow::Error) -> bool {
const CONTENDED: [i32; 2] = [32, 33];
error.chain().any(|cause| {
cause
.downcast_ref::<std::io::Error>()
.and_then(std::io::Error::raw_os_error)
.is_some_and(|code| CONTENDED.contains(&code))
})
}
fn is_conflict_exhausted(error: &anyhow::Error) -> bool {
error.chain().any(|cause| cause.is::<ConflictExhausted>())
}
pub fn is_index_error(error: &anyhow::Error) -> bool {
error
.downcast_ref::<lance::Error>()
.is_some_and(|err| matches!(err, lance::Error::Index { .. }))
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TableSizes {
pub sessions: u64,
pub messages: u64,
pub parts: u64,
pub other: u64,
pub sessions_data: DataLiveness,
pub messages_data: DataLiveness,
pub parts_data: DataLiveness,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct DataLiveness {
pub on_disk: u64,
pub live: Option<u64>,
}
impl DataLiveness {
pub fn dead(&self) -> Option<u64> {
self.live.map(|live| self.on_disk.saturating_sub(live))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ScalarValue {
String(String),
Int32(i32),
Raw(String),
}
impl From<&str> for ScalarValue {
fn from(value: &str) -> Self {
Self::String(value.to_owned())
}
}
impl From<String> for ScalarValue {
fn from(value: String) -> Self {
Self::String(value)
}
}
impl From<i32> for ScalarValue {
fn from(value: i32) -> Self {
Self::Int32(value)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Predicate {
Eq(&'static str, ScalarValue),
Ne(&'static str, ScalarValue),
IsNull(&'static str),
IsNotNull(&'static str),
In(&'static str, Vec<ScalarValue>),
LikeContains(&'static str, String),
Regex(&'static str, String),
Gte(&'static str, ScalarValue),
Lte(&'static str, ScalarValue),
And(Vec<Predicate>),
Or(Vec<Predicate>),
Not(Box<Predicate>),
}
impl Predicate {
pub fn to_lance(&self) -> String {
match self {
Self::Eq(column, value) => format!("{column} = {}", value.to_lance()),
Self::Ne(column, value) => format!("{column} <> {}", value.to_lance()),
Self::IsNull(column) => format!("{column} IS NULL"),
Self::IsNotNull(column) => format!("{column} IS NOT NULL"),
Self::In(column, values) => {
let values = values
.iter()
.map(ScalarValue::to_lance)
.collect::<Vec<_>>()
.join(", ");
format!("{column} IN ({values})")
}
Self::LikeContains(column, value) => {
format!("{column} LIKE {} ESCAPE '\\'", like_contains(value))
}
Self::Regex(column, pattern) => {
format!("regexp_like({column}, {})", quoted_string(pattern))
}
Self::Gte(column, value) => format!("{column} >= {}", value.to_lance()),
Self::Lte(column, value) => format!("{column} <= {}", value.to_lance()),
Self::And(predicates) => predicates
.iter()
.map(Self::to_lance)
.filter(|predicate| !predicate.is_empty())
.collect::<Vec<_>>()
.join(" AND "),
Self::Or(predicates) => {
let body = predicates
.iter()
.map(Self::to_lance)
.filter(|predicate| !predicate.is_empty())
.collect::<Vec<_>>()
.join(" OR ");
if body.is_empty() {
String::new()
} else {
format!("({body})")
}
}
Self::Not(inner) => {
let body = inner.to_lance();
if body.is_empty() {
String::new()
} else {
format!("NOT ({body})")
}
}
}
}
}
#[derive(Default)]
pub struct ScanOpts<'a> {
pub predicate: Option<&'a Predicate>,
pub projection: Option<&'a [&'a str]>,
}
impl<'a> ScanOpts<'a> {
pub fn project_only(projection: &'a [&'a str]) -> Self {
Self {
predicate: None,
projection: Some(projection),
}
}
pub fn with_predicate_and_projection(
predicate: &'a Predicate,
projection: &'a [&'a str],
) -> Self {
Self {
predicate: Some(predicate),
projection: Some(projection),
}
}
}
impl ScalarValue {
fn to_lance(&self) -> String {
match self {
Self::String(value) => quoted_string(value),
Self::Int32(value) => value.to_string(),
Self::Raw(value) => value.clone(),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RuntimeCaps {
pub index_cache_bytes: Option<usize>,
pub metadata_cache_bytes: Option<usize>,
}
impl RuntimeCaps {
pub fn from_config(config: &crate::config::RuntimeConfig) -> Self {
Self {
index_cache_bytes: config.index_cache_bytes,
metadata_cache_bytes: config.metadata_cache_bytes,
}
}
}
const LOCAL_INDEX_CACHE_BYTES: usize = 256 * 1024 * 1024;
const LOCAL_METADATA_CACHE_BYTES: usize = 128 * 1024 * 1024;
const REMOTE_INDEX_CACHE_BYTES: usize = 1024 * 1024 * 1024;
const REMOTE_METADATA_CACHE_BYTES: usize = 512 * 1024 * 1024;
fn resolve_cache_caps(location: &Url, caps: RuntimeCaps) -> (usize, usize) {
let (index_default, metadata_default) = if config::is_local(location) {
(LOCAL_INDEX_CACHE_BYTES, LOCAL_METADATA_CACHE_BYTES)
} else {
(REMOTE_INDEX_CACHE_BYTES, REMOTE_METADATA_CACHE_BYTES)
};
(
caps.index_cache_bytes.unwrap_or(index_default),
caps.metadata_cache_bytes.unwrap_or(metadata_default),
)
}
pub struct Handle {
datasets: DatasetSet,
retry: RetryPolicy,
#[allow(dead_code)]
session: Arc<Session>,
nm: Arc<dyn LanceNamespace>,
nm_ident: NamespaceIdent,
storage_options: HashMap<String, String>,
location: Url,
lazy_refresh_after: Duration,
store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
}
impl std::fmt::Debug for Handle {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("Handle")
.field("datasets", &self.datasets)
.field("retry", &self.retry)
.field("nm_ident", &self.nm_ident)
.field("storage_options", &self.storage_options)
.field("location", &self.location)
.finish()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Table {
Sessions,
Messages,
Parts,
}
impl Table {
pub fn as_str(self) -> &'static str {
self.label()
}
fn label(self) -> &'static str {
match self {
Self::Sessions => "sessions",
Self::Messages => "messages",
Self::Parts => "parts",
}
}
}
#[derive(Debug)]
struct DatasetSet {
sessions: OnceCell<Mutex<CachedDataset>>,
messages: Mutex<CachedDataset>,
parts: OnceCell<Mutex<CachedDataset>>,
}
#[derive(Debug)]
struct CachedDataset {
dataset: Dataset,
last_refresh: Instant,
refresh_after: Duration,
}
impl CachedDataset {
fn new(dataset: Dataset, refresh_after: Duration) -> Self {
Self {
dataset,
last_refresh: Instant::now(),
refresh_after,
}
}
async fn latest(&mut self) -> Result<Dataset> {
if self.last_refresh.elapsed() >= self.refresh_after {
self.dataset.checkout_latest().await?;
self.last_refresh = Instant::now();
}
Ok(self.dataset.clone())
}
fn replace(&mut self, dataset: Dataset) {
self.dataset = dataset;
self.last_refresh = Instant::now();
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct AppendStats {
pub rows: u64,
pub bytes_written: u64,
pub files_written: u64,
pub attempts: u32,
}
#[derive(Default)]
struct WriteAccum {
rows: std::sync::atomic::AtomicU64,
bytes: std::sync::atomic::AtomicU64,
files: std::sync::atomic::AtomicU64,
}
impl WriteAccum {
fn observe(&self, stats: &WriteStats) {
use std::sync::atomic::Ordering::Relaxed;
self.rows.fetch_max(stats.rows_written, Relaxed);
self.bytes.fetch_max(stats.bytes_written, Relaxed);
self.files.fetch_max(stats.files_written as u64, Relaxed);
}
fn rows(&self) -> u64 {
self.rows.load(std::sync::atomic::Ordering::Relaxed)
}
fn bytes(&self) -> u64 {
self.bytes.load(std::sync::atomic::Ordering::Relaxed)
}
fn files(&self) -> u64 {
self.files.load(std::sync::atomic::Ordering::Relaxed)
}
}
fn append_write_params() -> WriteParams {
let mut params = sessions::write_params_for_create();
params.mode = WriteMode::Append;
params.max_bytes_per_file = TARGET_FRAGMENT_BYTES as usize;
params
}
impl Handle {
pub async fn open(location: &Url) -> Result<Self> {
Self::open_with_options(location, HashMap::new(), RuntimeCaps::default()).await
}
pub fn lance_cache_bytes(&self) -> u64 {
self.session.size_bytes()
}
pub async fn open_with_options(
location: &Url,
storage_options: HashMap<String, String>,
caps: RuntimeCaps,
) -> Result<Self> {
Self::open_with_options_cached(location, storage_options, caps, None).await
}
pub async fn open_with_options_cached(
location: &Url,
mut storage_options: HashMap<String, String>,
caps: RuntimeCaps,
index_cache_dir: Option<PathBuf>,
) -> Result<Self> {
if let Some(path) = config::local_path(location) {
tokio::fs::create_dir_all(&path).await.with_context(|| {
format!(
"failed to create data dir {}; fix the storage destination ([storage].path in config) or re-run `pond init`",
path.display()
)
})?;
} else {
apply_remote_storage_defaults(&mut storage_options);
}
let (index_cache_bytes, metadata_cache_bytes) = resolve_cache_caps(location, caps);
let session = Arc::new(Session::new(
index_cache_bytes,
metadata_cache_bytes,
Arc::new(ObjectStoreRegistry::default()),
));
let root = location.as_str().trim_end_matches('/').to_string();
let mut connect = ConnectBuilder::new("dir")
.property("root", root)
.session(session.clone());
for (key, value) in &storage_options {
connect = connect.property(format!("storage.{key}"), value.clone());
}
let nm: Arc<dyn LanceNamespace> = connect
.connect()
.await
.context("failed to connect lance Directory namespace")?;
let nm_ident = NamespaceIdent::root();
let refresh_after = if config::is_local(location) {
Duration::ZERO
} else {
Duration::from_secs(5)
};
let wrapper = store_wrapper(location, index_cache_dir.as_deref());
let handle = Self {
datasets: DatasetSet {
sessions: OnceCell::new(),
messages: Mutex::new(CachedDataset::new(
open_or_create_via_ns(
&nm,
&nm_ident,
sessions::MESSAGES,
sessions::message_schema(),
&session,
&storage_options,
wrapper.clone(),
)
.await?,
refresh_after,
)),
parts: OnceCell::new(),
},
retry: RetryPolicy::default(),
session,
nm,
nm_ident,
storage_options,
location: location.clone(),
lazy_refresh_after: refresh_after,
store_wrapper: wrapper,
};
Ok(handle)
}
pub fn location(&self) -> &Url {
&self.location
}
pub fn storage_options(&self) -> &HashMap<String, String> {
&self.storage_options
}
fn export_uri(&self, name: &str) -> String {
format!(
"{}/exports/{name}",
self.location.as_str().trim_end_matches('/')
)
}
fn object_store_params(&self) -> ObjectStoreParams {
ObjectStoreParams {
storage_options_accessor: (!self.storage_options.is_empty()).then(|| {
Arc::new(StorageOptionsAccessor::with_static_options(
self.storage_options.clone(),
))
}),
..Default::default()
}
}
pub(crate) async fn export_write(&self, name: &str, bytes: &[u8]) -> Result<()> {
let uri = self.export_uri(name);
let registry = Arc::new(ObjectStoreRegistry::default());
let (store, path) =
ObjectStore::from_uri_and_params(registry, &uri, &self.object_store_params())
.await
.with_context(|| format!("failed to open object store for {uri}"))?;
store
.put(&path, bytes)
.await
.with_context(|| format!("failed to write export {uri}"))?;
Ok(())
}
pub(crate) async fn export_read(&self, name: &str) -> Result<Vec<u8>> {
let uri = self.export_uri(name);
let registry = Arc::new(ObjectStoreRegistry::default());
let (store, path) =
ObjectStore::from_uri_and_params(registry, &uri, &self.object_store_params())
.await
.with_context(|| format!("failed to open object store for {uri}"))?;
let bytes = store
.read_one_all(&path)
.await
.with_context(|| format!("failed to read export {uri}"))?;
Ok(bytes.to_vec())
}
pub(crate) fn export_local_path(&self, name: &str) -> Option<std::path::PathBuf> {
if self.location.scheme() != "file" {
return None;
}
let dir = self.location.to_file_path().ok()?;
Some(dir.join("exports").join(name))
}
pub async fn row_counts(&self) -> Result<(usize, usize, usize)> {
Ok((
self.count_rows(Table::Sessions).await?,
self.count_rows(Table::Messages).await?,
self.count_rows(Table::Parts).await?,
))
}
pub(crate) async fn merge_insert(
&self,
table: Table,
batch: RecordBatch,
row_count: usize,
) -> Result<u64> {
self.merge_insert_stats(table, batch, row_count)
.await
.map(|stats| stats.num_inserted_rows + stats.num_updated_rows)
}
pub(crate) async fn merge_insert_stats(
&self,
table: Table,
batch: RecordBatch,
row_count: usize,
) -> Result<MergeStats> {
self.merge(
table,
batch,
row_count,
"merge_insert",
WhenMatched::DoNothing,
WhenNotMatched::InsertAll,
)
.await
}
pub(crate) async fn merge_update(
&self,
table: Table,
batch: RecordBatch,
row_count: usize,
) -> Result<u64> {
self.merge(
table,
batch,
row_count,
"merge_update",
WhenMatched::UpdateAll,
WhenNotMatched::DoNothing,
)
.await
.map(|stats| stats.num_inserted_rows + stats.num_updated_rows)
}
async fn write_committed<E, Fut, P>(&self, table: Table, execute: E) -> Result<P>
where
E: Fn(Arc<Dataset>) -> Fut,
Fut: std::future::Future<Output = Result<(Dataset, P)>>,
{
self.write_committed_with(table, |_| true, execute).await
}
async fn write_committed_with<E, Fut, P, R>(
&self,
table: Table,
should_retry: R,
execute: E,
) -> Result<P>
where
E: Fn(Arc<Dataset>) -> Fut,
Fut: std::future::Future<Output = Result<(Dataset, P)>>,
R: Fn(&anyhow::Error) -> bool,
{
self.retry_lance_filtered(table.label(), should_retry, || {
let execute = &execute;
async move {
let mut cached = self.cached(table).await?.lock().await;
let existing = cached.latest().await?;
let (dataset, payload) = execute(Arc::new(existing)).await?;
cached.replace(dataset);
Ok(payload)
}
})
.await
}
async fn merge(
&self,
table: Table,
batch: RecordBatch,
row_count: usize,
op: &'static str,
when_matched: WhenMatched,
when_not_matched: WhenNotMatched,
) -> Result<MergeStats> {
if row_count == 0 {
return Ok(MergeStats::default());
}
let started = Instant::now();
let result = self
.write_committed(table, |existing| {
let batch = batch.clone();
let when_matched = when_matched.clone();
let when_not_matched = when_not_matched.clone();
async move {
let schema = batch.schema();
let reader = RecordBatchIterator::new([Ok(batch)], schema);
let mut builder = MergeInsertBuilder::try_new(existing, Vec::new())?;
builder.when_matched(when_matched);
builder.when_not_matched(when_not_matched);
builder.source_dedupe_behavior(SourceDedupeBehavior::FirstSeen);
builder.skip_auto_cleanup(true);
let (dataset, stats) = builder
.try_build()?
.execute_reader(Box::new(reader))
.await?;
Ok((dataset.as_ref().clone(), stats))
}
})
.await;
let skipped = result
.as_ref()
.map(|s| s.num_skipped_duplicates)
.unwrap_or(0);
tracing::info!(
target: "pond::perf",
op,
table = %table.label(),
rows = row_count,
elapsed_ms = started.elapsed().as_millis() as u64,
skipped,
"merge",
);
result
}
pub(crate) async fn append_stream<F, Fut>(
&self,
table: Table,
make_source: F,
) -> Result<AppendStats>
where
F: Fn() -> Fut,
Fut: std::future::Future<Output = Result<SendableRecordBatchStream>>,
{
let cum = Arc::new(WriteAccum::default());
let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
let started = Instant::now();
self.write_committed(table, |existing| {
let make_source = &make_source;
let cum = cum.clone();
let attempts = attempts.clone();
async move {
attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let stream = make_source().await?;
let dataset = InsertBuilder::new(existing)
.with_params(&append_write_params())
.progress(move |stats| cum.observe(&stats))
.execute_stream(stream)
.await?;
Ok((dataset, ()))
}
})
.await?;
let attempts = attempts.load(std::sync::atomic::Ordering::Relaxed);
let stats = AppendStats {
rows: cum.rows(),
bytes_written: cum.bytes(),
files_written: cum.files(),
attempts,
};
tracing::info!(
target: "pond::perf",
op = "append",
table = %table.label(),
rows = stats.rows,
files = stats.files_written,
attempts,
elapsed_ms = started.elapsed().as_millis() as u64,
"append",
);
Ok(stats)
}
pub(crate) async fn append_batches(
&self,
table: Table,
batches: Vec<RecordBatch>,
) -> Result<AppendStats> {
let total_rows: u64 = batches.iter().map(|batch| batch.num_rows() as u64).sum();
if total_rows == 0 {
return Ok(AppendStats::default());
}
let cum = Arc::new(WriteAccum::default());
let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
let started = Instant::now();
self.write_committed_with(table, is_commit_conflict, |existing| {
let cum = cum.clone();
let attempts = attempts.clone();
let batches = batches.clone();
async move {
attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let dataset = InsertBuilder::new(existing)
.with_params(&append_write_params())
.progress(move |stats| cum.observe(&stats))
.execute(batches)
.await?;
Ok((dataset, ()))
}
})
.await?;
let attempts = attempts.load(std::sync::atomic::Ordering::Relaxed);
let stats = AppendStats {
rows: total_rows,
bytes_written: cum.bytes(),
files_written: cum.files(),
attempts,
};
tracing::info!(
target: "pond::perf",
op = "append_batches",
table = %table.label(),
rows = stats.rows,
files = stats.files_written,
attempts,
elapsed_ms = started.elapsed().as_millis() as u64,
"append",
);
Ok(stats)
}
pub async fn optimize_table(
&self,
table: Table,
intents: &[IndexIntent],
progress: Option<&OptimizeProgressFn>,
policy: &MaintenancePolicy,
) -> TableOptimizeOutcome {
let compaction = self
.run_optimize_compact_phase(table, progress, policy)
.await;
let indices = self
.run_optimize_indices_phase(table, intents, progress, policy.fold_thresholds())
.await;
TableOptimizeOutcome {
table,
indices,
compaction,
}
}
pub async fn optimize_table_indices_only(
&self,
table: Table,
intents: &[IndexIntent],
progress: Option<&OptimizeProgressFn>,
) -> PhaseOutcome {
self.run_optimize_indices_phase(
table,
intents,
progress,
FoldThresholds {
scalar: 0,
index: 0,
},
)
.await
}
async fn run_optimize_indices_phase(
&self,
table: Table,
intents: &[IndexIntent],
progress: Option<&OptimizeProgressFn>,
folds: FoldThresholds,
) -> PhaseOutcome {
if intents.is_empty() {
return PhaseOutcome::Noop;
}
let result = self
.retry_lance(table.label(), || async {
let mut guard = self.cached(table).await?.lock().await;
let mut dataset = guard.latest().await?;
let did_work =
optimize_table_indices(&mut dataset, intents, table, progress, folds).await?;
guard.replace(dataset);
Ok::<_, anyhow::Error>(did_work)
})
.await;
match result {
Ok(true) => PhaseOutcome::Ok,
Ok(false) => PhaseOutcome::Noop,
Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
Err(error) => PhaseOutcome::Failed(error),
}
}
async fn run_optimize_compact_phase(
&self,
table: Table,
progress: Option<&OptimizeProgressFn>,
policy: &MaintenancePolicy,
) -> PhaseOutcome {
let result = self
.retry_lance(table.label(), || async {
let mut guard = self.cached(table).await?.lock().await;
let mut dataset = guard.latest().await?;
optimize_table_compact(&mut dataset, table, progress, policy).await?;
guard.replace(dataset);
Ok::<_, anyhow::Error>(())
})
.await;
match result {
Ok(()) => PhaseOutcome::Ok,
Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
Err(error) => PhaseOutcome::Failed(error),
}
}
pub async fn rebuild_index(
&self,
table: Table,
intent: &IndexIntent,
progress: Option<&OptimizeProgressFn>,
) -> Result<()> {
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::IndexRebuild,
detail: Some(intent.name.to_owned()),
},
);
let started = Instant::now();
let result = self
.retry_lance(table.label(), || async {
let mut guard = self.cached(table).await?.lock().await;
let mut dataset = guard.latest().await?;
rebuild_index(&mut dataset, intent, progress, table).await?;
guard.replace(dataset);
Ok(())
})
.await;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::IndexRebuild,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
result
}
pub async fn cleanup_table_versions(
&self,
table: Table,
older_than: chrono::Duration,
) -> Result<()> {
let mut guard = self.cached(table).await?.lock().await;
let dataset = guard.latest().await?;
dataset
.cleanup_old_versions(older_than, Some(false), Some(false))
.await
.with_context(|| format!("cleanup_old_versions failed for {}", table.label()))?;
Ok(())
}
pub async fn index_status(
&self,
table: Table,
intents: &[IndexIntent],
indexable_only: bool,
) -> Result<Vec<IndexStatus>> {
let dataset = self.dataset(table).await?;
index_status(table, &dataset, intents, indexable_only).await
}
pub(crate) async fn dataset(&self, table: Table) -> Result<Dataset> {
let mut cached = self.cached(table).await?.lock().await;
cached.latest().await
}
pub(crate) async fn scanner(
&self,
table: Table,
predicate: Option<&Predicate>,
) -> Result<lance::dataset::scanner::Scanner> {
let dataset = self.dataset(table).await?;
scanner_with_prefilter(&dataset, predicate)
}
pub async fn scan(
&self,
table: Table,
opts: ScanOpts<'_>,
) -> Result<lance::dataset::scanner::Scanner> {
let mut scanner = self.scanner(table, opts.predicate).await?;
if let Some(projection) = opts.projection {
scanner.project(projection)?;
}
Ok(scanner)
}
pub(crate) async fn scan_batch(
&self,
table: Table,
predicate: Option<&Predicate>,
projection: &[&str],
) -> Result<RecordBatch> {
let opts = ScanOpts {
predicate,
projection: (!projection.is_empty()).then_some(projection),
};
self.scan(table, opts)
.await?
.try_into_batch()
.await
.context("scan failed")
}
pub async fn count_rows(&self, table: Table) -> Result<usize> {
self.dataset(table)
.await?
.count_rows(None)
.await
.map_err(Into::into)
}
pub async fn collect_ids(&self, table: Table) -> Result<std::collections::HashSet<String>> {
let batch = self.scan_batch(table, None, &["id"]).await?;
let ids = batch
.column_by_name("id")
.context("scan projection dropped the id column")?
.as_any()
.downcast_ref::<StringArray>()
.context("id column is not Utf8")?;
Ok(ids.iter().flatten().map(str::to_owned).collect())
}
#[cfg(test)]
pub(crate) async fn messages_index_names(&self) -> Result<Vec<String>> {
let dataset = self.dataset(Table::Messages).await?;
let indices = dataset.load_indices().await?;
Ok(indices.iter().map(|index| index.name.clone()).collect())
}
pub(crate) async fn messages_has_index(&self, name: &str) -> Result<bool> {
let dataset = self.dataset(Table::Messages).await?;
let indices = dataset.load_indices().await?;
Ok(indices.iter().any(|index| index.name == name))
}
pub(crate) async fn messages_fast_search_ready(&self, name: &str) -> Result<bool> {
let dataset = self.dataset(Table::Messages).await?;
if !dataset
.load_indices()
.await?
.iter()
.any(|index| index.name == name)
{
return Ok(false);
}
let unindexed = dataset
.unindexed_fragments(name)
.await
.with_context(|| format!("unindexed_fragments failed for {name}"))?;
Ok(unindexed.is_empty())
}
pub(crate) async fn prune_index_cache(&self, cache_dir: &std::path::Path) {
if config::is_local(&self.location) {
return;
}
let root = cache_dir.join(store_key(&self.location)).join("indices");
if !root.exists() {
return;
}
let mut keep = std::collections::HashSet::new();
for table in [Table::Sessions, Table::Messages, Table::Parts] {
let Ok(dataset) = self.dataset(table).await else {
return;
};
let Ok(indices) = dataset.load_indices().await else {
return;
};
keep.extend(indices.iter().map(|index| index.uuid.to_string()));
}
prune_stale_uuid_dirs(&root, &keep);
}
pub(crate) async fn unindexed_row_count(
&self,
table: Table,
index_name: &str,
) -> Result<usize> {
let dataset = self.dataset(table).await?;
let fragments = dataset
.unindexed_fragments(index_name)
.await
.with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
Ok(fragments
.iter()
.map(|fragment| fragment.num_rows().unwrap_or(0))
.sum())
}
pub(crate) async fn find_index_owner(&self, name: &str) -> Result<Option<Table>> {
let list = |table: Table| async move {
let dataset = self.dataset(table).await?;
let names: Vec<String> = dataset
.load_indices()
.await
.with_context(|| format!("load_indices failed for {}", table.label()))?
.iter()
.map(|index| index.name.clone())
.collect();
Ok::<_, anyhow::Error>(names)
};
let (sessions, messages, parts) = tokio::try_join!(
list(Table::Sessions),
list(Table::Messages),
list(Table::Parts),
)?;
for (table, names) in [
(Table::Sessions, sessions),
(Table::Messages, messages),
(Table::Parts, parts),
] {
if names.iter().any(|n| n == name) {
return Ok(Some(table));
}
}
Ok(None)
}
pub(crate) async fn drop_index(&self, table: Table, name: &str) -> Result<()> {
let mut guard = self.cached(table).await?.lock().await;
let mut dataset = guard.latest().await?;
dataset
.drop_index(name)
.await
.with_context(|| format!("drop_index({name}) failed for {}", table.label()))?;
guard.replace(dataset);
Ok(())
}
async fn table_location(&self, table_name: &str) -> Result<String> {
let request = DescribeTableRequest {
id: Some(self.nm_ident.as_table_id(table_name)),
..Default::default()
};
let response = self
.nm
.describe_table(request)
.await
.with_context(|| format!("failed to describe table {table_name}"))?;
response
.location
.with_context(|| format!("namespace returned no location for table {table_name}"))
}
pub async fn initialized(&self) -> Result<bool> {
let request = DescribeTableRequest {
id: Some(self.nm_ident.as_table_id(sessions::PARTS)),
..Default::default()
};
match self.nm.describe_table(request).await {
Ok(_) => Ok(true),
Err(error) if is_namespace_error_code(&error, ErrorCode::TableNotFound) => Ok(false),
Err(error) => {
Err(anyhow::Error::from(error)).context("failed to probe table existence")
}
}
}
pub async fn table_sizes(&self) -> Result<TableSizes> {
let registry = Arc::new(ObjectStoreRegistry::default());
let params = self.object_store_params();
let sessions = self
.listed_size(
®istry,
¶ms,
&self.table_location(sessions::SESSIONS).await?,
)
.await?;
let messages = self
.listed_size(
®istry,
¶ms,
&self.table_location(sessions::MESSAGES).await?,
)
.await?;
let parts = self
.listed_size(
®istry,
¶ms,
&self.table_location(sessions::PARTS).await?,
)
.await?;
let root_total = self
.listed_size(®istry, ¶ms, self.location.as_str())
.await?;
let other = root_total.saturating_sub(sessions + messages + parts);
let sessions_data = self
.data_liveness(®istry, ¶ms, Table::Sessions, sessions::SESSIONS)
.await?;
let messages_data = self
.data_liveness(®istry, ¶ms, Table::Messages, sessions::MESSAGES)
.await?;
let parts_data = self
.data_liveness(®istry, ¶ms, Table::Parts, sessions::PARTS)
.await?;
Ok(TableSizes {
sessions,
messages,
parts,
other,
sessions_data,
messages_data,
parts_data,
})
}
async fn data_liveness(
&self,
registry: &Arc<ObjectStoreRegistry>,
params: &ObjectStoreParams,
table: Table,
table_name: &str,
) -> Result<DataLiveness> {
let location = self.table_location(table_name).await?;
let data_dir = format!("{}/data", location.trim_end_matches('/'));
let on_disk = self.listed_size(registry, params, &data_dir).await?;
let dataset = self.dataset(table).await?;
let live = dataset
.get_fragments()
.iter()
.try_fold(0u64, |total, fragment| {
Some(total + fragment_bytes(fragment.metadata())?)
});
Ok(DataLiveness { on_disk, live })
}
async fn listed_size(
&self,
registry: &Arc<ObjectStoreRegistry>,
params: &ObjectStoreParams,
uri: &str,
) -> Result<u64> {
let (store, base) = ObjectStore::from_uri_and_params(registry.clone(), uri, params)
.await
.with_context(|| format!("failed to open object store for {uri}"))?;
let mut listing = store.list(Some(base));
let mut total = 0u64;
while let Some(meta) = listing.next().await {
let meta = meta.with_context(|| format!("listing {uri} failed"))?;
total += meta.size;
}
Ok(total)
}
async fn cached(&self, table: Table) -> Result<&Mutex<CachedDataset>> {
match table {
Table::Sessions => self.sessions_cached().await,
Table::Messages => Ok(&self.datasets.messages),
Table::Parts => self.parts_cached().await,
}
}
async fn sessions_cached(&self) -> Result<&Mutex<CachedDataset>> {
self.lazy_cached(
&self.datasets.sessions,
sessions::SESSIONS,
sessions::session_schema,
)
.await
}
async fn parts_cached(&self) -> Result<&Mutex<CachedDataset>> {
self.lazy_cached(&self.datasets.parts, sessions::PARTS, sessions::part_schema)
.await
}
async fn lazy_cached<'a>(
&self,
cell: &'a OnceCell<Mutex<CachedDataset>>,
table_name: &str,
schema: fn() -> lance::deps::arrow_schema::SchemaRef,
) -> Result<&'a Mutex<CachedDataset>> {
cell.get_or_try_init(|| async {
let dataset = open_or_create_via_ns(
&self.nm,
&self.nm_ident,
table_name,
schema(),
&self.session,
&self.storage_options,
self.store_wrapper.clone(),
)
.await?;
Ok::<_, anyhow::Error>(Mutex::new(CachedDataset::new(
dataset,
self.lazy_refresh_after,
)))
})
.await
}
async fn retry_lance<T, Fut, Op>(&self, label: &str, operation: Op) -> Result<T>
where
Fut: std::future::Future<Output = Result<T>>,
Op: FnMut() -> Fut,
{
self.retry_lance_filtered(label, |_| true, operation).await
}
async fn retry_lance_filtered<T, Fut, Op, R>(
&self,
label: &str,
should_retry: R,
mut operation: Op,
) -> Result<T>
where
Fut: std::future::Future<Output = Result<T>>,
Op: FnMut() -> Fut,
R: Fn(&anyhow::Error) -> bool,
{
let mut attempt = 0u8;
loop {
attempt = attempt.saturating_add(1);
match operation().await {
Ok(value) => return Ok(value),
Err(error) if attempt < self.retry.attempts && should_retry(&error) => {
let backoff = self.backoff(attempt);
let error_chain = format!("{error:#}");
tracing::warn!(
label,
attempt,
?backoff,
error = %error_chain,
"retrying Lance operation"
);
tokio::time::sleep(backoff).await;
}
Err(error) => {
let error_chain = format!("{error:#}");
tracing::warn!(
label,
attempt,
error = %error_chain,
"Lance operation exhausted retries"
);
if is_commit_conflict(&error) {
return Err(error.context(ConflictExhausted { attempts: attempt }));
}
return Err(error);
}
}
}
}
fn backoff(&self, attempt: u8) -> Duration {
let shift = u32::from(attempt.saturating_sub(1));
let multiplier = 1u32.checked_shl(shift).unwrap_or(u32::MAX);
let base = self.retry.initial_backoff.saturating_mul(multiplier);
let factor = (1.0 + self.retry.jitter * (fastrand::f64() * 2.0 - 1.0)).max(0.0);
base.mul_f64(factor).min(self.retry.max_backoff)
}
}
async fn optimize_table_compact(
dataset: &mut Dataset,
table: Table,
progress: Option<&OptimizeProgressFn>,
policy: &MaintenancePolicy,
) -> Result<()> {
let stats: Vec<FragmentStat> = dataset
.get_fragments()
.iter()
.map(|fragment| fragment_stat(fragment.metadata()))
.collect();
let compaction = CompactionOptions {
target_rows_per_fragment: derived_target_rows(&stats),
max_bytes_per_file: Some(TARGET_FRAGMENT_BYTES as usize),
defer_index_remap: false,
compaction_mode: Some(CompactionMode::TryBinaryCopy),
..CompactionOptions::default()
};
let mut plan = plan_compaction(dataset, &compaction).await?;
if policy.compaction_fragment_cap > 0 {
let max_bytes_per_file = compaction
.max_bytes_per_file
.and_then(|bytes| u64::try_from(bytes).ok())
.unwrap_or_default();
plan.tasks.retain(|task| {
let task_stats: Vec<FragmentStat> = task.fragments.iter().map(fragment_stat).collect();
let reason = task_veto_reason(
&task_stats,
policy.compaction_fragment_cap,
compaction.materialize_deletions_threshold,
compaction.target_rows_per_fragment,
max_bytes_per_file,
);
if let Some(reason) = reason {
tracing::debug!(
target: "pond::perf",
table = table.as_str(),
fragments = task_stats.len(),
reason,
"compaction task vetoed",
);
}
reason.is_none()
});
}
if plan.tasks.is_empty() {
tracing::debug!(
target: "pond::perf",
table = table.as_str(),
"compaction skipped: no task to run",
);
} else {
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::Compact,
detail: None,
},
);
let started = Instant::now();
let mut completed = Vec::with_capacity(plan.tasks.len());
for task in plan.compaction_tasks() {
completed.push(task.execute(dataset).await?);
}
commit_compaction(
dataset,
completed,
Arc::new(DatasetIndexRemapperOptions::default()),
&compaction,
)
.await?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::Compact,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
}
if cleanup_due(dataset.version_id(), policy.cleanup_interval) {
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::Cleanup,
detail: None,
},
);
let started = Instant::now();
dataset
.cleanup_old_versions(policy.cleanup_older_than, Some(false), Some(false))
.await
.context("cleanup_old_versions failed during index optimize")?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::Cleanup,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
}
Ok(())
}
fn cleanup_due(version: u64, interval: u64) -> bool {
interval <= 1 || version.is_multiple_of(interval)
}
async fn optimize_table_indices(
dataset: &mut Dataset,
intents: &[IndexIntent],
table: Table,
progress: Option<&OptimizeProgressFn>,
folds: FoldThresholds,
) -> Result<bool> {
let existing = dataset.load_indices().await?;
let existing_names: std::collections::HashSet<String> =
existing.iter().map(|index| index.name.clone()).collect();
let mut append_indices: Vec<String> = Vec::new();
let mut did_work = false;
for intent in intents {
let exists = existing_names.contains(intent.name);
if !exists {
if !intent.trigger.should_create(dataset).await? {
continue;
}
let params = intent.params.build(dataset).await?;
let index_type = intent.params.index_type();
tracing::info!(
index = intent.name,
column = intent.column,
"creating Lance index (trigger fired)",
);
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::IndexCreate,
detail: Some(intent.name.to_owned()),
},
);
let started = Instant::now();
dataset
.create_index_builder(&[intent.column], index_type, params.as_ref())
.name(intent.name.to_owned())
.replace(false)
.progress(lance_progress(progress, table, intent.name))
.await
.with_context(|| format!("failed to create index {}", intent.name))?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::IndexCreate,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
did_work = true;
continue;
}
let unindexed = dataset.unindexed_fragments(intent.name).await?;
if unindexed.is_empty() {
continue;
}
let tail_rows: usize = unindexed
.iter()
.map(|fragment| fragment.num_rows().unwrap_or(0))
.sum();
let fold_threshold = match intent.params {
IndexParamsKind::Scalar(_) => folds.scalar,
IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. } => folds.index,
};
if fold_threshold > 0 && tail_rows < fold_threshold {
tracing::debug!(
target: "pond::perf",
index = intent.name,
tail_rows,
threshold = fold_threshold,
"deferring index fold (unindexed tail below threshold)",
);
continue;
}
if matches!(
intent.params,
IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. }
) && !column_has_values(dataset, intent.presence_column(), &unindexed).await?
{
tracing::debug!(
target: "pond::perf",
index = intent.name,
tail_rows,
"skipping content index fold (tail has no indexable values)",
);
continue;
}
append_indices.push(intent.name.to_owned());
}
if !append_indices.is_empty() {
let segment_count = |name: &str| {
existing
.iter()
.filter(|index| index.name.as_str() == name)
.count()
};
let (consolidate, to_append): (Vec<String>, Vec<String>) = append_indices
.iter()
.cloned()
.partition(|name| segment_count(name) >= DELTA_MERGE_THRESHOLD);
let mut fts_rebuilds: Vec<&IndexIntent> = Vec::new();
let mut to_merge: Vec<String> = Vec::new();
for name in consolidate {
let fts_intent = intents.iter().find(|intent| {
intent.name == name && matches!(intent.params, IndexParamsKind::InvertedFtsWord)
});
match fts_intent {
Some(intent) => fts_rebuilds.push(intent),
None => to_merge.push(name),
}
}
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::IndexAppend,
detail: Some(append_indices.join(", ")),
},
);
let started = Instant::now();
if !to_append.is_empty() {
dataset
.optimize_indices(&OptimizeOptions::append().index_names(to_append))
.await
.context("optimize_indices(append) failed during index optimize")?;
}
if !to_merge.is_empty() {
dataset
.optimize_indices(
&OptimizeOptions::merge(DELTA_MERGE_THRESHOLD).index_names(to_merge),
)
.await
.context("optimize_indices(merge) failed during index optimize")?;
}
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::IndexAppend,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
for intent in &fts_rebuilds {
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::IndexRebuild,
detail: Some(intent.name.to_owned()),
},
);
let rebuild_started = Instant::now();
rebuild_index(dataset, intent, progress, table).await?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::IndexRebuild,
elapsed_ms: rebuild_started.elapsed().as_millis() as u64,
},
);
}
tracing::debug!(
target: "pond::perf",
indices = ?append_indices,
rebuilt = ?fts_rebuilds,
"folded trailing fragments into indices",
);
did_work = true;
}
Ok(did_work)
}
fn non_null_scanner(
dataset: &Dataset,
column: &'static str,
fragments: &[lance::table::format::Fragment],
) -> Result<lance::dataset::scanner::Scanner> {
let mut scanner = dataset.scan();
scanner.with_fragments(fragments.to_vec());
scanner.filter(&Predicate::IsNotNull(column).to_lance())?;
Ok(scanner)
}
async fn column_has_values(
dataset: &Dataset,
column: &'static str,
fragments: &[lance::table::format::Fragment],
) -> Result<bool> {
let mut scanner = non_null_scanner(dataset, column, fragments)?;
scanner.project(&[column])?;
scanner.limit(Some(1), None)?;
let batch = scanner
.try_into_batch()
.await
.with_context(|| format!("non-null probe on {column} failed"))?;
Ok(batch.num_rows() > 0)
}
async fn column_value_count(
dataset: &Dataset,
column: &'static str,
fragments: &[lance::table::format::Fragment],
) -> Result<usize> {
let count = non_null_scanner(dataset, column, fragments)?
.count_rows()
.await
.with_context(|| format!("non-null count on {column} failed"))?;
Ok(count as usize)
}
async fn rebuild_index(
dataset: &mut Dataset,
intent: &IndexIntent,
progress: Option<&OptimizeProgressFn>,
table: Table,
) -> Result<()> {
if !intent.trigger.should_create(dataset).await? {
return Ok(());
}
let params = intent.params.build(dataset).await?;
dataset
.create_index_builder(
&[intent.column],
intent.params.index_type(),
params.as_ref(),
)
.name(intent.name.to_owned())
.replace(true)
.progress(lance_progress(progress, table, intent.name))
.await
.with_context(|| format!("failed to rebuild index {}", intent.name))?;
Ok(())
}
async fn index_status(
table: Table,
dataset: &Dataset,
intents: &[IndexIntent],
indexable_only: bool,
) -> Result<Vec<IndexStatus>> {
let existing = dataset.load_indices().await?;
let existing_names: std::collections::HashSet<String> =
existing.iter().map(|index| index.name.clone()).collect();
let total_fragments = dataset.get_fragments().len();
let total_rows = dataset.count_rows(None).await?;
let mut statuses = Vec::with_capacity(intents.len());
for intent in intents {
let exists = existing_names.contains(intent.name);
if !exists {
statuses.push(IndexStatus {
table,
intent_name: intent.name.to_owned(),
fragments_covered: 0,
unindexed_fragments: total_fragments,
unindexed_rows: total_rows,
exists,
});
continue;
}
let unindexed = dataset
.unindexed_fragments(intent.name)
.await
.with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
let unindexed_fragments = unindexed.len();
let mut unindexed_rows: usize = unindexed
.iter()
.map(|fragment| fragment.num_rows().unwrap_or(0))
.sum();
if indexable_only
&& unindexed_rows > 0
&& matches!(
intent.params,
IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. }
)
{
unindexed_rows =
column_value_count(dataset, intent.presence_column(), &unindexed).await?;
}
statuses.push(IndexStatus {
table,
intent_name: intent.name.to_owned(),
fragments_covered: total_fragments.saturating_sub(unindexed_fragments),
unindexed_fragments,
unindexed_rows,
exists,
});
}
Ok(statuses)
}
pub mod io_trace {
use lance_io::utils::tracking_store::{IOTracker, IoStats};
use std::sync::{Arc, OnceLock};
static TRACKER: OnceLock<IOTracker> = OnceLock::new();
pub fn enable() {
let _ = TRACKER.set(IOTracker::default());
}
pub(super) fn wrapper() -> Option<Arc<IOTracker>> {
TRACKER.get().map(|tracker| Arc::new(tracker.clone()))
}
pub fn take() -> Option<IoStats> {
TRACKER.get().map(IOTracker::incremental_stats)
}
}
pub mod index_cache {
use object_store::local::LocalFileSystem;
use object_store::path::Path as ObjPath;
use object_store::{
CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result as OsResult,
};
use std::collections::HashMap;
use std::ops::Range;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use bytes::Bytes;
use futures::stream::BoxStream;
use lance_io::object_store::WrappingObjectStore;
fn is_index_path(location: &ObjPath) -> bool {
AsRef::<str>::as_ref(location).contains("_indices/")
}
fn local_opts(options: &GetOptions) -> GetOptions {
GetOptions {
range: options.range.clone(),
head: options.head,
..Default::default()
}
}
#[derive(Debug)]
pub struct IndexDiskCache {
local: Arc<LocalFileSystem>,
inflight: Arc<Mutex<HashMap<ObjPath, Arc<tokio::sync::Mutex<()>>>>>,
}
impl IndexDiskCache {
pub fn new(root: PathBuf) -> std::io::Result<Self> {
std::fs::create_dir_all(&root)?;
Ok(Self {
local: Arc::new(LocalFileSystem::new_with_prefix(&root)?),
inflight: Arc::new(Mutex::new(HashMap::new())),
})
}
}
impl WrappingObjectStore for IndexDiskCache {
fn wrap(&self, _store_prefix: &str, inner: Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
Arc::new(CachingStore {
inner,
local: self.local.clone(),
inflight: self.inflight.clone(),
})
}
}
#[derive(Debug)]
struct CachingStore {
inner: Arc<dyn ObjectStore>,
local: Arc<LocalFileSystem>,
inflight: Arc<Mutex<HashMap<ObjPath, Arc<tokio::sync::Mutex<()>>>>>,
}
impl std::fmt::Display for CachingStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "CachingStore({})", self.inner)
}
}
impl CachingStore {
fn flight_lock(&self, location: &ObjPath) -> Arc<tokio::sync::Mutex<()>> {
self.inflight
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.entry(location.clone())
.or_default()
.clone()
}
async fn populate_and_serve(
&self,
location: &ObjPath,
options: GetOptions,
) -> OsResult<GetResult> {
let lock = self.flight_lock(location);
let _guard = lock.lock().await;
let result = self.fetch_under_flight(location, options).await;
self.inflight
.lock()
.unwrap_or_else(|p| p.into_inner())
.remove(location);
result
}
async fn fetch_under_flight(
&self,
location: &ObjPath,
options: GetOptions,
) -> OsResult<GetResult> {
if let Ok(result) = self.local.get_opts(location, local_opts(&options)).await {
return Ok(result);
}
let bytes = self.inner.get(location).await?.bytes().await?;
if self
.local
.put(location, PutPayload::from_bytes(bytes))
.await
.is_ok()
&& let Ok(result) = self.local.get_opts(location, local_opts(&options)).await
{
return Ok(result);
}
self.inner.get_opts(location, options).await
}
}
#[async_trait::async_trait]
impl ObjectStore for CachingStore {
async fn get_opts(&self, location: &ObjPath, options: GetOptions) -> OsResult<GetResult> {
if !is_index_path(location) {
return self.inner.get_opts(location, options).await;
}
match self.local.get_opts(location, local_opts(&options)).await {
Ok(result) => Ok(result),
Err(object_store::Error::NotFound { .. }) => {
self.populate_and_serve(location, options).await
}
Err(_) => self.inner.get_opts(location, options).await,
}
}
async fn put_opts(
&self,
location: &ObjPath,
payload: PutPayload,
opts: PutOptions,
) -> OsResult<PutResult> {
self.inner.put_opts(location, payload, opts).await
}
async fn put_multipart_opts(
&self,
location: &ObjPath,
opts: PutMultipartOptions,
) -> OsResult<Box<dyn MultipartUpload>> {
self.inner.put_multipart_opts(location, opts).await
}
async fn get_ranges(
&self,
location: &ObjPath,
ranges: &[Range<u64>],
) -> OsResult<Vec<Bytes>> {
if is_index_path(location) {
let mut out = Vec::with_capacity(ranges.len());
for range in ranges {
let opts = GetOptions {
range: Some(range.clone().into()),
..Default::default()
};
out.push(self.get_opts(location, opts).await?.bytes().await?);
}
return Ok(out);
}
self.inner.get_ranges(location, ranges).await
}
fn delete_stream(
&self,
locations: BoxStream<'static, OsResult<ObjPath>>,
) -> BoxStream<'static, OsResult<ObjPath>> {
self.inner.delete_stream(locations)
}
fn list(&self, prefix: Option<&ObjPath>) -> BoxStream<'static, OsResult<ObjectMeta>> {
self.inner.list(prefix)
}
fn list_with_offset(
&self,
prefix: Option<&ObjPath>,
offset: &ObjPath,
) -> BoxStream<'static, OsResult<ObjectMeta>> {
self.inner.list_with_offset(prefix, offset)
}
async fn list_with_delimiter(&self, prefix: Option<&ObjPath>) -> OsResult<ListResult> {
self.inner.list_with_delimiter(prefix).await
}
async fn copy_opts(&self, from: &ObjPath, to: &ObjPath, opts: CopyOptions) -> OsResult<()> {
self.inner.copy_opts(from, to, opts).await
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
use object_store::memory::InMemory;
async fn read(store: &Arc<dyn ObjectStore>, path: &ObjPath) -> Option<Vec<u8>> {
store
.get(path)
.await
.ok()?
.bytes()
.await
.ok()
.map(|b| b.to_vec())
}
#[tokio::test]
async fn caches_index_files_and_passes_data_through() {
let temp = tempfile::tempdir().unwrap();
let inner: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
let index_path = ObjPath::from("d/messages.lance/_indices/uuid1/index.idx");
let data_path = ObjPath::from("d/messages.lance/data/x.lance");
inner
.put(&index_path, PutPayload::from_static(b"INDEX"))
.await
.unwrap();
inner
.put(&data_path, PutPayload::from_static(b"DATA"))
.await
.unwrap();
let cache = IndexDiskCache::new(temp.path().join("indices")).unwrap();
let store = cache.wrap("test", inner.clone());
assert_eq!(
read(&store, &index_path).await.as_deref(),
Some(&b"INDEX"[..])
);
assert_eq!(
read(&store, &data_path).await.as_deref(),
Some(&b"DATA"[..])
);
inner.delete(&index_path).await.unwrap();
inner.delete(&data_path).await.unwrap();
assert_eq!(
read(&store, &index_path).await.as_deref(),
Some(&b"INDEX"[..])
);
assert_eq!(read(&store, &data_path).await, None);
let slice = store.get_range(&index_path, 1..4).await.unwrap();
assert_eq!(slice.as_ref(), b"NDE");
}
}
}
pub mod durability {
use std::fs::File;
#[cfg(unix)]
use std::io::ErrorKind;
use std::ops::Range;
use std::path::Path as FsPath;
use std::sync::Arc;
use bytes::Bytes;
use futures::stream::BoxStream;
use lance_io::object_store::WrappingObjectStore;
use object_store::path::Path as ObjPath;
use object_store::{
CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OsResult,
UploadPart,
};
#[derive(Debug, thiserror::Error)]
#[error("{op} {path}")]
struct FsyncFailed {
op: &'static str,
path: String,
source: std::io::Error,
}
fn durability_error(
op: &'static str,
path: &FsPath,
source: std::io::Error,
) -> object_store::Error {
object_store::Error::Generic {
store: "fsync-durability",
source: Box::new(FsyncFailed {
op,
path: path.display().to_string(),
source,
}),
}
}
fn terminal(error: object_store::Error) -> object_store::Error {
object_store::Error::Generic {
store: "fsync-durability",
source: format!("{error}").into(),
}
}
fn sync_file_and_parent(location: &ObjPath) -> OsResult<()> {
let local = lance_io::local::to_local_path(location);
let path = FsPath::new(&local);
#[cfg(windows)]
let opened = File::options().write(true).open(path);
#[cfg(not(windows))]
let opened = File::open(path);
let file = opened.map_err(|e| durability_error("open for fsync", path, e))?;
file.sync_all()
.map_err(|e| durability_error("fsync", path, e))?;
#[cfg(unix)]
if let Some(parent) = path.parent() {
match File::open(parent) {
Ok(dir) => dir
.sync_all()
.map_err(|e| durability_error("fsync dir", parent, e))?,
Err(e) if e.kind() == ErrorKind::NotFound => {}
Err(e) => return Err(durability_error("open dir for fsync", parent, e)),
}
}
Ok(())
}
#[derive(Debug)]
pub struct FsyncOnWrite;
impl WrappingObjectStore for FsyncOnWrite {
fn wrap(&self, _store_prefix: &str, inner: Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
Arc::new(FsyncStore { inner })
}
}
#[derive(Debug)]
struct FsyncStore {
inner: Arc<dyn ObjectStore>,
}
impl std::fmt::Display for FsyncStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "FsyncStore({})", self.inner)
}
}
#[async_trait::async_trait]
impl ObjectStore for FsyncStore {
async fn put_opts(
&self,
location: &ObjPath,
payload: PutPayload,
opts: PutOptions,
) -> OsResult<PutResult> {
let result = self.inner.put_opts(location, payload, opts).await?;
sync_file_and_parent(location)?;
Ok(result)
}
async fn put_multipart_opts(
&self,
location: &ObjPath,
opts: PutMultipartOptions,
) -> OsResult<Box<dyn MultipartUpload>> {
let upload = self.inner.put_multipart_opts(location, opts).await?;
Ok(Box::new(FsyncUpload {
inner: upload,
location: location.clone(),
}))
}
async fn get_opts(&self, location: &ObjPath, options: GetOptions) -> OsResult<GetResult> {
self.inner.get_opts(location, options).await
}
async fn get_ranges(
&self,
location: &ObjPath,
ranges: &[Range<u64>],
) -> OsResult<Vec<Bytes>> {
self.inner.get_ranges(location, ranges).await
}
fn delete_stream(
&self,
locations: BoxStream<'static, OsResult<ObjPath>>,
) -> BoxStream<'static, OsResult<ObjPath>> {
self.inner.delete_stream(locations)
}
fn list(&self, prefix: Option<&ObjPath>) -> BoxStream<'static, OsResult<ObjectMeta>> {
self.inner.list(prefix)
}
fn list_with_offset(
&self,
prefix: Option<&ObjPath>,
offset: &ObjPath,
) -> BoxStream<'static, OsResult<ObjectMeta>> {
self.inner.list_with_offset(prefix, offset)
}
async fn list_with_delimiter(&self, prefix: Option<&ObjPath>) -> OsResult<ListResult> {
self.inner.list_with_delimiter(prefix).await
}
async fn copy_opts(&self, from: &ObjPath, to: &ObjPath, opts: CopyOptions) -> OsResult<()> {
self.inner.copy_opts(from, to, opts).await?;
sync_file_and_parent(to)?;
Ok(())
}
async fn rename_opts(
&self,
from: &ObjPath,
to: &ObjPath,
opts: RenameOptions,
) -> OsResult<()> {
self.inner.rename_opts(from, to, opts).await?;
sync_file_and_parent(to).map_err(terminal)?;
Ok(())
}
}
#[derive(Debug)]
struct FsyncUpload {
inner: Box<dyn MultipartUpload>,
location: ObjPath,
}
#[async_trait::async_trait]
impl MultipartUpload for FsyncUpload {
fn put_part(&mut self, data: PutPayload) -> UploadPart {
self.inner.put_part(data)
}
async fn complete(&mut self) -> OsResult<PutResult> {
let result = self.inner.complete().await?;
sync_file_and_parent(&self.location)?;
Ok(result)
}
async fn abort(&mut self) -> OsResult<()> {
self.inner.abort().await
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
use object_store::ObjectStoreExt;
fn wrapped() -> Arc<dyn ObjectStore> {
let inner: Arc<dyn ObjectStore> = Arc::new(object_store::local::LocalFileSystem::new());
FsyncOnWrite.wrap("test", inner)
}
fn obj_path(root: &FsPath, name: &str) -> ObjPath {
ObjPath::from_absolute_path(root.join(name)).unwrap()
}
#[test]
fn a_contended_flush_is_retryable_unless_the_write_already_published() {
let sharing_violation = || std::io::Error::from_raw_os_error(32);
let classify = |error: object_store::Error| {
crate::substrate::is_transient_sharing_violation(&anyhow::Error::from(
lance::Error::from(error),
))
};
assert!(classify(durability_error(
"open for fsync",
FsPath::new("/store/data.lance"),
sharing_violation(),
)));
assert!(!classify(terminal(durability_error(
"open for fsync",
FsPath::new("/store/data.lance"),
sharing_violation(),
))));
}
#[tokio::test]
async fn put_through_wrapper_round_trips_and_lands_on_disk() {
let temp = tempfile::tempdir().unwrap();
let store = wrapped();
let path = obj_path(temp.path(), "sub/dir/manifest");
store
.put(&path, PutPayload::from_static(b"DURABLE"))
.await
.unwrap();
let got = store.get(&path).await.unwrap().bytes().await.unwrap();
assert_eq!(got.as_ref(), b"DURABLE");
assert_eq!(
std::fs::read(temp.path().join("sub/dir/manifest")).unwrap(),
b"DURABLE",
);
}
#[tokio::test]
async fn multipart_through_wrapper_completes_and_round_trips() {
let temp = tempfile::tempdir().unwrap();
let store = wrapped();
let path = obj_path(temp.path(), "data/part.lance");
let mut upload = store.put_multipart(&path).await.unwrap();
upload
.put_part(PutPayload::from_static(b"AB"))
.await
.unwrap();
upload
.put_part(PutPayload::from_static(b"CD"))
.await
.unwrap();
upload.complete().await.unwrap();
let got = store.get(&path).await.unwrap().bytes().await.unwrap();
assert_eq!(got.as_ref(), b"ABCD");
}
}
}
pub fn store_key(location: &Url) -> String {
blake3::hash(location.as_str().as_bytes()).to_hex()[..16].to_owned()
}
fn prune_stale_uuid_dirs(dir: &std::path::Path, keep: &std::collections::HashSet<String>) {
let Ok(entries) = std::fs::read_dir(dir) else {
return;
};
for entry in entries.flatten() {
let path = entry.path();
if !path.is_dir() {
continue;
}
if entry.file_name() == "_indices" {
let Ok(children) = std::fs::read_dir(&path) else {
continue;
};
for child in children.flatten() {
if child.path().is_dir()
&& !keep.contains(child.file_name().to_string_lossy().as_ref())
&& let Err(error) = std::fs::remove_dir_all(child.path())
{
tracing::debug!(
%error,
path = %child.path().display(),
"stale index cache dir not reclaimed; retried on the next prune",
);
}
}
} else {
prune_stale_uuid_dirs(&path, keep);
}
}
}
fn store_wrapper(
location: &Url,
index_cache_dir: Option<&std::path::Path>,
) -> Option<Arc<dyn WrappingObjectStore>> {
let mut wrappers: Vec<Arc<dyn WrappingObjectStore>> = Vec::new();
if config::is_local(location) {
wrappers.push(Arc::new(durability::FsyncOnWrite));
}
if let Some(dir) = index_cache_dir
&& !config::is_local(location)
{
let root = dir.join(store_key(location)).join("indices");
match index_cache::IndexDiskCache::new(root) {
Ok(cache) => wrappers.push(Arc::new(cache)),
Err(error) => tracing::warn!(%error, "index disk cache disabled; reads hit the store"),
}
}
if let Some(tracker) = io_trace::wrapper() {
wrappers.push(tracker);
}
match wrappers.len() {
0 => None,
1 => Some(wrappers.remove(0)),
_ => Some(Arc::new(ChainedWrappingObjectStore::new(wrappers))),
}
}
async fn open_or_create_via_ns(
nm: &Arc<dyn LanceNamespace>,
nm_ident: &NamespaceIdent,
table_name: &str,
schema: lance::deps::arrow_schema::SchemaRef,
session: &Arc<Session>,
storage_options: &HashMap<String, String>,
wrapper: Option<Arc<dyn WrappingObjectStore>>,
) -> Result<Dataset> {
let table_id = nm_ident.as_table_id(table_name);
let request = DescribeTableRequest {
id: Some(table_id.clone()),
..Default::default()
};
match nm.describe_table(request).await {
Ok(response) => {
let location = response.location.with_context(|| {
format!("namespace returned no location for table {table_name}")
})?;
let builder = apply_open_params(
DatasetBuilder::from_uri(&location).with_session(session.clone()),
&wrapper,
storage_options,
);
let mut dataset = match builder.load().await {
Ok(dataset) => dataset,
Err(load_error) => {
let load_error = anyhow::Error::new(load_error)
.context(format!("failed to open table {table_name}"));
match config::local_path(&uri_to_url(&location)?) {
Some(table_root) => {
heal_local_dataset(
&location,
&table_root,
table_name,
session,
storage_options,
&wrapper,
load_error,
)
.await?
}
None => return Err(load_error),
}
}
};
ensure_current_schema(&mut dataset, schema.as_ref(), table_name).await?;
return Ok(dataset);
}
Err(error) => match &error {
error if is_namespace_error_code(error, ErrorCode::TableNotFound) => {
}
_ => {
return Err(anyhow::Error::from(error))
.with_context(|| format!("failed to describe table {table_name}"));
}
},
}
let mut write_params = sessions::write_params_for_create();
write_params.session = Some(session.clone());
write_params.mode = WriteMode::Create;
if wrapper.is_some() || !storage_options.is_empty() {
write_params.store_params = Some(ObjectStoreParams {
object_store_wrapper: wrapper.clone(),
storage_options_accessor: (!storage_options.is_empty()).then(|| {
Arc::new(StorageOptionsAccessor::with_static_options(
storage_options.clone(),
))
}),
..Default::default()
});
}
let reader = sessions::empty_reader(schema)?;
Dataset::write_into_namespace(reader, nm.clone(), table_id, Some(write_params))
.await
.with_context(|| format!("failed to create table {table_name}"))
}
fn apply_open_params(
builder: DatasetBuilder,
wrapper: &Option<Arc<dyn WrappingObjectStore>>,
storage_options: &HashMap<String, String>,
) -> DatasetBuilder {
match wrapper {
Some(wrapper) => builder.with_store_params(ObjectStoreParams {
object_store_wrapper: Some(wrapper.clone()),
storage_options_accessor: (!storage_options.is_empty()).then(|| {
Arc::new(StorageOptionsAccessor::with_static_options(
storage_options.clone(),
))
}),
..Default::default()
}),
None if !storage_options.is_empty() => {
builder.with_storage_options(storage_options.clone())
}
None => builder,
}
}
const VERSIONS_DIR_NAME: &str = "_versions";
const HEAL_MAX_PROBES: usize = 32;
fn parse_manifest_version(filename: &str) -> Option<u64> {
if filename.starts_with('d') {
return None;
}
let stem = filename.strip_suffix(".manifest")?;
if stem.len() == 20 {
stem.parse::<u64>().ok().map(|inverted| u64::MAX - inverted)
} else {
stem.parse::<u64>().ok()
}
}
async fn scan_verify_version(
table_uri: &str,
version: u64,
session: &Arc<Session>,
storage_options: &HashMap<String, String>,
wrapper: &Option<Arc<dyn WrappingObjectStore>>,
) -> Result<()> {
let builder = apply_open_params(
DatasetBuilder::from_uri(table_uri)
.with_session(session.clone())
.with_version(version),
wrapper,
storage_options,
);
let dataset = builder.load().await?;
let scanner = dataset.scan();
let mut stream = scanner.try_into_stream().await?;
while let Some(batch) = stream.next().await {
batch?;
}
Ok(())
}
async fn heal_local_dataset(
table_uri: &str,
table_root: &std::path::Path,
table_name: &str,
session: &Arc<Session>,
storage_options: &HashMap<String, String>,
wrapper: &Option<Arc<dyn WrappingObjectStore>>,
load_error: anyhow::Error,
) -> Result<Dataset> {
let versions_dir = table_root.join(VERSIONS_DIR_NAME);
let entries = match std::fs::read_dir(&versions_dir) {
Ok(entries) => entries,
Err(_) => {
return Err(enriched_open_error(
table_name,
format!(
"open failed and no {} directory exists at {} - not a crash-damaged manifest",
VERSIONS_DIR_NAME,
versions_dir.display()
),
load_error,
));
}
};
let mut manifests: Vec<(u64, PathBuf)> = Vec::new();
for entry in entries {
let entry = entry.with_context(|| format!("listing {}", versions_dir.display()))?;
if let Some(version) = parse_manifest_version(&entry.file_name().to_string_lossy()) {
manifests.push((version, entry.path()));
}
}
manifests.sort_by_key(|(version, _)| std::cmp::Reverse(*version));
let Some((_, newest_path)) = manifests.first().cloned() else {
return Err(enriched_open_error(
table_name,
format!(
"open failed and no manifest files exist under {} - not a crash-damaged manifest",
versions_dir.display()
),
load_error,
));
};
let newest_desc = || {
let name = newest_path
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_default();
let bytes = std::fs::metadata(&newest_path).map(|m| m.len()).ok();
match bytes {
Some(bytes) => format!("newest manifest {name} ({bytes} bytes)"),
None => format!("newest manifest {name}"),
}
};
let mut doomed: Vec<PathBuf> = Vec::new();
let mut landed: Option<u64> = None;
let mut probes = 0usize;
for (version, path) in &manifests {
let len = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
if len < 16 {
doomed.push(path.clone());
continue;
}
if probes >= HEAL_MAX_PROBES {
break;
}
probes += 1;
match scan_verify_version(table_uri, *version, session, storage_options, wrapper).await {
Ok(()) => {
landed = Some(*version);
break;
}
Err(_) => doomed.push(path.clone()),
}
}
let Some(landed_version) = landed else {
return Err(enriched_open_error(
table_name,
format!(
"{} is unreadable (interrupted commit during a hard host stop) and no older version passed a scan-verify probe; nothing was quarantined",
newest_desc()
),
load_error,
));
};
if doomed.is_empty() {
return Err(enriched_open_error(
table_name,
format!(
"open failed but the manifest head under {} is readable - not a crash-damaged manifest",
versions_dir.display()
),
load_error,
));
}
let mut quarantined: Vec<String> = Vec::new();
for path in &doomed {
let mut corrupt = path.clone().into_os_string();
corrupt.push(".corrupt");
match std::fs::rename(path, &corrupt) {
Ok(()) => {
if let Some(name) = path.file_name() {
quarantined.push(name.to_string_lossy().into_owned());
}
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
Err(error) => {
return Err(anyhow::Error::new(error).context(format!(
"table {table_name}: failed to quarantine corrupt manifest {}. If another pond process (or a scanner) has it open, close it and re-run - heal retries on the next open.",
path.display()
)));
}
}
}
let builder = apply_open_params(
DatasetBuilder::from_uri(table_uri).with_session(session.clone()),
wrapper,
storage_options,
);
let dataset = builder.load().await.with_context(|| {
format!(
"table {table_name}: open still failed after quarantining {} corrupt manifest(s); restore from a `pond copy` replica or re-run `pond init`",
quarantined.len()
)
})?;
tracing::warn!(
"pond self-healed local table {table_name}: quarantined {} unreadable manifest(s) ({}) to {VERSIONS_DIR_NAME}/*.corrupt and rolled back to version {landed_version}. The interrupted commit's rows are reconstructed on the next `pond sync` from source histories.",
quarantined.len(),
quarantined.join(", "),
);
Ok(dataset)
}
fn enriched_open_error(
table_name: &str,
finding: String,
load_error: anyhow::Error,
) -> anyhow::Error {
load_error.context(format!(
"table {table_name}: {finding}. Restore this store from a `pond copy` replica or re-run `pond init` to re-sync from source histories"
))
}
fn is_namespace_error_code(error: &lance::Error, code: ErrorCode) -> bool {
if !matches!(error, lance::Error::Namespace { .. }) {
return false;
}
std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |link| {
link.source()
})
.filter_map(|link| link.downcast_ref::<NamespaceError>())
.any(|inner| inner.code() == code)
}
fn scanner_with_prefilter(
dataset: &Dataset,
predicate: Option<&Predicate>,
) -> Result<lance::dataset::scanner::Scanner> {
let mut scanner = dataset.scan();
scanner.prefilter(true);
if let Some(predicate) = predicate {
let filter = predicate.to_lance();
if !filter.is_empty() {
scanner.filter(&filter)?;
}
}
Ok(scanner)
}
enum SchemaFit {
Match,
MissingNullable(Vec<lance::deps::arrow_schema::Field>),
UnknownExtra(Vec<String>),
}
fn classify_schema(
actual: &lance::deps::arrow_schema::Schema,
expected: &lance::deps::arrow_schema::Schema,
table_name: &str,
) -> Result<SchemaFit> {
use std::collections::BTreeSet;
let actual_names: BTreeSet<&str> = actual.fields().iter().map(|f| f.name().as_str()).collect();
let expected_names: BTreeSet<&str> = expected
.fields()
.iter()
.map(|f| f.name().as_str())
.collect();
let missing: Vec<_> = expected
.fields()
.iter()
.filter(|f| !actual_names.contains(f.name().as_str()))
.map(|f| f.as_ref().clone())
.collect();
let extra: Vec<String> = actual_names
.difference(&expected_names)
.map(|name| (*name).to_owned())
.collect();
match (missing.is_empty(), extra.is_empty()) {
(true, true) => Ok(SchemaFit::Match),
(false, true) if missing.iter().all(|f| f.is_nullable()) => {
Ok(SchemaFit::MissingNullable(missing))
}
(true, false) => Ok(SchemaFit::UnknownExtra(extra)),
_ => anyhow::bail!(
"table {table_name} has columns {actual_names:?} but this pond build expects \
{expected_names:?}, and the difference is not an additive nullable-column \
change this build can migrate - upgrade pond, or restore the store from a \
`pond copy` snapshot taken by the version that wrote it",
),
}
}
async fn ensure_current_schema(
dataset: &mut Dataset,
expected: &lance::deps::arrow_schema::Schema,
table_name: &str,
) -> Result<()> {
use lance::deps::arrow_schema::DataType;
const MAX_MIGRATION_ATTEMPTS: usize = 3;
for _ in 0..MAX_MIGRATION_ATTEMPTS {
let actual = lance::deps::arrow_schema::Schema::from(dataset.schema());
match classify_schema(&actual, expected, table_name)? {
SchemaFit::MissingNullable(missing) => {
backfill_missing_columns(dataset, table_name, missing).await?;
continue;
}
SchemaFit::Match => {}
SchemaFit::UnknownExtra(extra) => {
tracing::warn!(
table = table_name,
?extra,
"store carries columns unknown to this pond build (written by a newer \
version); reads proceed, writes need the newer pond",
);
}
}
for actual_field in actual.fields() {
let Some(expected_field) = expected.field_with_name(actual_field.name()).ok() else {
continue;
};
if let (
DataType::FixedSizeList(_, actual_dim),
DataType::FixedSizeList(_, expected_dim),
) = (actual_field.data_type(), expected_field.data_type())
&& actual_dim != expected_dim
{
tracing::warn!(
table = table_name,
column = actual_field.name(),
actual_dim,
expected_dim,
"embedding dimension differs from config; open proceeds because model swaps are operator-driven",
);
}
}
return Ok(());
}
anyhow::bail!(
"schema migration for table {table_name} did not converge after \
{MAX_MIGRATION_ATTEMPTS} attempts (a concurrent writer kept changing \
the schema); re-run once the other pond process finishes",
)
}
async fn backfill_missing_columns(
dataset: &mut Dataset,
table_name: &str,
missing: Vec<lance::deps::arrow_schema::Field>,
) -> Result<()> {
use lance::dataset::{BatchUDF, NewColumnTransform};
let names: Vec<&str> = missing.iter().map(|f| f.name().as_str()).collect();
let spec = sessions::column_backfill(table_name, &missing)?;
let _ = crate::output::line_err(&format!(
"migrating {table_name}: backfilling {names:?} from stored data (one-time, in place)...",
));
let started = std::time::Instant::now();
let mapper = spec.mapper;
let migration: std::pin::Pin<
Box<dyn std::future::Future<Output = lance::Result<()>> + Send + '_>,
> = Box::pin(dataset.add_columns(
NewColumnTransform::BatchUDF(BatchUDF {
mapper: Box::new(move |batch| {
mapper(batch).map_err(|error| lance::Error::io(format!("{error:#}")))
}),
output_schema: spec.output_schema,
result_checkpoint: None,
}),
Some(spec.read_columns),
None,
));
let result = migration.await;
match result {
Ok(()) => {
let _ = crate::output::line_err(&format!(
"migrated {table_name} in {:.1}s",
started.elapsed().as_secs_f64(),
));
Ok(())
}
Err(error) => {
let error = anyhow::Error::from(error);
if is_commit_conflict(&error) {
dataset.checkout_latest().await?;
Ok(())
} else {
Err(error).with_context(|| {
format!(
"schema backfill failed for {table_name}; the one-time migration \
writes new column files, so it needs write access to the store - \
re-run any pond command with write-capable credentials to complete it",
)
})
}
}
}
}
fn apply_remote_storage_defaults(options: &mut HashMap<String, String>) {
fn set_default(options: &mut HashMap<String, String>, aliases: &[&str], value: &str) {
if aliases
.iter()
.any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)))
{
return;
}
options.insert(aliases[0].to_owned(), value.to_owned());
}
set_default(options, &["pool_idle_timeout"], "300 seconds");
set_default(options, &["connect_timeout"], "10 seconds");
set_default(options, &["request_timeout"], "60 seconds");
let has_custom_endpoint = ["aws_endpoint", "endpoint"]
.iter()
.any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)));
if has_custom_endpoint {
set_default(
options,
&["aws_unsigned_payload", "unsigned_payload"],
"true",
);
}
}
fn quoted_string(value: &str) -> String {
format!("'{}'", value.replace('\'', "''"))
}
fn like_contains(value: &str) -> String {
let escaped = value
.replace('\\', "\\\\")
.replace('%', "\\%")
.replace('_', "\\_")
.replace('\'', "''");
format!("'%{escaped}%'")
}
#[cfg(test)]
mod tests {
#![allow(clippy::expect_used, clippy::unwrap_used)]
use super::*;
use tempfile::TempDir;
#[test]
fn a_contended_windows_commit_reads_through_the_whole_error_chain() {
let contended = |raw: i32| {
let source = object_store::Error::Generic {
store: "LocalFileSystem",
source: Box::new(std::io::Error::from_raw_os_error(raw)),
};
anyhow::Error::from(lance::Error::from(source)).context("commit sessions")
};
assert!(is_transient_sharing_violation(&contended(32)));
assert!(is_transient_sharing_violation(&contended(33)));
assert!(!is_transient_sharing_violation(&contended(5)));
assert!(!is_transient_sharing_violation(&anyhow::anyhow!(
"no io in this chain"
)));
}
#[test]
fn is_index_error_matches_lance_index_class_through_context() {
let index_fault = anyhow::Error::from(lance::Error::index(
"cannot merge inverted index segments with different posting tail codecs",
))
.context("optimize_indices(merge) failed during index optimize");
assert!(is_index_error(&index_fault));
let io_fault = anyhow::anyhow!("connection reset").context("optimize_indices failed");
assert!(!is_index_error(&io_fault));
}
#[test]
fn secret_command_with_space_round_trips() {
let result = run_secret_command("test", "field", "echo hello").unwrap();
assert_eq!(result, "hello");
}
#[test]
fn secret_command_with_space_and_quote_round_trips() {
let result = run_secret_command("test2", "q", r#"echo "hello world""#).unwrap();
#[cfg(windows)]
assert_eq!(
result, r#""hello world""#,
"cmd echo must preserve outer quotes"
);
#[cfg(not(windows))]
assert_eq!(result, "hello world", "sh must strip outer quotes");
}
#[test]
fn prune_keeps_live_uuid_dirs_and_drops_dead_ones() {
let temp = TempDir::new().unwrap();
let indices = temp.path().join("bkt/messages.lance/_indices");
for uuid in ["live", "dead"] {
std::fs::create_dir_all(indices.join(uuid)).unwrap();
std::fs::write(indices.join(uuid).join("index.idx"), b"x").unwrap();
}
let keep = std::collections::HashSet::from(["live".to_owned()]);
prune_stale_uuid_dirs(temp.path(), &keep);
assert!(indices.join("live").exists());
assert!(!indices.join("dead").exists());
}
#[test]
fn store_wrapper_present_for_local_absent_for_remote() {
let local = Url::parse("file:///tmp/pond-wrapper-test").unwrap();
assert!(store_wrapper(&local, None).is_some());
for remote in ["memory:///pond-wrapper-test", "s3://bucket/prefix"] {
let url = Url::parse(remote).unwrap();
assert!(
store_wrapper(&url, None).is_none(),
"remote store must not carry the fsync wrapper: {remote}",
);
}
}
fn set(scope: Option<&str>) -> CredsSet {
CredsSet {
scope: scope.map(str::to_owned),
access_key_id: Some("AKIA".to_owned()),
secret_access_key: Some("shh".to_owned()),
..CredsSet::default()
}
}
fn opts(resolved: &ResolvedStorage, key: &str) -> Option<String> {
resolved.options.get(key).cloned()
}
#[test]
fn storage_url_translation_table() {
let (local_input, local_expected) = if cfg!(windows) {
("C:\\srv\\pond", "file:///C:/srv/pond/")
} else {
("/srv/pond", "file:///srv/pond/")
};
let local = StorageUrl::parse(local_input).unwrap();
assert_eq!(local.lance_url().as_str(), local_expected);
assert!(local.is_local());
assert!(local.scheme_options.is_empty());
let aws = StorageUrl::parse("s3://bucket/prefix").unwrap();
assert_eq!(aws.lance_url().as_str(), "s3://bucket/prefix");
assert!(aws.scheme_options.is_empty());
let fat = StorageUrl::parse("s3+https://nbg1.example.com/my-pond/sub").unwrap();
assert_eq!(fat.lance_url().as_str(), "s3://my-pond/sub");
assert_eq!(
fat.scheme_options,
vec![
("allow_http", "false".to_owned()),
("virtual_hosted_style_request", "true".to_owned()),
("region", "us-east-1".to_owned()),
],
);
let resolved = fat.resolve(&BTreeMap::new()).unwrap();
assert_eq!(
opts(&resolved, "endpoint").as_deref(),
Some("https://my-pond.nbg1.example.com"),
);
assert_eq!(opts(&resolved, "region").as_deref(), Some("us-east-1"));
let plain = StorageUrl::parse("s3+http://127.0.0.1:9000/pond").unwrap();
assert_eq!(plain.lance_url().as_str(), "s3://pond/");
assert_eq!(plain.scheme_options[0], ("allow_http", "true".to_owned()));
assert_eq!(
plain.scheme_options[1],
("virtual_hosted_style_request", "false".to_owned()),
);
let resolved = plain.resolve(&BTreeMap::new()).unwrap();
assert_eq!(
opts(&resolved, "endpoint").as_deref(),
Some("http://127.0.0.1:9000"),
);
let mut pinned = BTreeMap::new();
pinned.insert(
"default".to_owned(),
CredsSet {
extra: [(
"endpoint".to_owned(),
"https://pinned.example.com".to_owned(),
)]
.into_iter()
.collect(),
..CredsSet::default()
},
);
let resolved = fat.resolve(&pinned).unwrap();
assert_eq!(
opts(&resolved, "endpoint").as_deref(),
Some("https://pinned.example.com"),
);
let gcs = StorageUrl::parse("gs://bucket/p").unwrap();
assert_eq!(gcs.lance_url().as_str(), "gs://bucket/p");
let azure = StorageUrl::parse("az://acct/container/p").unwrap();
assert_eq!(azure.lance_url().as_str(), "az://container/p");
assert_eq!(
azure.scheme_options,
vec![("account_name", "acct".to_owned())]
);
let shared = StorageUrl::parse("shared-memory://pond-test-x/").unwrap();
assert_eq!(shared.lance_url().as_str(), "shared-memory://pond-test-x/");
}
#[test]
fn storage_url_rejects_bad_shapes() {
let err = StorageUrl::parse("s3+https://user:pass@host/bucket")
.expect_err("userinfo must be rejected")
.to_string();
assert!(
err.contains("creds"),
"error must name the alternative: {err}"
);
assert!(StorageUrl::parse("s3+https://host").is_err());
assert!(StorageUrl::parse("az://acct").is_err());
let err = StorageUrl::parse("ftp://host/x")
.expect_err("ftp")
.to_string();
assert!(err.contains("s3+https"), "got: {err}");
let err = StorageUrl::parse("s3://b/p?regoin=x")
.expect_err("typo")
.to_string();
assert!(err.contains("regoin"), "got: {err}");
let err = StorageUrl::parse("memory://x?creds=y")
.expect_err("memory query")
.to_string();
assert!(err.contains("query params"), "got: {err}");
let err = StorageUrl::parse("file:///x?creds=y")
.expect_err("file query")
.to_string();
assert!(err.contains("query params"), "got: {err}");
assert!(StorageUrl::parse("/tmp/a?b").is_ok());
}
#[cfg(windows)]
#[test]
fn child_uri_round_trips_a_windows_drive_path() {
let store = StorageUrl::parse(r"C:\srv\pond").unwrap();
let child = crate::config::child_uri(store.lance_url(), "sessions.lance");
assert!(child.starts_with(r"C:\srv\pond\"), "got: {child}");
assert!(StorageUrl::parse(&child).is_ok(), "got: {child}");
}
#[test]
fn storage_url_refuses_windows_network_paths() {
for input in [
r"\\fileserver\share\pond",
"file://fileserver/share/pond",
r"\\?\UNC\fileserver\share\pond",
] {
let err = StorageUrl::parse(input)
.expect_err("a network path must be refused at the seam")
.to_string();
assert!(err.contains("Windows network path"), "got: {err}");
}
let err = StorageUrl::parse(r"\\?\C:\srv\pond")
.expect_err("extended-length paths are refused")
.to_string();
assert!(err.contains("extended-length"), "got: {err}");
assert!(!err.contains("network path"), "got: {err}");
assert!(StorageUrl::parse("file:///C:/srv/pond").is_ok());
assert!(StorageUrl::parse("file://localhost/srv/pond").is_ok());
#[cfg(windows)]
{
assert!(StorageUrl::parse(r"C:\srv\pond").is_ok());
assert!(StorageUrl::parse("C:/srv/pond").is_ok());
assert!(StorageUrl::parse("//fileserver/share/pond").is_err());
}
}
#[test]
fn storage_url_canonicalizes_ports_and_keeps_percent_encoding() {
let with_port = StorageUrl::parse("s3+https://host:443/bucket/p").unwrap();
let without = StorageUrl::parse("s3+https://host/bucket/p").unwrap();
assert_eq!(with_port.canonical(), without.canonical());
let odd = StorageUrl::parse("s3+https://host:8443/bucket").unwrap();
let resolved = odd.resolve(&BTreeMap::new()).unwrap();
assert_eq!(
resolved.options.get("endpoint").map(String::as_str),
Some("https://bucket.host:8443"),
);
let encoded = StorageUrl::parse("s3+https://host/bucket/pre%20fix").unwrap();
assert_eq!(encoded.lance_url().as_str(), "s3://bucket/pre%20fix");
}
#[test]
fn query_params_strip_and_apply_over_set_fields() {
let mut creds = BTreeMap::new();
creds.insert(
"default".to_owned(),
CredsSet {
region: Some("from-set".to_owned()),
virtual_hosted_style_request: Some(false),
..set(None)
},
);
let url = StorageUrl::parse(
"s3+https://host/bucket/p?region=from-query&virtual_hosted_style_request=true",
)
.unwrap();
assert_eq!(url.lance_url().as_str(), "s3://bucket/p");
assert!(url.canonical().query().is_none());
let resolved = url.resolve(&creds).unwrap();
assert_eq!(opts(&resolved, "region").as_deref(), Some("from-query"));
assert_eq!(
opts(&resolved, "virtual_hosted_style_request").as_deref(),
Some("true"),
);
assert_eq!(
opts(&resolved, "endpoint").as_deref(),
Some("https://bucket.host"),
);
}
#[test]
fn scope_matching_binds_by_longest_prefix_at_segment_boundaries() {
let mut creds = BTreeMap::new();
creds.insert("all".to_owned(), set(None));
creds.insert("bucket".to_owned(), set(Some("s3+https://host/pond/")));
creds.insert("deep".to_owned(), set(Some("s3+https://host/pond/sub")));
let bind = |input: &str| {
StorageUrl::parse(input)
.unwrap()
.resolve(&creds)
.unwrap()
.binding
};
assert_eq!(
bind("s3+https://host/pond/sub/x"),
CredsBinding::Set {
name: "deep".to_owned(),
via: BindVia::Scope
},
);
assert_eq!(
bind("s3+https://host/pond/other"),
CredsBinding::Set {
name: "bucket".to_owned(),
via: BindVia::Scope
},
);
assert_eq!(
bind("s3+https://host/pond-2"),
CredsBinding::Set {
name: "all".to_owned(),
via: BindVia::CatchAll
},
);
assert_eq!(
bind("s3://pond/sub"),
CredsBinding::Set {
name: "all".to_owned(),
via: BindVia::CatchAll
},
);
assert_eq!(
bind("s3+https://host:443/pond/x"),
CredsBinding::Set {
name: "bucket".to_owned(),
via: BindVia::Scope
},
);
assert_eq!(
bind("s3+https://host/pond/sub/x?creds=all"),
CredsBinding::Set {
name: "all".to_owned(),
via: BindVia::Pointer
},
);
let err = StorageUrl::parse("s3://b/p?creds=nope")
.unwrap()
.resolve(&creds)
.expect_err("missing set")
.to_string();
assert!(err.contains("creds=nope"), "got: {err}");
let empty = BTreeMap::new();
assert_eq!(
StorageUrl::parse("s3://b/p")
.unwrap()
.resolve(&empty)
.unwrap()
.binding,
CredsBinding::Ambient,
);
assert_eq!(
StorageUrl::parse("/srv/pond")
.unwrap()
.resolve(&creds)
.unwrap()
.binding,
CredsBinding::NotApplicable,
);
}
#[test]
fn unmatched_sets_are_reported_only_on_remote_invocations() {
let mut creds = BTreeMap::new();
creds.insert("used".to_owned(), set(Some("s3://bucket/")));
creds.insert("idle".to_owned(), set(Some("s3://other/")));
let remote = StorageUrl::parse("s3://bucket/p")
.unwrap()
.resolve(&creds)
.unwrap();
assert_eq!(unmatched_creds_sets(&[&remote], &creds), vec!["idle"]);
let local = StorageUrl::parse("/srv/pond")
.unwrap()
.resolve(&creds)
.unwrap();
assert!(unmatched_creds_sets(&[&local], &creds).is_empty());
}
#[test]
fn secrets_materialize_from_file_and_command() {
let dir = TempDir::new().unwrap();
let key_path = dir.path().join("key");
std::fs::write(&key_path, "from-file\n").unwrap();
let mut creds = BTreeMap::new();
creds.insert(
"default".to_owned(),
CredsSet {
access_key_id_file: Some(key_path),
secret_access_key_command: Some(
if cfg!(windows) {
"echo from-command&echo."
} else {
"printf 'from-command\\n\\n'"
}
.to_owned(),
),
..CredsSet::default()
},
);
let url = StorageUrl::parse("s3://bucket/p").unwrap();
let resolved = url.resolve(&creds).unwrap();
assert_eq!(
opts(&resolved, "access_key_id").as_deref(),
Some("from-file")
);
assert_eq!(
opts(&resolved, "secret_access_key").as_deref(),
Some(if cfg!(windows) {
"from-command\r\n"
} else {
"from-command\n"
}),
);
let mut failing = BTreeMap::new();
failing.insert(
"default".to_owned(),
CredsSet {
secret_access_key_command: Some("exit 3".to_owned()),
..CredsSet::default()
},
);
let err = url
.resolve(&failing)
.expect_err("command must fail")
.to_string();
assert!(err.contains("exit 3"), "got: {err}");
let marker = dir.path().join("runs");
let command = format!("echo run >> {} && echo secret", marker.display());
let mut counted = BTreeMap::new();
counted.insert(
"default".to_owned(),
CredsSet {
secret_access_key_command: Some(command),
..CredsSet::default()
},
);
url.resolve(&counted).unwrap();
url.resolve(&counted).unwrap();
let runs = std::fs::read_to_string(&marker).unwrap();
assert_eq!(runs.lines().count(), 1, "command must run exactly once");
}
#[test]
fn check_errors_classify_by_kind_and_binding() {
let auth_error = || object_store::Error::Unauthenticated {
path: "k".to_owned(),
source: "denied".into(),
};
let bound = CredsBinding::Set {
name: "work".to_owned(),
via: BindVia::Scope,
};
match classify_check_error(auth_error(), &bound, "put") {
CheckFailure::Auth { set, .. } => assert_eq!(set, "work"),
other => panic!("want Auth, got {other:?}"),
}
assert!(matches!(
classify_check_error(auth_error(), &CredsBinding::Ambient, "put"),
CheckFailure::NoCreds { .. },
));
let denied = object_store::Error::PermissionDenied {
path: "k".to_owned(),
source: "403".into(),
};
assert!(matches!(
classify_check_error(denied, &bound, "put"),
CheckFailure::Auth { .. },
));
let missing = object_store::Error::NotFound {
path: "k".to_owned(),
source: "404".into(),
};
assert!(matches!(
classify_check_error(missing, &bound, "get"),
CheckFailure::Io { .. },
));
let no_creds = || object_store::Error::Generic {
store: "S3",
source: "Failed to get AWS credentials: CredentialsNotLoaded".into(),
};
assert!(matches!(
classify_check_error(no_creds(), &bound, "put"),
CheckFailure::Auth { .. },
));
assert!(matches!(
classify_check_error(no_creds(), &CredsBinding::Ambient, "put"),
CheckFailure::NoCreds { .. },
));
}
#[test]
fn concise_cause_strips_upstream_noise_to_one_line() {
let inner = "Encountered internal error. Please file a bug report at \
https://github.com/lance-format/lance/issues. Failed to get AWS \
credentials: CredentialsNotLoaded, <WORKSPACE>/src/object_store/providers/aws.rs:401:21: \
Encountered internal error. Please file a bug report at \
https://github.com/lance-format/lance/issues. Failed to get AWS \
credentials: CredentialsNotLoaded";
let failure = CheckFailure::NoCreds {
source: anyhow!(inner.to_owned()).context("initial conditional put"),
};
let cause = failure.concise_cause().expect("auth-class carries a cause");
assert_eq!(cause, "Failed to get AWS credentials: CredentialsNotLoaded");
assert!(
!failure.to_string().contains("file a bug report"),
"lead must not trail the chain: {failure}"
);
let occ = CheckFailure::OccUnsupported {
detail: "put-if-none-match ignored".to_owned(),
};
assert!(occ.concise_cause().is_none());
let long = CheckFailure::Io {
source: anyhow!(format!("{} dns error: lookup failed", "x".repeat(500))),
};
let cause = long.concise_cause().expect("io carries a cause");
assert!(cause.contains(" ... "), "long causes truncate: {cause}");
assert!(
cause.ends_with("dns error: lookup failed"),
"the tail survives: {cause}"
);
}
#[tokio::test]
async fn storage_check_passes_on_memory_backend() {
let resolved = StorageUrl::parse("memory://check/probe")
.unwrap()
.resolve(&BTreeMap::new())
.unwrap();
storage_check(&resolved).await.expect("memory probe passes");
}
fn fragment(bytes: u64, rows: u64, deleted_rows: u64) -> FragmentStat {
FragmentStat {
bytes: Some(bytes),
rows,
deleted_rows,
}
}
fn stat(bytes: u64) -> FragmentStat {
fragment(bytes, bytes / 1_000, 0)
}
fn task_is_kept(stats: &[FragmentStat], target_rows_per_fragment: usize) -> bool {
task_veto_reason(
stats,
64,
0.1,
target_rows_per_fragment,
TARGET_FRAGMENT_BYTES,
)
.is_none()
}
#[test]
fn compaction_veto_blocks_absorb_keeps_peers() {
let absorb = [stat(100_000_000), stat(1_000_000), stat(2_000_000)];
assert!(!task_is_kept(&absorb, derived_target_rows(&absorb)));
let peers = [stat(100_000_000), stat(100_000_000)];
assert!(task_is_kept(&peers, derived_target_rows(&peers)));
let tiered = [stat(400_000), stat(60_000), stat(40_000)];
assert!(task_is_kept(&tiered, derived_target_rows(&tiered)));
}
#[test]
fn compaction_veto_passes_deletions_and_cap() {
let mut deleting = stat(665_000_000);
deleting.deleted_rows = deleting.rows / 5;
let deleting_task = [deleting, stat(1_000)];
assert!(task_is_kept(
&deleting_task,
derived_target_rows(&deleting_task),
));
let wide: Vec<FragmentStat> = std::iter::once(stat(100_000_000))
.chain(std::iter::repeat_with(|| stat(100_000)).take(63))
.collect();
assert!(task_is_kept(&wide, derived_target_rows(&wide)));
}
#[test]
fn compaction_veto_fails_closed_on_unknown_sizes() {
let mut unknown = stat(665_000_000);
unknown.bytes = None;
let task = [unknown, stat(665_000_000)];
assert_eq!(
task_veto_reason(
&task,
64,
0.1,
derived_target_rows(&task),
TARGET_FRAGMENT_BYTES
),
Some("missing_sizes"),
);
}
#[test]
fn compaction_veto_uses_physical_rows_after_deletions() {
let partially_deleted = || fragment(100_000_000, 100_000, 9_000);
let task: Vec<FragmentStat> = std::iter::repeat_with(partially_deleted).take(3).collect();
assert!(task_is_kept(&task, derived_target_rows(&task)));
}
#[test]
fn compaction_filter_keeps_above_budget_off_boundary_task() {
let table = [
fragment(100_000_000, 150_000, 0),
fragment(100_000_000, 150_000, 0),
fragment(100_000_000, 150_000, 0),
fragment(100_000_000, 50_000, 0),
];
let task = &table[..3];
let target = derived_target_rows(&table);
let total_bytes = task.iter().map(|stat| stat.bytes.unwrap()).sum::<u64>();
assert!(total_bytes > TARGET_FRAGMENT_BYTES);
assert!(target < derived_target_rows(task));
assert!(task_is_kept(task, target));
}
#[test]
fn compaction_filter_keeps_real_mixed_width_tasks_below_budget() {
let four_fragment_task = [
fragment(26_612_870, 9_281, 0),
fragment(14_242_314, 4_400, 0),
fragment(54_111_122, 20_988, 0),
fragment(517_923, 166, 0),
];
assert!(task_is_kept(&four_fragment_task, 58_468));
let nine_fragment_task = [
fragment(13_547_709, 3_946, 0),
fragment(344_320, 155, 0),
fragment(134_209, 34, 0),
fragment(54_624_719, 15_759, 0),
fragment(1_364_162, 292, 0),
fragment(110_225_801, 32_826, 0),
fragment(8_151_118, 1_840, 0),
fragment(728_590, 79, 0),
fragment(685_184, 128, 0),
];
assert!(task_is_kept(&nine_fragment_task, 58_468));
}
#[test]
fn compaction_veto_rejects_byte_capped_second_cycle() {
let wide_peer = || fragment(200_000_000, 2_000, 0);
let first_cycle: Vec<FragmentStat> = std::iter::repeat_with(wide_peer).take(5).collect();
let expected_outputs_by_bytes = 1_000_000_000u64.div_ceil(TARGET_FRAGMENT_BYTES) as usize;
assert_eq!(expected_outputs_by_bytes, 4);
assert!(expected_outputs_by_bytes < first_cycle.len());
let cap_sized_task: Vec<FragmentStat> =
std::iter::repeat_with(wide_peer).take(64).collect();
let second_cycle = [
fragment(TARGET_FRAGMENT_BYTES, 2_684, 0),
fragment(TARGET_FRAGMENT_BYTES, 2_684, 0),
fragment(TARGET_FRAGMENT_BYTES, 2_684, 0),
fragment(194_693_632, 1_948, 0),
];
assert_eq!(
task_veto_reason(&first_cycle, 64, 0.1, 66_000, TARGET_FRAGMENT_BYTES),
Some("row_target_unattainable"),
);
assert!(task_is_kept(&cap_sized_task, 66_000));
assert_eq!(
task_veto_reason(&second_cycle, 64, 0.1, 66_000, TARGET_FRAGMENT_BYTES),
Some("cannot_shrink"),
);
}
#[test]
fn parse_manifest_version_handles_all_naming_schemes() {
assert_eq!(parse_manifest_version("5.manifest"), Some(5));
assert_eq!(parse_manifest_version("0.manifest"), Some(0));
assert_eq!(
parse_manifest_version("18446744073709551615.manifest"),
Some(0)
);
assert_eq!(
parse_manifest_version("18446744073709551610.manifest"),
Some(5)
);
assert_eq!(parse_manifest_version("d123.manifest"), None);
assert_eq!(parse_manifest_version("5.manifest.corrupt"), None);
assert_eq!(
parse_manifest_version(".tmp_7.manifest_9c100374-3298-4537-afc6-f5ee7913666d"),
None
);
assert_eq!(parse_manifest_version("data.lance"), None);
assert_eq!(parse_manifest_version("notanumber.manifest"), None);
}
#[tokio::test]
async fn scan_verify_rejects_zeroed_column_add_data_file() {
let temp = tempfile::tempdir().unwrap();
let uri_owned = temp.path().join("t.lance");
let uri = uri_owned.to_str().unwrap();
let schema = Arc::new(lance::deps::arrow_schema::Schema::new(vec![
lance::deps::arrow_schema::Field::new(
"id",
lance::deps::arrow_schema::DataType::Utf8,
false,
),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(StringArray::from(vec!["a", "b", "c"]))],
)
.unwrap();
let reader = RecordBatchIterator::new([Ok(batch)], schema);
let mut dataset = Dataset::write(reader, uri, None).await.unwrap();
let data_files = || -> std::collections::BTreeSet<PathBuf> {
std::fs::read_dir(uri_owned.join("data"))
.unwrap()
.map(|entry| entry.unwrap().path())
.collect()
};
let before = data_files();
dataset
.add_columns(
lance::dataset::NewColumnTransform::SqlExpressions(vec![(
"extra".to_string(),
"id".to_string(),
)]),
None,
None,
)
.await
.unwrap();
let column_add_file = data_files()
.difference(&before)
.next()
.cloned()
.expect("add_columns writes a new per-fragment data file");
let version = dataset.version().version;
drop(dataset);
let fresh = || Arc::new(Session::new(0, 0, Arc::new(ObjectStoreRegistry::default())));
scan_verify_version(uri, version, &fresh(), &HashMap::new(), &None)
.await
.expect("intact version must pass scan-verify");
std::fs::write(&column_add_file, b"").unwrap();
let verdict = scan_verify_version(uri, version, &fresh(), &HashMap::new(), &None).await;
assert!(
verdict.is_err(),
"zeroed column-add data file must fail scan-verify",
);
}
#[test]
fn cleanup_due_gates_on_version_interval() {
assert!(cleanup_due(0, 1));
assert!(cleanup_due(7, 1));
assert!(cleanup_due(5, 0));
assert!(cleanup_due(0, 16));
assert!(cleanup_due(16, 16));
assert!(cleanup_due(48, 16));
assert!(!cleanup_due(15, 16));
assert!(!cleanup_due(17, 16));
assert!(!cleanup_due(31, 16));
}
#[test]
fn derived_target_rows_tracks_row_size_and_byte_cap() {
let parts_like = [FragmentStat {
bytes: Some(665_000_000),
rows: 511_000,
deleted_rows: 0,
}];
let target = derived_target_rows(&parts_like);
assert!((80_000..150_000).contains(&target), "{target}");
let unknown = [FragmentStat {
bytes: None,
rows: 511_000,
deleted_rows: 0,
}];
assert_eq!(
derived_target_rows(&unknown),
MAX_TARGET_ROWS_PER_FRAGMENT as usize
);
let tiny = [FragmentStat {
bytes: Some(1_000_000),
rows: 100_000,
deleted_rows: 0,
}];
assert_eq!(
derived_target_rows(&tiny),
MAX_TARGET_ROWS_PER_FRAGMENT as usize
);
let incident_parts = [FragmentStat {
bytes: Some(1_027_449_798),
rows: 105_087,
deleted_rows: 0,
}];
let incident_target = derived_target_rows(&incident_parts);
assert_eq!(incident_target, 13_727);
assert!(
u128::from(incident_parts[0].bytes.unwrap()) * incident_target as u128 * 2
<= u128::from(incident_parts[0].rows) * u128::from(TARGET_FRAGMENT_BYTES)
);
let huge = [FragmentStat {
bytes: Some(1_000_000_000),
rows: 1,
deleted_rows: 0,
}];
assert_eq!(derived_target_rows(&huge), 1);
}
#[test]
fn namespace_error_code_walks_wrapped_chain() {
let direct = lance::Error::namespace_source(Box::new(NamespaceError::TableNotFound {
message: "missing".into(),
}));
assert!(is_namespace_error_code(&direct, ErrorCode::TableNotFound));
let wrapped = lance::Error::namespace_source(Box::new(direct));
assert!(is_namespace_error_code(&wrapped, ErrorCode::TableNotFound));
let other_code =
lance::Error::namespace_source(Box::new(NamespaceError::NamespaceNotFound {
message: "nope".into(),
}));
assert!(!is_namespace_error_code(
&other_code,
ErrorCode::TableNotFound
));
let not_namespace = lance::Error::internal("unrelated");
assert!(!is_namespace_error_code(
¬_namespace,
ErrorCode::TableNotFound
));
}
#[tokio::test]
async fn store_opens_via_namespace_and_scan_works() -> Result<()> {
let temp = TempDir::new()?;
let url = Url::from_directory_path(temp.path())
.map_err(|()| anyhow::anyhow!("temp path is not absolute"))?;
let handle = Handle::open(&url).await?;
let cases: [(Table, &[&str]); 3] = [
(Table::Sessions, &["id"]),
(Table::Messages, &["id"]),
(Table::Parts, &["id"]),
];
for (table, projection) in cases {
let scanner = handle
.scan(table, ScanOpts::project_only(projection))
.await?;
let batch = scanner.try_into_batch().await?;
assert_eq!(batch.num_rows(), 0, "fresh table should be empty");
}
Ok(())
}
}