1use crate::{
5 RetryPolicy,
6 config::{self, CredsSet},
7 handlers::NamespaceIdent,
8 sessions::{self},
9};
10use anyhow::{Context, Result, anyhow, bail};
11use lance::Dataset;
12use lance::dataset::builder::DatasetBuilder;
13use lance::dataset::index::DatasetIndexRemapperOptions;
14use lance::dataset::optimize::{
15 CompactionMode, CompactionOptions, commit_compaction, plan_compaction,
16};
17pub use lance::dataset::write::merge_insert::MergeStats;
18use lance::dataset::write::merge_insert::SourceDedupeBehavior;
19use lance::dataset::{InsertBuilder, MergeInsertBuilder, WhenMatched, WhenNotMatched, WriteMode};
20pub use lance::dataset::{WriteParams, WriteStats};
21use lance::deps::arrow_array::{Array, RecordBatch, RecordBatchIterator, StringArray};
22use lance::deps::datafusion::physical_plan::SendableRecordBatchStream;
23use lance::index::DatasetIndexExt;
24use lance::index::DatasetIndexInternalExt;
25use lance::index::vector::VectorIndexParams;
26use lance::session::Session;
27use lance_index::IndexType;
28use lance_index::optimize::OptimizeOptions;
29use lance_index::scalar::{BuiltinIndexType, InvertedIndexParams, ScalarIndexParams};
30use lance_index::vector::ivf::IvfBuildParams;
31use lance_index::vector::sq::builder::SQBuildParams;
32use lance_io::object_store::{
33 ChainedWrappingObjectStore, ObjectStore, ObjectStoreParams, ObjectStoreRegistry,
34 StorageOptionsAccessor, WrappingObjectStore, uri_to_url,
35};
36use lance_linalg::distance::MetricType;
37use lance_namespace::LanceNamespace;
38use lance_namespace::error::{ErrorCode, NamespaceError};
39use lance_namespace::models::DescribeTableRequest;
40use lance_namespace_impls::ConnectBuilder;
41use std::{
42 collections::{BTreeMap, HashMap},
43 path::PathBuf,
44 sync::Arc,
45 time::{Duration, Instant},
46};
47use tokio::sync::{Mutex, OnceCell};
48use tokio_stream::StreamExt;
49use url::Url;
50pub const VECTOR_INDEX_ACTIVATION_ROWS: usize = 100_000;
55
56pub const DELTA_MERGE_THRESHOLD: usize = 4;
63
64#[derive(Debug, Clone, PartialEq)]
74pub struct StorageUrl {
75 canonical: Url,
79 lance: Url,
81 scheme_options: Vec<(&'static str, String)>,
83 query_options: Vec<(&'static str, String)>,
85 creds_pointer: Option<String>,
87 endpoint: Option<S3Endpoint>,
92}
93
94#[derive(Debug, Clone, PartialEq)]
95struct S3Endpoint {
96 scheme: &'static str,
97 authority: String,
99 bucket: String,
100}
101
102const RECOGNIZED_QUERY_PARAMS: [&str; 3] = ["creds", "region", "virtual_hosted_style_request"];
106
107impl StorageUrl {
108 pub fn parse(input: &str) -> Result<Self> {
112 let trimmed = input.trim();
113 if trimmed.is_empty() {
114 bail!("storage path is empty");
115 }
116 if !trimmed.contains("://") || trimmed.starts_with("file://") {
119 let url =
120 uri_to_url(trimmed).with_context(|| format!("invalid storage path {trimmed:?}"))?;
121 if url.query().is_some() {
126 bail!("storage URL {trimmed:?} carries query params; local URLs take none");
127 }
128 return Ok(Self::plain(url));
129 }
130 let url =
131 Url::parse(trimmed).with_context(|| format!("invalid storage URL {trimmed:?}"))?;
132 if !url.username().is_empty() || url.password().is_some() {
134 bail!(
135 "storage URL {trimmed:?} embeds credentials; put them in [creds.*] (or POND_CREDS_*) instead"
136 );
137 }
138 match url.scheme() {
139 "memory" | "shared-memory" => {
140 if url.query().is_some() {
141 bail!(
142 "storage URL {trimmed:?} carries query params; {}:// URLs take none",
143 url.scheme(),
144 );
145 }
146 Ok(Self::plain(url))
147 }
148 "s3" | "gs" => {
149 let (canonical, query_options, creds_pointer) = strip_query(url)?;
150 let mut lance = canonical.clone();
151 lance.set_query(None);
152 Ok(Self {
153 canonical,
154 lance,
155 scheme_options: Vec::new(),
156 query_options,
157 creds_pointer,
158 endpoint: None,
159 })
160 }
161 "s3+https" | "s3+http" => {
162 let (mut canonical, query_options, creds_pointer) = strip_query(url)?;
163 let tls = canonical.scheme() == "s3+https";
164 if canonical.port() == Some(if tls { 443 } else { 80 }) {
168 let _ = canonical.set_port(None);
169 }
170 let host = canonical
171 .host_str()
172 .ok_or_else(|| anyhow!("storage URL {trimmed:?} has no endpoint host"))?;
173 let endpoint_authority = match canonical.port() {
174 Some(port) => format!("{host}:{port}"),
175 None => host.to_owned(),
176 };
177 let mut segments = canonical.path().trim_start_matches('/').splitn(2, '/');
178 let bucket = segments.next().unwrap_or_default().to_owned();
179 let prefix = segments.next().unwrap_or_default().to_owned();
180 if bucket.is_empty() {
181 bail!(
182 "storage URL {trimmed:?} is missing the bucket: the form is {}://host/bucket/prefix",
183 canonical.scheme(),
184 );
185 }
186 let lance = Url::parse(&format!("s3://{bucket}/{prefix}")).with_context(|| {
187 format!("storage URL {trimmed:?}: bucket/prefix do not form a valid s3:// URL")
188 })?;
189 let scheme = if tls { "https" } else { "http" };
190 let virtual_hosted = host.parse::<std::net::IpAddr>().is_err()
199 && !matches!(canonical.host(), Some(url::Host::Ipv6(_)));
200 let scheme_options = vec![
201 ("allow_http", (!tls).to_string()),
202 ("virtual_hosted_style_request", virtual_hosted.to_string()),
203 ("region", "us-east-1".to_owned()),
210 ];
211 Ok(Self {
212 canonical,
213 lance,
214 scheme_options,
215 query_options,
216 creds_pointer,
217 endpoint: Some(S3Endpoint {
218 scheme,
219 authority: endpoint_authority,
220 bucket,
221 }),
222 })
223 }
224 "az" => {
225 let (canonical, query_options, creds_pointer) = strip_query(url)?;
226 let account = canonical
227 .host_str()
228 .ok_or_else(|| anyhow!("storage URL {trimmed:?} has no account: the form is az://account/container/prefix"))?
229 .to_owned();
230 let mut segments = canonical.path().trim_start_matches('/').splitn(2, '/');
231 let container = segments.next().unwrap_or_default();
232 if container.is_empty() {
233 bail!(
234 "storage URL {trimmed:?} is missing the container: the form is az://account/container/prefix"
235 );
236 }
237 let prefix = segments.next().unwrap_or_default();
238 let lance = Url::parse(&format!("az://{container}/{prefix}"))
239 .with_context(|| format!("storage URL {trimmed:?}: container/prefix do not form a valid az:// URL"))?;
240 Ok(Self {
241 canonical,
242 lance,
243 scheme_options: vec![("account_name", account)],
244 query_options,
245 creds_pointer,
246 endpoint: None,
247 })
248 }
249 other => bail!(
250 "storage URL scheme {other:?} not recognized; use a local path, s3://, s3+https://, s3+http://, gs://, or az://"
251 ),
252 }
253 }
254
255 fn plain(url: Url) -> Self {
257 Self {
258 canonical: url.clone(),
259 lance: url,
260 scheme_options: Vec::new(),
261 query_options: Vec::new(),
262 creds_pointer: None,
263 endpoint: None,
264 }
265 }
266
267 pub fn lance_url(&self) -> &Url {
269 &self.lance
270 }
271
272 pub fn canonical(&self) -> &Url {
275 &self.canonical
276 }
277
278 pub fn is_local(&self) -> bool {
279 config::is_local(&self.canonical)
280 }
281
282 pub fn display(&self) -> String {
284 config::display(&self.canonical)
285 }
286
287 fn takes_credentials(&self) -> bool {
290 !matches!(
291 self.canonical.scheme(),
292 "file" | "file+uring" | "memory" | "shared-memory"
293 )
294 }
295
296 pub fn resolve(&self, creds: &BTreeMap<String, CredsSet>) -> Result<ResolvedStorage> {
303 if !self.takes_credentials() {
304 return Ok(ResolvedStorage {
305 storage: self.clone(),
306 options: HashMap::new(),
307 binding: CredsBinding::NotApplicable,
308 });
309 }
310 let matched: Option<(&String, &CredsSet, BindVia)> = match &self.creds_pointer {
311 Some(name) => {
312 let set = creds.get(name).ok_or_else(|| {
313 anyhow!(
314 "URL names ?creds={name} but no [creds.{name}] set is configured; define it or drop the pointer"
315 )
316 })?;
317 Some((name, set, BindVia::Pointer))
318 }
319 None => {
320 let mut best: Option<(&String, &CredsSet, String)> = None;
321 for (name, set) in creds {
322 let Some(scope) = &set.scope else { continue };
323 let scope_url = parse_scope(scope).with_context(|| {
324 format!("[creds.{name}] scope {scope:?} is not a valid URL prefix")
325 })?;
326 if scope_matches(&scope_url, &self.canonical)
327 && best
328 .as_ref()
329 .is_none_or(|(_, _, len)| scope_url.as_str().len() > len.len())
330 {
331 best = Some((name, set, scope_url.as_str().to_owned()));
332 }
333 }
334 match best {
335 Some((name, set, _)) => Some((name, set, BindVia::Scope)),
336 None => creds
337 .iter()
338 .find(|(_, set)| set.scope.is_none())
339 .map(|(name, set)| (name, set, BindVia::CatchAll)),
340 }
341 }
342 };
343 let mut options: HashMap<String, String> = self
344 .scheme_options
345 .iter()
346 .map(|(key, value)| ((*key).to_owned(), value.clone()))
347 .collect();
348 let binding = match matched {
349 None => CredsBinding::Ambient,
350 Some((name, set, via)) => {
351 if let Some(region) = &set.region {
352 options.insert("region".to_owned(), region.clone());
353 }
354 if let Some(virtual_hosted) = set.virtual_hosted_style_request {
355 options.insert(
356 "virtual_hosted_style_request".to_owned(),
357 virtual_hosted.to_string(),
358 );
359 }
360 for (key, value) in &set.extra {
361 options.insert(key.clone(), value.clone());
362 }
363 if let Some(value) = materialize_secret(
364 name,
365 "access_key_id",
366 set.access_key_id.as_deref(),
367 set.access_key_id_file.as_deref(),
368 None,
369 )? {
370 options.insert("access_key_id".to_owned(), value);
371 }
372 if let Some(value) = materialize_secret(
373 name,
374 "secret_access_key",
375 set.secret_access_key.as_deref(),
376 set.secret_access_key_file.as_deref(),
377 set.secret_access_key_command.as_deref(),
378 )? {
379 options.insert("secret_access_key".to_owned(), value);
380 }
381 CredsBinding::Set {
382 name: name.clone(),
383 via,
384 }
385 }
386 };
387 for (key, value) in &self.query_options {
388 options.insert((*key).to_owned(), value.clone());
389 }
390 if let Some(endpoint) = &self.endpoint
395 && !options.keys().any(|key| {
396 key.eq_ignore_ascii_case("endpoint") || key.eq_ignore_ascii_case("aws_endpoint")
397 })
398 {
399 let virtual_hosted = options
400 .get("virtual_hosted_style_request")
401 .is_some_and(|value| value == "true");
402 let url = if virtual_hosted {
403 format!(
404 "{}://{}.{}",
405 endpoint.scheme, endpoint.bucket, endpoint.authority
406 )
407 } else {
408 format!("{}://{}", endpoint.scheme, endpoint.authority)
409 };
410 options.insert("endpoint".to_owned(), url);
411 }
412 Ok(ResolvedStorage {
413 storage: self.clone(),
414 options,
415 binding,
416 })
417 }
418}
419
420type StrippedQuery = (Url, Vec<(&'static str, String)>, Option<String>);
422
423fn strip_query(url: Url) -> Result<StrippedQuery> {
425 let mut query_options = Vec::new();
426 let mut creds_pointer = None;
427 for (key, value) in url.query_pairs() {
428 match RECOGNIZED_QUERY_PARAMS
429 .iter()
430 .find(|known| **known == key.as_ref())
431 {
432 Some(&"creds") => creds_pointer = Some(value.into_owned()),
433 Some(known) => query_options.push((*known, value.into_owned())),
434 None => bail!(
435 "storage URL query param {key:?} not recognized (known: {})",
436 RECOGNIZED_QUERY_PARAMS.join(", "),
437 ),
438 }
439 }
440 let mut canonical = url;
441 canonical.set_query(None);
442 Ok((canonical, query_options, creds_pointer))
443}
444
445pub(crate) fn parse_scope(scope: &str) -> Result<Url> {
448 let mut url = Url::parse(scope.trim())?;
449 if !url.username().is_empty() || url.password().is_some() {
450 bail!("scope embeds credentials");
451 }
452 if url.query().is_some() {
453 bail!("scope carries query params; scopes are plain URL prefixes");
454 }
455 match (url.scheme(), url.port()) {
456 ("s3+https", Some(443)) | ("s3+http", Some(80)) => {
457 let _ = url.set_port(None);
458 }
459 _ => {}
460 }
461 Ok(url)
462}
463
464fn scope_matches(scope: &Url, address: &Url) -> bool {
469 if scope.scheme() != address.scheme()
470 || scope.host_str() != address.host_str()
471 || scope.port() != address.port()
472 {
473 return false;
474 }
475 let scope_path = scope.path().trim_end_matches('/');
476 let address_path = address.path().trim_end_matches('/');
477 address_path == scope_path
478 || address_path
479 .strip_prefix(scope_path)
480 .is_some_and(|rest| rest.starts_with('/'))
481}
482
483#[derive(Debug, Clone, Copy, PartialEq, Eq)]
486pub enum BindVia {
487 Pointer,
489 Scope,
491 CatchAll,
493}
494
495#[derive(Debug, Clone, PartialEq)]
496pub enum CredsBinding {
497 Set { name: String, via: BindVia },
499 Ambient,
504 NotApplicable,
506}
507
508impl CredsBinding {
509 pub fn describe(&self) -> String {
511 match self {
512 Self::Set { name, via } => {
513 let via = match via {
514 BindVia::Pointer => "?creds",
515 BindVia::Scope => "scope match",
516 BindVia::CatchAll => "catch-all",
517 };
518 format!("creds {name} ({via})")
519 }
520 Self::Ambient => "ambient chain".to_owned(),
521 Self::NotApplicable => "local (no credentials)".to_owned(),
522 }
523 }
524}
525
526#[derive(Debug, Clone)]
530pub struct ResolvedStorage {
531 storage: StorageUrl,
532 pub options: HashMap<String, String>,
533 pub binding: CredsBinding,
534}
535
536impl ResolvedStorage {
537 pub fn lance_url(&self) -> &Url {
538 self.storage.lance_url()
539 }
540
541 pub fn display(&self) -> String {
542 self.storage.display()
543 }
544}
545
546pub fn unmatched_creds_sets<'c>(
551 resolved: &[&ResolvedStorage],
552 creds: &'c BTreeMap<String, CredsSet>,
553) -> Vec<&'c str> {
554 if resolved
555 .iter()
556 .all(|entry| matches!(entry.binding, CredsBinding::NotApplicable))
557 {
558 return Vec::new();
559 }
560 creds
561 .keys()
562 .filter(|name| {
563 !resolved.iter().any(|entry| {
564 matches!(&entry.binding, CredsBinding::Set { name: bound, .. } if bound == *name)
565 })
566 })
567 .map(String::as_str)
568 .collect()
569}
570
571fn materialize_secret(
574 set: &str,
575 field: &str,
576 inline: Option<&str>,
577 file: Option<&std::path::Path>,
578 command: Option<&str>,
579) -> Result<Option<String>> {
580 if let Some(value) = inline {
581 return Ok(Some(value.to_owned()));
582 }
583 if let Some(path) = file {
584 let text = std::fs::read_to_string(path).with_context(|| {
585 format!(
586 "[creds.{set}] {field}_file: failed to read {}",
587 path.display()
588 )
589 })?;
590 return Ok(Some(strip_one_newline(text)));
591 }
592 if let Some(command) = command {
593 return Ok(Some(run_secret_command(set, field, command)?));
594 }
595 Ok(None)
596}
597
598fn run_secret_command(set: &str, field: &str, command: &str) -> Result<String> {
601 static CACHE: std::sync::OnceLock<std::sync::Mutex<HashMap<String, String>>> =
602 std::sync::OnceLock::new();
603 let cache = CACHE.get_or_init(Default::default);
604 if let Some(hit) = cache
605 .lock()
606 .unwrap_or_else(std::sync::PoisonError::into_inner)
607 .get(command)
608 {
609 return Ok(hit.clone());
610 }
611 let output = std::process::Command::new("sh")
612 .arg("-c")
613 .arg(command)
614 .output()
615 .with_context(|| format!("[creds.{set}] {field}_command failed to spawn: {command}"))?;
616 if !output.status.success() {
617 bail!(
618 "[creds.{set}] {field}_command exited {}: {command}\n{}",
619 output.status,
620 String::from_utf8_lossy(&output.stderr).trim_end(),
621 );
622 }
623 let value = strip_one_newline(
624 String::from_utf8(output.stdout)
625 .with_context(|| format!("[creds.{set}] {field}_command output is not UTF-8"))?,
626 );
627 cache
628 .lock()
629 .unwrap_or_else(std::sync::PoisonError::into_inner)
630 .insert(command.to_owned(), value.clone());
631 Ok(value)
632}
633
634fn strip_one_newline(mut text: String) -> String {
637 if text.ends_with('\n') {
638 text.pop();
639 if text.ends_with('\r') {
640 text.pop();
641 }
642 }
643 text
644}
645
646#[derive(Debug, thiserror::Error)]
654pub enum CheckFailure {
655 #[error(
656 "authentication failed and no creds set matched this URL; add one with `pond creds add` (or set POND_CREDS_*), or provide ambient AWS_* credentials"
657 )]
658 NoCreds { source: anyhow::Error },
659 #[error("authentication failed using creds set {set:?}; check its keys and scope")]
660 Auth { set: String, source: anyhow::Error },
661 #[error(
662 "backend does not enforce conditional writes (If-None-Match); concurrent pond writers would corrupt each other - {detail}"
663 )]
664 OccUnsupported { detail: String },
665 #[error("storage probe failed")]
666 Io { source: anyhow::Error },
667}
668
669impl CheckFailure {
670 pub fn concise_cause(&self) -> Option<String> {
676 let source = match self {
677 Self::NoCreds { source } | Self::Auth { source, .. } | Self::Io { source } => source,
678 Self::OccUnsupported { .. } => return None,
679 };
680 Some(condense_error_chain(source))
681 }
682}
683
684fn condense_error_chain(error: &anyhow::Error) -> String {
691 let mut text = error
692 .chain()
693 .last()
694 .map(ToString::to_string)
695 .unwrap_or_else(|| format!("{error:#}"));
696 if let Some(pos) = text.find(", <WORKSPACE>") {
697 text.truncate(pos);
698 }
699 text = text.replace(
700 "Encountered internal error. Please file a bug report at https://github.com/lance-format/lance/issues. ",
701 "",
702 );
703 let line = text.split_whitespace().collect::<Vec<_>>().join(" ");
704 const HEAD: usize = 120;
705 const TAIL: usize = 120;
706 let chars: Vec<char> = line.chars().collect();
707 if chars.len() > HEAD + TAIL + 5 {
708 let head: String = chars[..HEAD].iter().collect();
709 let tail: String = chars[chars.len() - TAIL..].iter().collect();
710 format!("{head} ... {tail}")
711 } else {
712 line
713 }
714}
715
716pub async fn storage_check(resolved: &ResolvedStorage) -> std::result::Result<(), CheckFailure> {
721 use object_store::{Error as OsError, ObjectStoreExt, PutMode, PutOptions, PutPayload};
722
723 let classify =
724 |error: OsError, step: &str| classify_check_error(error, &resolved.binding, step);
725
726 let probe_uri = format!(
727 "{}/_config-check/{}",
728 resolved.lance_url().as_str().trim_end_matches('/'),
729 uuid::Uuid::now_v7(),
730 );
731 let params = ObjectStoreParams {
732 storage_options_accessor: (!resolved.options.is_empty()).then(|| {
733 Arc::new(StorageOptionsAccessor::with_static_options(
734 resolved.options.clone(),
735 ))
736 }),
737 ..Default::default()
738 };
739 let registry = Arc::new(ObjectStoreRegistry::default());
740 let (store, path) = ObjectStore::from_uri_and_params(registry, &probe_uri, ¶ms)
741 .await
742 .map_err(|error| CheckFailure::Io {
743 source: anyhow!(error).context(format!("failed to open object store for {probe_uri}")),
744 })?;
745
746 let body: &[u8] = b"pond storage check";
747 let create = PutOptions::from(PutMode::Create);
748 store
749 .inner
750 .put_opts(&path, PutPayload::from_static(body), create.clone())
751 .await
752 .map_err(|error| classify(error, "initial conditional put"))?;
753 let outcome = async {
757 match store
762 .inner
763 .put_opts(&path, PutPayload::from_static(body), create)
764 .await
765 {
766 Err(OsError::AlreadyExists { .. }) => {}
767 Ok(_) => {
768 return Err(CheckFailure::OccUnsupported {
769 detail: "a second create over an existing key succeeded".to_owned(),
770 });
771 }
772 Err(OsError::NotImplemented { .. }) => {
773 return Err(CheckFailure::OccUnsupported {
774 detail: "the backend rejects conditional puts as unimplemented".to_owned(),
775 });
776 }
777 Err(error) => return Err(classify(error, "conditional-put probe")),
778 }
779 let read_back = store
780 .inner
781 .get(&path)
782 .await
783 .map_err(|error| classify(error, "read-back"))?
784 .bytes()
785 .await
786 .map_err(|error| classify(error, "read-back body"))?;
787 if read_back.as_ref() != body {
788 return Err(CheckFailure::Io {
789 source: anyhow!("read-back returned different bytes than written"),
790 });
791 }
792 Ok(())
793 }
794 .await;
795 let cleanup = store.inner.delete(&path).await;
796 outcome?;
797 cleanup.map_err(|error| classify(error, "cleanup delete"))?;
798 Ok(())
799}
800
801fn classify_check_error(
805 error: object_store::Error,
806 binding: &CredsBinding,
807 step: &str,
808) -> CheckFailure {
809 use object_store::Error as OsError;
810 let auth_class = matches!(
815 error,
816 OsError::Unauthenticated { .. } | OsError::PermissionDenied { .. }
817 ) || {
818 let rendered = error.to_string();
819 rendered.contains("CredentialsNotLoaded")
820 || rendered.contains("no providers in chain provided credentials")
821 };
822 match (auth_class, binding) {
823 (true, CredsBinding::Set { name, .. }) => CheckFailure::Auth {
824 set: name.clone(),
825 source: anyhow!(error).context(step.to_owned()),
826 },
827 (true, _) => CheckFailure::NoCreds {
828 source: anyhow!(error).context(step.to_owned()),
829 },
830 (false, _) => CheckFailure::Io {
831 source: anyhow!(error).context(step.to_owned()),
832 },
833 }
834}
835
836pub const DEFAULT_COMPACTION_FRAGMENT_CAP: usize = 64;
840
841pub const TARGET_FRAGMENT_BYTES: u64 = 256 * 1024 * 1024;
845
846const MIN_TARGET_ROWS_PER_FRAGMENT: u64 = 50_000;
847const MAX_TARGET_ROWS_PER_FRAGMENT: u64 = 1024 * 1024;
849
850pub const COMPACTION_ABSORB_FACTOR: u64 = 4;
853
854pub fn default_cleanup_older_than() -> chrono::Duration {
860 chrono::Duration::hours(1)
864}
865
866pub const DEFAULT_SYNC_CLEANUP_INTERVAL: u64 = 16;
873
874pub const DEFAULT_SYNC_SCALAR_FOLD_ROWS: usize = 50_000;
883
884pub const DEFAULT_SYNC_INDEX_FOLD_ROWS: usize = 5_000;
894
895#[derive(Debug, Clone, Copy)]
900pub struct MaintenancePolicy {
901 pub compaction_fragment_cap: usize,
903 pub cleanup_older_than: chrono::Duration,
905 pub cleanup_interval: u64,
910 pub scalar_fold_row_threshold: usize,
915 pub index_fold_row_threshold: usize,
921}
922
923impl MaintenancePolicy {
924 pub fn always_compact() -> Self {
926 Self {
927 compaction_fragment_cap: 0,
928 cleanup_older_than: default_cleanup_older_than(),
929 cleanup_interval: 1,
930 scalar_fold_row_threshold: 0,
931 index_fold_row_threshold: 0,
932 }
933 }
934
935 #[must_use]
938 pub fn with_cleanup_interval(mut self, interval: u64) -> Self {
939 self.cleanup_interval = interval.max(1);
940 self
941 }
942
943 #[must_use]
947 pub fn with_scalar_fold_row_threshold(mut self, threshold: usize) -> Self {
948 self.scalar_fold_row_threshold = threshold;
949 self
950 }
951
952 #[must_use]
956 pub fn with_index_fold_row_threshold(mut self, threshold: usize) -> Self {
957 self.index_fold_row_threshold = threshold;
958 self
959 }
960
961 fn fold_thresholds(&self) -> FoldThresholds {
964 FoldThresholds {
965 scalar: self.scalar_fold_row_threshold,
966 index: self.index_fold_row_threshold,
967 }
968 }
969}
970
971#[derive(Debug, Clone, Copy)]
974struct FoldThresholds {
975 scalar: usize,
976 index: usize,
977}
978
979struct FragmentStat {
980 bytes: Option<u64>,
982 rows: u64,
983 deleted_rows: u64,
984}
985
986fn fragment_bytes(fragment: &lance::table::format::Fragment) -> Option<u64> {
989 fragment.files.iter().try_fold(0u64, |total, file| {
990 Some(total + file.file_size_bytes.get()?.get())
991 })
992}
993
994fn fragment_stat(fragment: &lance::table::format::Fragment) -> FragmentStat {
995 FragmentStat {
996 bytes: fragment_bytes(fragment),
997 rows: fragment.physical_rows.unwrap_or(0) as u64,
998 deleted_rows: fragment
999 .deletion_file
1000 .as_ref()
1001 .and_then(|deletions| deletions.num_deleted_rows)
1002 .unwrap_or(0) as u64,
1003 }
1004}
1005
1006fn derived_target_rows(stats: &[FragmentStat]) -> usize {
1017 let (mut bytes, mut rows) = (0u64, 0u64);
1018 for stat in stats {
1019 if let Some(fragment_bytes) = stat.bytes
1020 && stat.rows > 0
1021 {
1022 bytes += fragment_bytes;
1023 rows += stat.rows;
1024 }
1025 }
1026 if bytes == 0 || rows == 0 {
1027 return MAX_TARGET_ROWS_PER_FRAGMENT as usize;
1028 }
1029 let avg_row_bytes = (bytes / rows).max(1);
1030 (TARGET_FRAGMENT_BYTES / 2 / avg_row_bytes)
1031 .clamp(MIN_TARGET_ROWS_PER_FRAGMENT, MAX_TARGET_ROWS_PER_FRAGMENT) as usize
1032}
1033
1034fn keep_task(stats: &[FragmentStat], cap: usize, deletion_threshold: f32) -> bool {
1039 if stats.iter().any(|stat| {
1040 stat.rows > 0 && (stat.deleted_rows as f32 / stat.rows as f32) > deletion_threshold
1041 }) {
1042 return true;
1043 }
1044 if stats.len() >= cap {
1045 return true;
1046 }
1047 let weights: Vec<u64> = if stats.iter().all(|stat| stat.bytes.is_some()) {
1048 stats.iter().filter_map(|stat| stat.bytes).collect()
1049 } else {
1050 stats.iter().map(|stat| stat.rows).collect()
1051 };
1052 let total: u64 = weights.iter().sum();
1053 let largest = weights.iter().copied().max().unwrap_or(0);
1054 (total - largest) * COMPACTION_ABSORB_FACTOR >= largest
1055}
1056
1057#[derive(Debug, Clone)]
1060pub struct IndexIntent {
1061 pub name: &'static str,
1064 pub column: &'static str,
1066 pub trigger: IndexTrigger,
1068 pub params: IndexParamsKind,
1071}
1072
1073#[derive(Debug, Clone)]
1075pub enum IndexTrigger {
1076 OnAnyRows,
1079 OnNonNullCount {
1082 column: &'static str,
1083 threshold: usize,
1084 },
1085}
1086
1087#[derive(Debug, Clone)]
1090pub enum IndexParamsKind {
1091 Scalar(BuiltinIndexType),
1094 InvertedFtsWord,
1100 IvfSqCosine { num_bits: u16, max_iters: usize },
1109}
1110
1111impl IndexTrigger {
1112 async fn should_create(&self, dataset: &Dataset) -> Result<bool> {
1113 match self {
1114 Self::OnAnyRows => Ok(dataset.count_rows(None).await? > 0),
1115 Self::OnNonNullCount { column, threshold } => {
1116 let count = dataset
1117 .count_rows(Some(format!("{column} IS NOT NULL")))
1118 .await?;
1119 Ok(count >= *threshold)
1120 }
1121 }
1122 }
1123}
1124
1125impl IndexParamsKind {
1126 fn index_type(&self) -> IndexType {
1127 match self {
1128 Self::Scalar(BuiltinIndexType::Bitmap) => IndexType::Bitmap,
1129 Self::Scalar(BuiltinIndexType::ZoneMap) => IndexType::ZoneMap,
1130 Self::Scalar(_) => IndexType::BTree,
1131 Self::InvertedFtsWord => IndexType::Inverted,
1132 Self::IvfSqCosine { .. } => IndexType::Vector,
1133 }
1134 }
1135
1136 async fn build(&self, dataset: &Dataset) -> Result<Box<dyn lance::index::IndexParams>> {
1137 match self {
1138 Self::Scalar(kind) => Ok(Box::new(ScalarIndexParams::for_builtin(kind.clone()))),
1139 Self::InvertedFtsWord => Ok(Box::new(
1140 InvertedIndexParams::default()
1141 .base_tokenizer("simple".to_owned())
1142 .stem(true)
1143 .remove_stop_words(false),
1144 )),
1145 Self::IvfSqCosine {
1146 num_bits,
1147 max_iters,
1148 } => {
1149 let count = dataset
1150 .count_rows(Some("vector IS NOT NULL".to_owned()))
1151 .await?;
1152 let partitions = count.checked_div(4096).unwrap_or(0).max(1);
1153 let mut ivf = IvfBuildParams::new(partitions);
1154 ivf.max_iters = *max_iters;
1155 let sq = SQBuildParams {
1156 num_bits: *num_bits,
1157 ..Default::default()
1158 };
1159 Ok(Box::new(VectorIndexParams::with_ivf_sq_params(
1160 MetricType::Cosine,
1161 ivf,
1162 sq,
1163 )))
1164 }
1165 }
1166 }
1167}
1168
1169#[derive(Debug, Clone, PartialEq, Eq)]
1170pub struct IndexStatus {
1171 pub table: Table,
1172 pub intent_name: String,
1173 pub fragments_covered: usize,
1174 pub unindexed_fragments: usize,
1175 pub unindexed_rows: usize,
1176 pub exists: bool,
1177}
1178
1179#[derive(Debug, Clone, Copy)]
1184pub struct ConflictExhausted {
1185 pub attempts: u8,
1186}
1187
1188impl std::fmt::Display for ConflictExhausted {
1189 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1190 write!(
1191 formatter,
1192 "commit conflict exhausted after {} attempt(s)",
1193 self.attempts
1194 )
1195 }
1196}
1197
1198impl std::error::Error for ConflictExhausted {}
1199
1200#[derive(Debug)]
1205pub enum PhaseOutcome {
1206 Ok,
1208 Noop,
1210 SkippedConflict,
1213 Failed(anyhow::Error),
1215 NotAttempted,
1218}
1219
1220impl PhaseOutcome {
1221 pub fn is_failed(&self) -> bool {
1222 matches!(self, Self::Failed(_))
1223 }
1224}
1225
1226#[derive(Debug)]
1228pub struct TableOptimizeOutcome {
1229 pub table: Table,
1230 pub indices: PhaseOutcome,
1231 pub compaction: PhaseOutcome,
1232}
1233
1234#[derive(Debug, Clone)]
1237pub enum OptimizeEvent {
1238 PhaseStart {
1239 table: Table,
1240 phase: OptimizePhase,
1241 detail: Option<String>,
1242 },
1243 PhaseDone {
1244 table: Table,
1245 phase: OptimizePhase,
1246 elapsed_ms: u64,
1247 },
1248 IndexStage {
1253 table: Table,
1254 index: String,
1255 stage: String,
1256 completed: u64,
1257 total: Option<u64>,
1258 unit: String,
1259 },
1260}
1261
1262#[derive(Debug, Clone, Copy)]
1263pub enum OptimizePhase {
1264 Compact,
1265 Cleanup,
1266 IndexCreate,
1267 IndexRebuild,
1268 IndexAppend,
1269}
1270
1271impl OptimizePhase {
1272 pub fn label(self) -> &'static str {
1273 match self {
1274 Self::Compact => "compact",
1275 Self::Cleanup => "cleanup",
1276 Self::IndexCreate => "index-create",
1277 Self::IndexRebuild => "index-rebuild",
1278 Self::IndexAppend => "index-append",
1279 }
1280 }
1281}
1282
1283pub type OptimizeProgressFn = Arc<dyn Fn(OptimizeEvent) + Send + Sync>;
1287
1288fn emit(progress: Option<&OptimizeProgressFn>, event: OptimizeEvent) {
1289 if let Some(callback) = progress {
1290 callback(event);
1291 }
1292}
1293
1294struct PondIndexProgress {
1303 callback: OptimizeProgressFn,
1304 table: Table,
1305 index: String,
1306 state: std::sync::Mutex<PondIndexStageState>,
1307}
1308
1309impl std::fmt::Debug for PondIndexProgress {
1312 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1313 f.debug_struct("PondIndexProgress")
1314 .field("table", &self.table)
1315 .field("index", &self.index)
1316 .finish_non_exhaustive()
1317 }
1318}
1319
1320#[derive(Debug, Default)]
1321struct PondIndexStageState {
1322 total: Option<u64>,
1323 unit: String,
1324 last_emit: Option<Instant>,
1325}
1326
1327impl PondIndexProgress {
1328 fn new(callback: OptimizeProgressFn, table: Table, index: String) -> Arc<Self> {
1329 Arc::new(Self {
1330 callback,
1331 table,
1332 index,
1333 state: std::sync::Mutex::new(PondIndexStageState::default()),
1334 })
1335 }
1336}
1337
1338#[async_trait::async_trait]
1339impl lance_index::progress::IndexBuildProgress for PondIndexProgress {
1340 async fn stage_start(&self, stage: &str, total: Option<u64>, unit: &str) -> lance::Result<()> {
1341 if let Ok(mut state) = self.state.lock() {
1342 state.total = total;
1343 state.unit = unit.to_owned();
1344 state.last_emit = Some(Instant::now());
1345 }
1346 (self.callback)(OptimizeEvent::IndexStage {
1347 table: self.table,
1348 index: self.index.clone(),
1349 stage: stage.to_owned(),
1350 completed: 0,
1351 total,
1352 unit: unit.to_owned(),
1353 });
1354 Ok(())
1355 }
1356
1357 async fn stage_progress(&self, stage: &str, completed: u64) -> lance::Result<()> {
1358 let (total, unit) = {
1359 let Ok(mut state) = self.state.lock() else {
1360 return Ok(());
1361 };
1362 let now = Instant::now();
1363 if let Some(prev) = state.last_emit
1364 && now.duration_since(prev) < Duration::from_millis(100)
1365 {
1366 return Ok(());
1367 }
1368 state.last_emit = Some(now);
1369 (state.total, state.unit.clone())
1370 };
1371 (self.callback)(OptimizeEvent::IndexStage {
1372 table: self.table,
1373 index: self.index.clone(),
1374 stage: stage.to_owned(),
1375 completed,
1376 total,
1377 unit,
1378 });
1379 Ok(())
1380 }
1381
1382 async fn stage_complete(&self, stage: &str) -> lance::Result<()> {
1383 let (total, unit) = {
1384 let Ok(state) = self.state.lock() else {
1385 return Ok(());
1386 };
1387 (state.total, state.unit.clone())
1388 };
1389 (self.callback)(OptimizeEvent::IndexStage {
1390 table: self.table,
1391 index: self.index.clone(),
1392 stage: stage.to_owned(),
1393 completed: total.unwrap_or(0),
1394 total,
1395 unit,
1396 });
1397 Ok(())
1398 }
1399}
1400
1401fn lance_progress(
1402 progress: Option<&OptimizeProgressFn>,
1403 table: Table,
1404 index: &str,
1405) -> Arc<dyn lance_index::progress::IndexBuildProgress> {
1406 match progress {
1407 Some(callback) => PondIndexProgress::new(callback.clone(), table, index.to_owned()),
1408 None => Arc::new(lance_index::progress::NoopIndexBuildProgress),
1409 }
1410}
1411
1412pub fn is_commit_conflict(error: &anyhow::Error) -> bool {
1416 error.downcast_ref::<lance::Error>().is_some_and(|err| {
1417 matches!(
1418 err,
1419 lance::Error::CommitConflict { .. }
1420 | lance::Error::RetryableCommitConflict { .. }
1421 | lance::Error::TooMuchWriteContention { .. }
1422 )
1423 })
1424}
1425
1426fn is_conflict_exhausted(error: &anyhow::Error) -> bool {
1429 error.chain().any(|cause| cause.is::<ConflictExhausted>())
1430}
1431
1432pub fn is_index_error(error: &anyhow::Error) -> bool {
1436 error
1437 .downcast_ref::<lance::Error>()
1438 .is_some_and(|err| matches!(err, lance::Error::Index { .. }))
1439}
1440
1441#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1445pub struct TableSizes {
1446 pub sessions: u64,
1447 pub messages: u64,
1448 pub parts: u64,
1449 pub other: u64,
1450 pub sessions_data: DataLiveness,
1451 pub messages_data: DataLiveness,
1452 pub parts_data: DataLiveness,
1453}
1454
1455#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1458pub struct DataLiveness {
1459 pub on_disk: u64,
1460 pub live: Option<u64>,
1462}
1463
1464impl DataLiveness {
1465 pub fn dead(&self) -> Option<u64> {
1466 self.live.map(|live| self.on_disk.saturating_sub(live))
1467 }
1468}
1469
1470#[derive(Debug, Clone, PartialEq, Eq)]
1471pub enum ScalarValue {
1472 String(String),
1473 Int32(i32),
1474 Raw(String),
1475}
1476impl From<&str> for ScalarValue {
1477 fn from(value: &str) -> Self {
1478 Self::String(value.to_owned())
1479 }
1480}
1481impl From<String> for ScalarValue {
1482 fn from(value: String) -> Self {
1483 Self::String(value)
1484 }
1485}
1486impl From<i32> for ScalarValue {
1487 fn from(value: i32) -> Self {
1488 Self::Int32(value)
1489 }
1490}
1491#[derive(Debug, Clone, PartialEq, Eq)]
1492pub enum Predicate {
1493 Eq(&'static str, ScalarValue),
1494 Ne(&'static str, ScalarValue),
1495 IsNull(&'static str),
1496 IsNotNull(&'static str),
1497 In(&'static str, Vec<ScalarValue>),
1498 LikeContains(&'static str, String),
1499 Regex(&'static str, String),
1504 Gte(&'static str, ScalarValue),
1505 Lte(&'static str, ScalarValue),
1506 And(Vec<Predicate>),
1507 Or(Vec<Predicate>),
1508 Not(Box<Predicate>),
1509}
1510impl Predicate {
1511 pub fn to_lance(&self) -> String {
1512 match self {
1513 Self::Eq(column, value) => format!("{column} = {}", value.to_lance()),
1514 Self::Ne(column, value) => format!("{column} <> {}", value.to_lance()),
1515 Self::IsNull(column) => format!("{column} IS NULL"),
1516 Self::IsNotNull(column) => format!("{column} IS NOT NULL"),
1517 Self::In(column, values) => {
1518 let values = values
1519 .iter()
1520 .map(ScalarValue::to_lance)
1521 .collect::<Vec<_>>()
1522 .join(", ");
1523 format!("{column} IN ({values})")
1524 }
1525 Self::LikeContains(column, value) => {
1526 format!("{column} LIKE {} ESCAPE '\\'", like_contains(value))
1527 }
1528 Self::Regex(column, pattern) => {
1529 format!("regexp_like({column}, {})", quoted_string(pattern))
1530 }
1531 Self::Gte(column, value) => format!("{column} >= {}", value.to_lance()),
1532 Self::Lte(column, value) => format!("{column} <= {}", value.to_lance()),
1533 Self::And(predicates) => predicates
1534 .iter()
1535 .map(Self::to_lance)
1536 .filter(|predicate| !predicate.is_empty())
1537 .collect::<Vec<_>>()
1538 .join(" AND "),
1539 Self::Or(predicates) => {
1540 let body = predicates
1543 .iter()
1544 .map(Self::to_lance)
1545 .filter(|predicate| !predicate.is_empty())
1546 .collect::<Vec<_>>()
1547 .join(" OR ");
1548 if body.is_empty() {
1549 String::new()
1550 } else {
1551 format!("({body})")
1552 }
1553 }
1554 Self::Not(inner) => {
1555 let body = inner.to_lance();
1556 if body.is_empty() {
1557 String::new()
1558 } else {
1559 format!("NOT ({body})")
1560 }
1561 }
1562 }
1563 }
1564}
1565#[derive(Default)]
1568pub struct ScanOpts<'a> {
1569 pub predicate: Option<&'a Predicate>,
1570 pub projection: Option<&'a [&'a str]>,
1571}
1572
1573impl<'a> ScanOpts<'a> {
1574 pub fn project_only(projection: &'a [&'a str]) -> Self {
1575 Self {
1576 predicate: None,
1577 projection: Some(projection),
1578 }
1579 }
1580 pub fn with_predicate_and_projection(
1581 predicate: &'a Predicate,
1582 projection: &'a [&'a str],
1583 ) -> Self {
1584 Self {
1585 predicate: Some(predicate),
1586 projection: Some(projection),
1587 }
1588 }
1589}
1590
1591impl ScalarValue {
1592 fn to_lance(&self) -> String {
1593 match self {
1594 Self::String(value) => quoted_string(value),
1595 Self::Int32(value) => value.to_string(),
1596 Self::Raw(value) => value.clone(),
1597 }
1598 }
1599}
1600#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1604pub struct RuntimeCaps {
1605 pub index_cache_bytes: Option<usize>,
1606 pub metadata_cache_bytes: Option<usize>,
1607}
1608
1609impl RuntimeCaps {
1610 pub fn from_config(config: &crate::config::RuntimeConfig) -> Self {
1611 Self {
1612 index_cache_bytes: config.index_cache_bytes,
1613 metadata_cache_bytes: config.metadata_cache_bytes,
1614 }
1615 }
1616}
1617
1618const LOCAL_INDEX_CACHE_BYTES: usize = 256 * 1024 * 1024;
1622const LOCAL_METADATA_CACHE_BYTES: usize = 128 * 1024 * 1024;
1623const REMOTE_INDEX_CACHE_BYTES: usize = 1024 * 1024 * 1024;
1628const REMOTE_METADATA_CACHE_BYTES: usize = 512 * 1024 * 1024;
1629
1630fn resolve_cache_caps(location: &Url, caps: RuntimeCaps) -> (usize, usize) {
1631 let (index_default, metadata_default) = if config::is_local(location) {
1632 (LOCAL_INDEX_CACHE_BYTES, LOCAL_METADATA_CACHE_BYTES)
1633 } else {
1634 (REMOTE_INDEX_CACHE_BYTES, REMOTE_METADATA_CACHE_BYTES)
1635 };
1636 (
1637 caps.index_cache_bytes.unwrap_or(index_default),
1638 caps.metadata_cache_bytes.unwrap_or(metadata_default),
1639 )
1640}
1641
1642pub struct Handle {
1643 datasets: DatasetSet,
1644 retry: RetryPolicy,
1645 #[allow(dead_code)]
1653 session: Arc<Session>,
1654 nm: Arc<dyn LanceNamespace>,
1658 nm_ident: NamespaceIdent,
1662 storage_options: HashMap<String, String>,
1667 location: Url,
1671 lazy_refresh_after: Duration,
1675 store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
1679}
1680
1681impl std::fmt::Debug for Handle {
1682 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1683 formatter
1684 .debug_struct("Handle")
1685 .field("datasets", &self.datasets)
1686 .field("retry", &self.retry)
1687 .field("nm_ident", &self.nm_ident)
1688 .field("storage_options", &self.storage_options)
1689 .field("location", &self.location)
1690 .finish()
1691 }
1692}
1693
1694#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1695pub enum Table {
1696 Sessions,
1697 Messages,
1698 Parts,
1699}
1700impl Table {
1701 pub fn as_str(self) -> &'static str {
1702 self.label()
1703 }
1704
1705 fn label(self) -> &'static str {
1706 match self {
1707 Self::Sessions => "sessions",
1708 Self::Messages => "messages",
1709 Self::Parts => "parts",
1710 }
1711 }
1712}
1713#[derive(Debug)]
1714struct DatasetSet {
1715 sessions: OnceCell<Mutex<CachedDataset>>,
1719 messages: Mutex<CachedDataset>,
1720 parts: OnceCell<Mutex<CachedDataset>>,
1728}
1729#[derive(Debug)]
1730struct CachedDataset {
1731 dataset: Dataset,
1732 last_refresh: Instant,
1733 refresh_after: Duration,
1734}
1735impl CachedDataset {
1736 fn new(dataset: Dataset, refresh_after: Duration) -> Self {
1737 Self {
1738 dataset,
1739 last_refresh: Instant::now(),
1740 refresh_after,
1741 }
1742 }
1743 async fn latest(&mut self) -> Result<Dataset> {
1744 if self.last_refresh.elapsed() >= self.refresh_after {
1745 self.dataset.checkout_latest().await?;
1746 self.last_refresh = Instant::now();
1747 }
1748 Ok(self.dataset.clone())
1749 }
1750 fn replace(&mut self, dataset: Dataset) {
1751 self.dataset = dataset;
1752 self.last_refresh = Instant::now();
1753 }
1754}
1755
1756#[derive(Debug, Clone, Copy, Default)]
1761pub struct AppendStats {
1762 pub rows: u64,
1763 pub bytes_written: u64,
1764 pub files_written: u64,
1765 pub attempts: u32,
1766}
1767
1768#[derive(Default)]
1774struct WriteAccum {
1775 rows: std::sync::atomic::AtomicU64,
1776 bytes: std::sync::atomic::AtomicU64,
1777 files: std::sync::atomic::AtomicU64,
1778}
1779
1780impl WriteAccum {
1781 fn observe(&self, stats: &WriteStats) {
1782 use std::sync::atomic::Ordering::Relaxed;
1783 self.rows.fetch_max(stats.rows_written, Relaxed);
1784 self.bytes.fetch_max(stats.bytes_written, Relaxed);
1785 self.files.fetch_max(stats.files_written as u64, Relaxed);
1786 }
1787 fn rows(&self) -> u64 {
1788 self.rows.load(std::sync::atomic::Ordering::Relaxed)
1789 }
1790 fn bytes(&self) -> u64 {
1791 self.bytes.load(std::sync::atomic::Ordering::Relaxed)
1792 }
1793 fn files(&self) -> u64 {
1794 self.files.load(std::sync::atomic::Ordering::Relaxed)
1795 }
1796}
1797
1798fn append_write_params() -> WriteParams {
1803 let mut params = sessions::write_params_for_create();
1804 params.mode = WriteMode::Append;
1805 params.max_bytes_per_file = TARGET_FRAGMENT_BYTES as usize;
1806 params
1807}
1808
1809impl Handle {
1810 pub async fn open(location: &Url) -> Result<Self> {
1813 Self::open_with_options(location, HashMap::new(), RuntimeCaps::default()).await
1814 }
1815
1816 pub fn lance_cache_bytes(&self) -> u64 {
1819 self.session.size_bytes()
1820 }
1821
1822 pub async fn open_with_options(
1829 location: &Url,
1830 storage_options: HashMap<String, String>,
1831 caps: RuntimeCaps,
1832 ) -> Result<Self> {
1833 Self::open_with_options_cached(location, storage_options, caps, None).await
1834 }
1835
1836 pub async fn open_with_options_cached(
1840 location: &Url,
1841 mut storage_options: HashMap<String, String>,
1842 caps: RuntimeCaps,
1843 index_cache_dir: Option<PathBuf>,
1844 ) -> Result<Self> {
1845 if let Some(path) = config::local_path(location) {
1846 tokio::fs::create_dir_all(&path).await.with_context(|| {
1847 format!(
1848 "failed to create data dir {}; fix the storage destination ([storage].path in config) or re-run `pond init`",
1849 path.display()
1850 )
1851 })?;
1852 } else {
1853 apply_remote_storage_defaults(&mut storage_options);
1854 }
1855 let (index_cache_bytes, metadata_cache_bytes) = resolve_cache_caps(location, caps);
1861 let session = Arc::new(Session::new(
1862 index_cache_bytes,
1863 metadata_cache_bytes,
1864 Arc::new(ObjectStoreRegistry::default()),
1865 ));
1866 let root = location.as_str().trim_end_matches('/').to_string();
1872 let mut connect = ConnectBuilder::new("dir")
1873 .property("root", root)
1874 .session(session.clone());
1875 for (key, value) in &storage_options {
1879 connect = connect.property(format!("storage.{key}"), value.clone());
1880 }
1881 let nm: Arc<dyn LanceNamespace> = connect
1882 .connect()
1883 .await
1884 .context("failed to connect lance Directory namespace")?;
1885 let nm_ident = NamespaceIdent::root();
1886 let refresh_after = if config::is_local(location) {
1892 Duration::ZERO
1893 } else {
1894 Duration::from_secs(5)
1895 };
1896 let wrapper = store_wrapper(location, index_cache_dir.as_deref());
1897 let handle = Self {
1898 datasets: DatasetSet {
1899 sessions: OnceCell::new(),
1900 messages: Mutex::new(CachedDataset::new(
1901 open_or_create_via_ns(
1902 &nm,
1903 &nm_ident,
1904 sessions::MESSAGES,
1905 sessions::message_schema(),
1906 &session,
1907 &storage_options,
1908 wrapper.clone(),
1909 )
1910 .await?,
1911 refresh_after,
1912 )),
1913 parts: OnceCell::new(),
1914 },
1915 retry: RetryPolicy::default(),
1916 session,
1917 nm,
1918 nm_ident,
1919 storage_options,
1920 location: location.clone(),
1921 lazy_refresh_after: refresh_after,
1922 store_wrapper: wrapper,
1923 };
1924 Ok(handle)
1925 }
1926
1927 pub fn location(&self) -> &Url {
1928 &self.location
1929 }
1930
1931 pub fn storage_options(&self) -> &HashMap<String, String> {
1935 &self.storage_options
1936 }
1937
1938 fn export_uri(&self, name: &str) -> String {
1944 format!(
1945 "{}/exports/{name}",
1946 self.location.as_str().trim_end_matches('/')
1947 )
1948 }
1949
1950 fn object_store_params(&self) -> ObjectStoreParams {
1954 ObjectStoreParams {
1955 storage_options_accessor: (!self.storage_options.is_empty()).then(|| {
1956 Arc::new(StorageOptionsAccessor::with_static_options(
1957 self.storage_options.clone(),
1958 ))
1959 }),
1960 ..Default::default()
1961 }
1962 }
1963
1964 pub(crate) async fn export_write(&self, name: &str, bytes: &[u8]) -> Result<()> {
1967 let uri = self.export_uri(name);
1968 let registry = Arc::new(ObjectStoreRegistry::default());
1969 let (store, path) =
1970 ObjectStore::from_uri_and_params(registry, &uri, &self.object_store_params())
1971 .await
1972 .with_context(|| format!("failed to open object store for {uri}"))?;
1973 store
1974 .put(&path, bytes)
1975 .await
1976 .with_context(|| format!("failed to write export {uri}"))?;
1977 Ok(())
1978 }
1979
1980 pub(crate) async fn export_read(&self, name: &str) -> Result<Vec<u8>> {
1983 let uri = self.export_uri(name);
1984 let registry = Arc::new(ObjectStoreRegistry::default());
1985 let (store, path) =
1986 ObjectStore::from_uri_and_params(registry, &uri, &self.object_store_params())
1987 .await
1988 .with_context(|| format!("failed to open object store for {uri}"))?;
1989 let bytes = store
1990 .read_one_all(&path)
1991 .await
1992 .with_context(|| format!("failed to read export {uri}"))?;
1993 Ok(bytes.to_vec())
1994 }
1995
1996 pub(crate) fn export_local_path(&self, name: &str) -> Option<std::path::PathBuf> {
2001 if self.location.scheme() != "file" {
2002 return None;
2003 }
2004 let dir = self.location.to_file_path().ok()?;
2005 Some(dir.join("exports").join(name))
2006 }
2007
2008 pub async fn row_counts(&self) -> Result<(usize, usize, usize)> {
2009 Ok((
2010 self.count_rows(Table::Sessions).await?,
2011 self.count_rows(Table::Messages).await?,
2012 self.count_rows(Table::Parts).await?,
2013 ))
2014 }
2015
2016 pub(crate) async fn merge_insert(
2020 &self,
2021 table: Table,
2022 batch: RecordBatch,
2023 row_count: usize,
2024 ) -> Result<u64> {
2025 self.merge_insert_stats(table, batch, row_count)
2026 .await
2027 .map(|stats| stats.num_inserted_rows + stats.num_updated_rows)
2028 }
2029
2030 pub(crate) async fn merge_insert_stats(
2035 &self,
2036 table: Table,
2037 batch: RecordBatch,
2038 row_count: usize,
2039 ) -> Result<MergeStats> {
2040 self.merge(
2041 table,
2042 batch,
2043 row_count,
2044 "merge_insert",
2045 WhenMatched::DoNothing,
2046 WhenNotMatched::InsertAll,
2047 )
2048 .await
2049 }
2050
2051 pub(crate) async fn merge_update(
2054 &self,
2055 table: Table,
2056 batch: RecordBatch,
2057 row_count: usize,
2058 ) -> Result<u64> {
2059 self.merge(
2060 table,
2061 batch,
2062 row_count,
2063 "merge_update",
2064 WhenMatched::UpdateAll,
2065 WhenNotMatched::DoNothing,
2066 )
2067 .await
2068 .map(|stats| stats.num_inserted_rows + stats.num_updated_rows)
2069 }
2070
2071 async fn write_committed<E, Fut, P>(&self, table: Table, execute: E) -> Result<P>
2080 where
2081 E: Fn(Arc<Dataset>) -> Fut,
2082 Fut: std::future::Future<Output = Result<(Dataset, P)>>,
2083 {
2084 self.write_committed_with(table, |_| true, execute).await
2085 }
2086
2087 async fn write_committed_with<E, Fut, P, R>(
2093 &self,
2094 table: Table,
2095 should_retry: R,
2096 execute: E,
2097 ) -> Result<P>
2098 where
2099 E: Fn(Arc<Dataset>) -> Fut,
2100 Fut: std::future::Future<Output = Result<(Dataset, P)>>,
2101 R: Fn(&anyhow::Error) -> bool,
2102 {
2103 self.retry_lance_filtered(table.label(), should_retry, || {
2104 let execute = &execute;
2105 async move {
2106 let mut cached = self.cached(table).await?.lock().await;
2107 let existing = cached.latest().await?;
2108 let (dataset, payload) = execute(Arc::new(existing)).await?;
2109 cached.replace(dataset);
2110 Ok(payload)
2111 }
2112 })
2113 .await
2114 }
2115
2116 async fn merge(
2122 &self,
2123 table: Table,
2124 batch: RecordBatch,
2125 row_count: usize,
2126 op: &'static str,
2127 when_matched: WhenMatched,
2128 when_not_matched: WhenNotMatched,
2129 ) -> Result<MergeStats> {
2130 if row_count == 0 {
2131 return Ok(MergeStats::default());
2132 }
2133 let started = Instant::now();
2134 let result = self
2135 .write_committed(table, |existing| {
2136 let batch = batch.clone();
2137 let when_matched = when_matched.clone();
2138 let when_not_matched = when_not_matched.clone();
2139 async move {
2140 let schema = batch.schema();
2141 let reader = RecordBatchIterator::new([Ok(batch)], schema);
2142 let mut builder = MergeInsertBuilder::try_new(existing, Vec::new())?;
2143 builder.when_matched(when_matched);
2144 builder.when_not_matched(when_not_matched);
2145 builder.source_dedupe_behavior(SourceDedupeBehavior::FirstSeen);
2148 builder.skip_auto_cleanup(true);
2152 let (dataset, stats) = builder
2153 .try_build()?
2154 .execute_reader(Box::new(reader))
2155 .await?;
2156 Ok((dataset.as_ref().clone(), stats))
2157 }
2158 })
2159 .await;
2160 let skipped = result
2161 .as_ref()
2162 .map(|s| s.num_skipped_duplicates)
2163 .unwrap_or(0);
2164 tracing::info!(
2165 target: "pond::perf",
2166 op,
2167 table = %table.label(),
2168 rows = row_count,
2169 elapsed_ms = started.elapsed().as_millis() as u64,
2170 skipped,
2171 "merge",
2172 );
2173 result
2174 }
2175
2176 pub(crate) async fn append_stream<F, Fut>(
2197 &self,
2198 table: Table,
2199 make_source: F,
2200 ) -> Result<AppendStats>
2201 where
2202 F: Fn() -> Fut,
2203 Fut: std::future::Future<Output = Result<SendableRecordBatchStream>>,
2204 {
2205 let cum = Arc::new(WriteAccum::default());
2206 let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
2207 let started = Instant::now();
2208 self.write_committed(table, |existing| {
2209 let make_source = &make_source;
2210 let cum = cum.clone();
2211 let attempts = attempts.clone();
2212 async move {
2213 attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2214 let stream = make_source().await?;
2215 let dataset = InsertBuilder::new(existing)
2216 .with_params(&append_write_params())
2217 .progress(move |stats| cum.observe(&stats))
2218 .execute_stream(stream)
2219 .await?;
2220 Ok((dataset, ()))
2221 }
2222 })
2223 .await?;
2224
2225 let attempts = attempts.load(std::sync::atomic::Ordering::Relaxed);
2226 let stats = AppendStats {
2227 rows: cum.rows(),
2228 bytes_written: cum.bytes(),
2229 files_written: cum.files(),
2230 attempts,
2231 };
2232 tracing::info!(
2233 target: "pond::perf",
2234 op = "append",
2235 table = %table.label(),
2236 rows = stats.rows,
2237 files = stats.files_written,
2238 attempts,
2239 elapsed_ms = started.elapsed().as_millis() as u64,
2240 "append",
2241 );
2242 Ok(stats)
2243 }
2244
2245 pub(crate) async fn append_batches(
2256 &self,
2257 table: Table,
2258 batches: Vec<RecordBatch>,
2259 ) -> Result<AppendStats> {
2260 let total_rows: u64 = batches.iter().map(|batch| batch.num_rows() as u64).sum();
2261 if total_rows == 0 {
2262 return Ok(AppendStats::default());
2263 }
2264 let cum = Arc::new(WriteAccum::default());
2265 let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
2266 let started = Instant::now();
2267 self.write_committed_with(table, is_commit_conflict, |existing| {
2268 let cum = cum.clone();
2269 let attempts = attempts.clone();
2270 let batches = batches.clone();
2271 async move {
2272 attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2273 let dataset = InsertBuilder::new(existing)
2274 .with_params(&append_write_params())
2275 .progress(move |stats| cum.observe(&stats))
2276 .execute(batches)
2277 .await?;
2278 Ok((dataset, ()))
2279 }
2280 })
2281 .await?;
2282
2283 let attempts = attempts.load(std::sync::atomic::Ordering::Relaxed);
2284 let stats = AppendStats {
2285 rows: total_rows,
2286 bytes_written: cum.bytes(),
2287 files_written: cum.files(),
2288 attempts,
2289 };
2290 tracing::info!(
2291 target: "pond::perf",
2292 op = "append_batches",
2293 table = %table.label(),
2294 rows = stats.rows,
2295 files = stats.files_written,
2296 attempts,
2297 elapsed_ms = started.elapsed().as_millis() as u64,
2298 "append",
2299 );
2300 Ok(stats)
2301 }
2302
2303 pub async fn optimize_table(
2312 &self,
2313 table: Table,
2314 intents: &[IndexIntent],
2315 progress: Option<&OptimizeProgressFn>,
2316 policy: &MaintenancePolicy,
2317 ) -> TableOptimizeOutcome {
2318 let compaction = self
2319 .run_optimize_compact_phase(table, progress, policy)
2320 .await;
2321 let indices = self
2322 .run_optimize_indices_phase(table, intents, progress, policy.fold_thresholds())
2323 .await;
2324 TableOptimizeOutcome {
2325 table,
2326 indices,
2327 compaction,
2328 }
2329 }
2330
2331 pub async fn optimize_table_indices_only(
2336 &self,
2337 table: Table,
2338 intents: &[IndexIntent],
2339 progress: Option<&OptimizeProgressFn>,
2340 ) -> PhaseOutcome {
2341 self.run_optimize_indices_phase(
2345 table,
2346 intents,
2347 progress,
2348 FoldThresholds {
2349 scalar: 0,
2350 index: 0,
2351 },
2352 )
2353 .await
2354 }
2355
2356 async fn run_optimize_indices_phase(
2357 &self,
2358 table: Table,
2359 intents: &[IndexIntent],
2360 progress: Option<&OptimizeProgressFn>,
2361 folds: FoldThresholds,
2362 ) -> PhaseOutcome {
2363 if intents.is_empty() {
2364 return PhaseOutcome::Noop;
2365 }
2366 let result = self
2367 .retry_lance(table.label(), || async {
2368 let mut guard = self.cached(table).await?.lock().await;
2369 let mut dataset = guard.latest().await?;
2370 let did_work =
2371 optimize_table_indices(&mut dataset, intents, table, progress, folds).await?;
2372 guard.replace(dataset);
2373 Ok::<_, anyhow::Error>(did_work)
2374 })
2375 .await;
2376 match result {
2377 Ok(true) => PhaseOutcome::Ok,
2378 Ok(false) => PhaseOutcome::Noop,
2379 Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
2380 Err(error) => PhaseOutcome::Failed(error),
2381 }
2382 }
2383
2384 async fn run_optimize_compact_phase(
2385 &self,
2386 table: Table,
2387 progress: Option<&OptimizeProgressFn>,
2388 policy: &MaintenancePolicy,
2389 ) -> PhaseOutcome {
2390 let result = self
2391 .retry_lance(table.label(), || async {
2392 let mut guard = self.cached(table).await?.lock().await;
2393 let mut dataset = guard.latest().await?;
2394 optimize_table_compact(&mut dataset, table, progress, policy).await?;
2395 guard.replace(dataset);
2396 Ok::<_, anyhow::Error>(())
2397 })
2398 .await;
2399 match result {
2400 Ok(()) => PhaseOutcome::Ok,
2401 Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
2402 Err(error) => PhaseOutcome::Failed(error),
2403 }
2404 }
2405
2406 pub async fn rebuild_index(
2407 &self,
2408 table: Table,
2409 intent: &IndexIntent,
2410 progress: Option<&OptimizeProgressFn>,
2411 ) -> Result<()> {
2412 emit(
2413 progress,
2414 OptimizeEvent::PhaseStart {
2415 table,
2416 phase: OptimizePhase::IndexRebuild,
2417 detail: Some(intent.name.to_owned()),
2418 },
2419 );
2420 let started = Instant::now();
2421 let result = self
2422 .retry_lance(table.label(), || async {
2423 let mut guard = self.cached(table).await?.lock().await;
2424 let mut dataset = guard.latest().await?;
2425 rebuild_index(&mut dataset, intent, progress, table).await?;
2426 guard.replace(dataset);
2427 Ok(())
2428 })
2429 .await;
2430 emit(
2431 progress,
2432 OptimizeEvent::PhaseDone {
2433 table,
2434 phase: OptimizePhase::IndexRebuild,
2435 elapsed_ms: started.elapsed().as_millis() as u64,
2436 },
2437 );
2438 result
2439 }
2440
2441 pub async fn cleanup_table_versions(
2445 &self,
2446 table: Table,
2447 older_than: chrono::Duration,
2448 ) -> Result<()> {
2449 let mut guard = self.cached(table).await?.lock().await;
2450 let dataset = guard.latest().await?;
2451 dataset
2452 .cleanup_old_versions(older_than, Some(false), Some(false))
2453 .await
2454 .with_context(|| format!("cleanup_old_versions failed for {}", table.label()))?;
2455 Ok(())
2456 }
2457
2458 pub async fn index_status(
2459 &self,
2460 table: Table,
2461 intents: &[IndexIntent],
2462 indexable_only: bool,
2463 ) -> Result<Vec<IndexStatus>> {
2464 let dataset = self.dataset(table).await?;
2465 index_status(table, &dataset, intents, indexable_only).await
2466 }
2467
2468 pub(crate) async fn dataset(&self, table: Table) -> Result<Dataset> {
2469 let mut cached = self.cached(table).await?.lock().await;
2470 cached.latest().await
2471 }
2472 pub(crate) async fn scanner(
2477 &self,
2478 table: Table,
2479 predicate: Option<&Predicate>,
2480 ) -> Result<lance::dataset::scanner::Scanner> {
2481 let dataset = self.dataset(table).await?;
2482 scanner_with_prefilter(&dataset, predicate)
2483 }
2484 pub async fn scan(
2487 &self,
2488 table: Table,
2489 opts: ScanOpts<'_>,
2490 ) -> Result<lance::dataset::scanner::Scanner> {
2491 let mut scanner = self.scanner(table, opts.predicate).await?;
2492 if let Some(projection) = opts.projection {
2493 scanner.project(projection)?;
2494 }
2495 Ok(scanner)
2496 }
2497 pub(crate) async fn scan_batch(
2498 &self,
2499 table: Table,
2500 predicate: Option<&Predicate>,
2501 projection: &[&str],
2502 ) -> Result<RecordBatch> {
2503 let opts = ScanOpts {
2504 predicate,
2505 projection: (!projection.is_empty()).then_some(projection),
2506 };
2507 self.scan(table, opts)
2508 .await?
2509 .try_into_batch()
2510 .await
2511 .context("scan failed")
2512 }
2513 pub async fn count_rows(&self, table: Table) -> Result<usize> {
2514 self.dataset(table)
2515 .await?
2516 .count_rows(None)
2517 .await
2518 .map_err(Into::into)
2519 }
2520 pub async fn collect_ids(&self, table: Table) -> Result<std::collections::HashSet<String>> {
2526 let batch = self.scan_batch(table, None, &["id"]).await?;
2527 let ids = batch
2528 .column_by_name("id")
2529 .context("scan projection dropped the id column")?
2530 .as_any()
2531 .downcast_ref::<StringArray>()
2532 .context("id column is not Utf8")?;
2533 Ok(ids.iter().flatten().map(str::to_owned).collect())
2534 }
2535 #[cfg(test)]
2537 pub(crate) async fn messages_index_names(&self) -> Result<Vec<String>> {
2538 let dataset = self.dataset(Table::Messages).await?;
2539 let indices = dataset.load_indices().await?;
2540 Ok(indices.iter().map(|index| index.name.clone()).collect())
2541 }
2542
2543 pub(crate) async fn messages_has_index(&self, name: &str) -> Result<bool> {
2549 let dataset = self.dataset(Table::Messages).await?;
2550 let indices = dataset.load_indices().await?;
2551 Ok(indices.iter().any(|index| index.name == name))
2552 }
2553
2554 pub(crate) async fn messages_fast_search_ready(&self, name: &str) -> Result<bool> {
2561 let dataset = self.dataset(Table::Messages).await?;
2562 if !dataset
2563 .load_indices()
2564 .await?
2565 .iter()
2566 .any(|index| index.name == name)
2567 {
2568 return Ok(false);
2569 }
2570 let unindexed = dataset
2571 .unindexed_fragments(name)
2572 .await
2573 .with_context(|| format!("unindexed_fragments failed for {name}"))?;
2574 Ok(unindexed.is_empty())
2575 }
2576
2577 pub(crate) async fn prune_index_cache(&self, cache_dir: &std::path::Path) {
2582 if config::is_local(&self.location) {
2583 return;
2584 }
2585 let root = cache_dir.join(store_key(&self.location)).join("indices");
2586 if !root.exists() {
2587 return;
2588 }
2589 let mut keep = std::collections::HashSet::new();
2590 for table in [Table::Sessions, Table::Messages, Table::Parts] {
2591 let Ok(dataset) = self.dataset(table).await else {
2592 return;
2593 };
2594 let Ok(indices) = dataset.load_indices().await else {
2595 return;
2596 };
2597 keep.extend(indices.iter().map(|index| index.uuid.to_string()));
2598 }
2599 prune_stale_uuid_dirs(&root, &keep);
2600 }
2601
2602 pub(crate) async fn unindexed_row_count(
2605 &self,
2606 table: Table,
2607 index_name: &str,
2608 ) -> Result<usize> {
2609 let dataset = self.dataset(table).await?;
2610 let fragments = dataset
2611 .unindexed_fragments(index_name)
2612 .await
2613 .with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
2614 Ok(fragments
2615 .iter()
2616 .map(|fragment| fragment.num_rows().unwrap_or(0))
2617 .sum())
2618 }
2619
2620 pub(crate) async fn find_index_owner(&self, name: &str) -> Result<Option<Table>> {
2627 let list = |table: Table| async move {
2628 let dataset = self.dataset(table).await?;
2629 let names: Vec<String> = dataset
2630 .load_indices()
2631 .await
2632 .with_context(|| format!("load_indices failed for {}", table.label()))?
2633 .iter()
2634 .map(|index| index.name.clone())
2635 .collect();
2636 Ok::<_, anyhow::Error>(names)
2637 };
2638 let (sessions, messages, parts) = tokio::try_join!(
2639 list(Table::Sessions),
2640 list(Table::Messages),
2641 list(Table::Parts),
2642 )?;
2643 for (table, names) in [
2644 (Table::Sessions, sessions),
2645 (Table::Messages, messages),
2646 (Table::Parts, parts),
2647 ] {
2648 if names.iter().any(|n| n == name) {
2649 return Ok(Some(table));
2650 }
2651 }
2652 Ok(None)
2653 }
2654
2655 pub(crate) async fn drop_index(&self, table: Table, name: &str) -> Result<()> {
2661 let mut guard = self.cached(table).await?.lock().await;
2662 let mut dataset = guard.latest().await?;
2663 dataset
2664 .drop_index(name)
2665 .await
2666 .with_context(|| format!("drop_index({name}) failed for {}", table.label()))?;
2667 guard.replace(dataset);
2668 Ok(())
2669 }
2670
2671 async fn table_location(&self, table_name: &str) -> Result<String> {
2674 let request = DescribeTableRequest {
2675 id: Some(self.nm_ident.as_table_id(table_name)),
2676 ..Default::default()
2677 };
2678 let response = self
2679 .nm
2680 .describe_table(request)
2681 .await
2682 .with_context(|| format!("failed to describe table {table_name}"))?;
2683 response
2684 .location
2685 .with_context(|| format!("namespace returned no location for table {table_name}"))
2686 }
2687
2688 pub async fn initialized(&self) -> Result<bool> {
2694 let request = DescribeTableRequest {
2695 id: Some(self.nm_ident.as_table_id(sessions::PARTS)),
2696 ..Default::default()
2697 };
2698 match self.nm.describe_table(request).await {
2699 Ok(_) => Ok(true),
2700 Err(error) if is_namespace_error_code(&error, ErrorCode::TableNotFound) => Ok(false),
2701 Err(error) => {
2702 Err(anyhow::Error::from(error)).context("failed to probe table existence")
2703 }
2704 }
2705 }
2706
2707 pub async fn table_sizes(&self) -> Result<TableSizes> {
2711 let registry = Arc::new(ObjectStoreRegistry::default());
2712 let params = self.object_store_params();
2713
2714 let sessions = self
2715 .listed_size(
2716 ®istry,
2717 ¶ms,
2718 &self.table_location(sessions::SESSIONS).await?,
2719 )
2720 .await?;
2721 let messages = self
2722 .listed_size(
2723 ®istry,
2724 ¶ms,
2725 &self.table_location(sessions::MESSAGES).await?,
2726 )
2727 .await?;
2728 let parts = self
2729 .listed_size(
2730 ®istry,
2731 ¶ms,
2732 &self.table_location(sessions::PARTS).await?,
2733 )
2734 .await?;
2735 let root_total = self
2738 .listed_size(®istry, ¶ms, self.location.as_str())
2739 .await?;
2740 let other = root_total.saturating_sub(sessions + messages + parts);
2741 let sessions_data = self
2742 .data_liveness(®istry, ¶ms, Table::Sessions, sessions::SESSIONS)
2743 .await?;
2744 let messages_data = self
2745 .data_liveness(®istry, ¶ms, Table::Messages, sessions::MESSAGES)
2746 .await?;
2747 let parts_data = self
2748 .data_liveness(®istry, ¶ms, Table::Parts, sessions::PARTS)
2749 .await?;
2750 Ok(TableSizes {
2751 sessions,
2752 messages,
2753 parts,
2754 other,
2755 sessions_data,
2756 messages_data,
2757 parts_data,
2758 })
2759 }
2760
2761 async fn data_liveness(
2762 &self,
2763 registry: &Arc<ObjectStoreRegistry>,
2764 params: &ObjectStoreParams,
2765 table: Table,
2766 table_name: &str,
2767 ) -> Result<DataLiveness> {
2768 let location = self.table_location(table_name).await?;
2769 let data_dir = format!("{}/data", location.trim_end_matches('/'));
2770 let on_disk = self.listed_size(registry, params, &data_dir).await?;
2771 let dataset = self.dataset(table).await?;
2772 let live = dataset
2773 .get_fragments()
2774 .iter()
2775 .try_fold(0u64, |total, fragment| {
2776 Some(total + fragment_bytes(fragment.metadata())?)
2777 });
2778 Ok(DataLiveness { on_disk, live })
2779 }
2780
2781 async fn listed_size(
2783 &self,
2784 registry: &Arc<ObjectStoreRegistry>,
2785 params: &ObjectStoreParams,
2786 uri: &str,
2787 ) -> Result<u64> {
2788 let (store, base) = ObjectStore::from_uri_and_params(registry.clone(), uri, params)
2789 .await
2790 .with_context(|| format!("failed to open object store for {uri}"))?;
2791 let mut listing = store.list(Some(base));
2792 let mut total = 0u64;
2793 while let Some(meta) = listing.next().await {
2794 let meta = meta.with_context(|| format!("listing {uri} failed"))?;
2795 total += meta.size;
2796 }
2797 Ok(total)
2798 }
2799 async fn cached(&self, table: Table) -> Result<&Mutex<CachedDataset>> {
2800 match table {
2801 Table::Sessions => self.sessions_cached().await,
2802 Table::Messages => Ok(&self.datasets.messages),
2803 Table::Parts => self.parts_cached().await,
2804 }
2805 }
2806
2807 async fn sessions_cached(&self) -> Result<&Mutex<CachedDataset>> {
2812 self.lazy_cached(
2813 &self.datasets.sessions,
2814 sessions::SESSIONS,
2815 sessions::session_schema,
2816 )
2817 .await
2818 }
2819
2820 async fn parts_cached(&self) -> Result<&Mutex<CachedDataset>> {
2823 self.lazy_cached(&self.datasets.parts, sessions::PARTS, sessions::part_schema)
2824 .await
2825 }
2826
2827 async fn lazy_cached<'a>(
2831 &self,
2832 cell: &'a OnceCell<Mutex<CachedDataset>>,
2833 table_name: &str,
2834 schema: fn() -> lance::deps::arrow_schema::SchemaRef,
2835 ) -> Result<&'a Mutex<CachedDataset>> {
2836 cell.get_or_try_init(|| async {
2837 let dataset = open_or_create_via_ns(
2838 &self.nm,
2839 &self.nm_ident,
2840 table_name,
2841 schema(),
2842 &self.session,
2843 &self.storage_options,
2844 self.store_wrapper.clone(),
2845 )
2846 .await?;
2847 Ok::<_, anyhow::Error>(Mutex::new(CachedDataset::new(
2848 dataset,
2849 self.lazy_refresh_after,
2850 )))
2851 })
2852 .await
2853 }
2854 async fn retry_lance<T, Fut, Op>(&self, label: &str, operation: Op) -> Result<T>
2855 where
2856 Fut: std::future::Future<Output = Result<T>>,
2857 Op: FnMut() -> Fut,
2858 {
2859 self.retry_lance_filtered(label, |_| true, operation).await
2861 }
2862
2863 async fn retry_lance_filtered<T, Fut, Op, R>(
2873 &self,
2874 label: &str,
2875 should_retry: R,
2876 mut operation: Op,
2877 ) -> Result<T>
2878 where
2879 Fut: std::future::Future<Output = Result<T>>,
2880 Op: FnMut() -> Fut,
2881 R: Fn(&anyhow::Error) -> bool,
2882 {
2883 let mut attempt = 0u8;
2884 loop {
2885 attempt = attempt.saturating_add(1);
2886 match operation().await {
2887 Ok(value) => return Ok(value),
2888 Err(error) if attempt < self.retry.attempts && should_retry(&error) => {
2889 let backoff = self.backoff(attempt);
2890 let error_chain = format!("{error:#}");
2893 tracing::warn!(
2894 label,
2895 attempt,
2896 ?backoff,
2897 error = %error_chain,
2898 "retrying Lance operation"
2899 );
2900 tokio::time::sleep(backoff).await;
2901 }
2902 Err(error) => {
2903 let error_chain = format!("{error:#}");
2904 tracing::warn!(
2905 label,
2906 attempt,
2907 error = %error_chain,
2908 "Lance operation exhausted retries"
2909 );
2910 if is_commit_conflict(&error) {
2917 return Err(error.context(ConflictExhausted { attempts: attempt }));
2918 }
2919 return Err(error);
2920 }
2921 }
2922 }
2923 }
2924 fn backoff(&self, attempt: u8) -> Duration {
2925 let shift = u32::from(attempt.saturating_sub(1));
2926 let multiplier = 1u32.checked_shl(shift).unwrap_or(u32::MAX);
2927 let base = self.retry.initial_backoff.saturating_mul(multiplier);
2928 let factor = (1.0 + self.retry.jitter * (fastrand::f64() * 2.0 - 1.0)).max(0.0);
2931 base.mul_f64(factor).min(self.retry.max_backoff)
2932 }
2933}
2934async fn optimize_table_compact(
2955 dataset: &mut Dataset,
2956 table: Table,
2957 progress: Option<&OptimizeProgressFn>,
2958 policy: &MaintenancePolicy,
2959) -> Result<()> {
2960 let stats: Vec<FragmentStat> = dataset
2961 .get_fragments()
2962 .iter()
2963 .map(|fragment| fragment_stat(fragment.metadata()))
2964 .collect();
2965 let compaction = CompactionOptions {
2966 target_rows_per_fragment: derived_target_rows(&stats),
2967 max_bytes_per_file: Some(TARGET_FRAGMENT_BYTES as usize),
2968 defer_index_remap: false,
2969 compaction_mode: Some(CompactionMode::TryBinaryCopy),
2974 ..CompactionOptions::default()
2975 };
2976
2977 let mut plan = plan_compaction(dataset, &compaction).await?;
2978 if policy.compaction_fragment_cap > 0 {
2979 plan.tasks.retain(|task| {
2980 let task_stats: Vec<FragmentStat> = task.fragments.iter().map(fragment_stat).collect();
2981 let keep = keep_task(
2982 &task_stats,
2983 policy.compaction_fragment_cap,
2984 compaction.materialize_deletions_threshold,
2985 );
2986 if !keep {
2987 tracing::debug!(
2988 target: "pond::perf",
2989 table = table.as_str(),
2990 fragments = task_stats.len(),
2991 "compaction task vetoed: merge dominated by one large fragment",
2992 );
2993 }
2994 keep
2995 });
2996 }
2997 if plan.tasks.is_empty() {
2998 tracing::debug!(
2999 target: "pond::perf",
3000 table = table.as_str(),
3001 "compaction skipped: no task to run",
3002 );
3003 } else {
3004 emit(
3005 progress,
3006 OptimizeEvent::PhaseStart {
3007 table,
3008 phase: OptimizePhase::Compact,
3009 detail: None,
3010 },
3011 );
3012 let started = Instant::now();
3013 let mut completed = Vec::with_capacity(plan.tasks.len());
3014 for task in plan.compaction_tasks() {
3015 completed.push(task.execute(dataset).await?);
3016 }
3017 commit_compaction(
3018 dataset,
3019 completed,
3020 Arc::new(DatasetIndexRemapperOptions::default()),
3021 &compaction,
3022 )
3023 .await?;
3024 emit(
3025 progress,
3026 OptimizeEvent::PhaseDone {
3027 table,
3028 phase: OptimizePhase::Compact,
3029 elapsed_ms: started.elapsed().as_millis() as u64,
3030 },
3031 );
3032 }
3033
3034 if cleanup_due(dataset.version_id(), policy.cleanup_interval) {
3045 emit(
3046 progress,
3047 OptimizeEvent::PhaseStart {
3048 table,
3049 phase: OptimizePhase::Cleanup,
3050 detail: None,
3051 },
3052 );
3053 let started = Instant::now();
3054 dataset
3063 .cleanup_old_versions(policy.cleanup_older_than, Some(false), Some(false))
3064 .await
3065 .context("cleanup_old_versions failed during index optimize")?;
3066 emit(
3067 progress,
3068 OptimizeEvent::PhaseDone {
3069 table,
3070 phase: OptimizePhase::Cleanup,
3071 elapsed_ms: started.elapsed().as_millis() as u64,
3072 },
3073 );
3074 }
3075
3076 Ok(())
3077}
3078
3079fn cleanup_due(version: u64, interval: u64) -> bool {
3085 interval <= 1 || version.is_multiple_of(interval)
3086}
3087
3088async fn optimize_table_indices(
3093 dataset: &mut Dataset,
3094 intents: &[IndexIntent],
3095 table: Table,
3096 progress: Option<&OptimizeProgressFn>,
3097 folds: FoldThresholds,
3098) -> Result<bool> {
3099 let existing = dataset.load_indices().await?;
3100 let existing_names: std::collections::HashSet<String> =
3101 existing.iter().map(|index| index.name.clone()).collect();
3102
3103 let mut append_indices: Vec<String> = Vec::new();
3104 let mut did_work = false;
3105
3106 for intent in intents {
3107 let exists = existing_names.contains(intent.name);
3108
3109 if !exists {
3110 if !intent.trigger.should_create(dataset).await? {
3111 continue;
3112 }
3113 let params = intent.params.build(dataset).await?;
3114 let index_type = intent.params.index_type();
3115 tracing::info!(
3116 index = intent.name,
3117 column = intent.column,
3118 "creating Lance index (trigger fired)",
3119 );
3120 emit(
3121 progress,
3122 OptimizeEvent::PhaseStart {
3123 table,
3124 phase: OptimizePhase::IndexCreate,
3125 detail: Some(intent.name.to_owned()),
3126 },
3127 );
3128 let started = Instant::now();
3129 dataset
3130 .create_index_builder(&[intent.column], index_type, params.as_ref())
3131 .name(intent.name.to_owned())
3132 .replace(false)
3133 .progress(lance_progress(progress, table, intent.name))
3134 .await
3135 .with_context(|| format!("failed to create index {}", intent.name))?;
3136 emit(
3137 progress,
3138 OptimizeEvent::PhaseDone {
3139 table,
3140 phase: OptimizePhase::IndexCreate,
3141 elapsed_ms: started.elapsed().as_millis() as u64,
3142 },
3143 );
3144 did_work = true;
3145 continue;
3146 }
3147
3148 let unindexed = dataset.unindexed_fragments(intent.name).await?;
3154 if unindexed.is_empty() {
3155 continue;
3156 }
3157 let tail_rows: usize = unindexed
3158 .iter()
3159 .map(|fragment| fragment.num_rows().unwrap_or(0))
3160 .sum();
3161 let fold_threshold = match intent.params {
3170 IndexParamsKind::Scalar(_) => folds.scalar,
3171 IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. } => folds.index,
3172 };
3173 if fold_threshold > 0 && tail_rows < fold_threshold {
3174 tracing::debug!(
3175 target: "pond::perf",
3176 index = intent.name,
3177 tail_rows,
3178 threshold = fold_threshold,
3179 "deferring index fold (unindexed tail below threshold)",
3180 );
3181 continue;
3182 }
3183 if matches!(intent.params, IndexParamsKind::InvertedFtsWord)
3191 && !column_has_values(dataset, intent.column, &unindexed).await?
3192 {
3193 tracing::debug!(
3194 target: "pond::perf",
3195 index = intent.name,
3196 tail_rows,
3197 "skipping FTS fold (tail has no indexable values)",
3198 );
3199 continue;
3200 }
3201 append_indices.push(intent.name.to_owned());
3207 }
3208
3209 if !append_indices.is_empty() {
3210 let segment_count = |name: &str| {
3216 existing
3217 .iter()
3218 .filter(|index| index.name.as_str() == name)
3219 .count()
3220 };
3221 let (consolidate, to_append): (Vec<String>, Vec<String>) = append_indices
3222 .iter()
3223 .cloned()
3224 .partition(|name| segment_count(name) >= DELTA_MERGE_THRESHOLD);
3225 let mut fts_rebuilds: Vec<&IndexIntent> = Vec::new();
3234 let mut to_merge: Vec<String> = Vec::new();
3235 for name in consolidate {
3236 let fts_intent = intents.iter().find(|intent| {
3237 intent.name == name && matches!(intent.params, IndexParamsKind::InvertedFtsWord)
3238 });
3239 match fts_intent {
3240 Some(intent) => fts_rebuilds.push(intent),
3241 None => to_merge.push(name),
3242 }
3243 }
3244
3245 emit(
3246 progress,
3247 OptimizeEvent::PhaseStart {
3248 table,
3249 phase: OptimizePhase::IndexAppend,
3250 detail: Some(append_indices.join(", ")),
3251 },
3252 );
3253 let started = Instant::now();
3254 if !to_append.is_empty() {
3255 dataset
3256 .optimize_indices(&OptimizeOptions::append().index_names(to_append))
3257 .await
3258 .context("optimize_indices(append) failed during index optimize")?;
3259 }
3260 if !to_merge.is_empty() {
3261 dataset
3262 .optimize_indices(
3263 &OptimizeOptions::merge(DELTA_MERGE_THRESHOLD).index_names(to_merge),
3264 )
3265 .await
3266 .context("optimize_indices(merge) failed during index optimize")?;
3267 }
3268 emit(
3269 progress,
3270 OptimizeEvent::PhaseDone {
3271 table,
3272 phase: OptimizePhase::IndexAppend,
3273 elapsed_ms: started.elapsed().as_millis() as u64,
3274 },
3275 );
3276 for intent in &fts_rebuilds {
3277 emit(
3278 progress,
3279 OptimizeEvent::PhaseStart {
3280 table,
3281 phase: OptimizePhase::IndexRebuild,
3282 detail: Some(intent.name.to_owned()),
3283 },
3284 );
3285 let rebuild_started = Instant::now();
3286 rebuild_index(dataset, intent, progress, table).await?;
3287 emit(
3288 progress,
3289 OptimizeEvent::PhaseDone {
3290 table,
3291 phase: OptimizePhase::IndexRebuild,
3292 elapsed_ms: rebuild_started.elapsed().as_millis() as u64,
3293 },
3294 );
3295 }
3296 tracing::debug!(
3297 target: "pond::perf",
3298 indices = ?append_indices,
3299 rebuilt = ?fts_rebuilds,
3300 "folded trailing fragments into indices",
3301 );
3302 did_work = true;
3303 }
3304
3305 Ok(did_work)
3306}
3307
3308fn non_null_scanner(
3311 dataset: &Dataset,
3312 column: &'static str,
3313 fragments: &[lance::table::format::Fragment],
3314) -> Result<lance::dataset::scanner::Scanner> {
3315 let mut scanner = dataset.scan();
3316 scanner.with_fragments(fragments.to_vec());
3317 scanner.filter(&Predicate::IsNotNull(column).to_lance())?;
3318 Ok(scanner)
3319}
3320
3321async fn column_has_values(
3326 dataset: &Dataset,
3327 column: &'static str,
3328 fragments: &[lance::table::format::Fragment],
3329) -> Result<bool> {
3330 let mut scanner = non_null_scanner(dataset, column, fragments)?;
3331 scanner.project(&[column])?;
3332 scanner.limit(Some(1), None)?;
3333 let batch = scanner
3334 .try_into_batch()
3335 .await
3336 .with_context(|| format!("non-null probe on {column} failed"))?;
3337 Ok(batch.num_rows() > 0)
3338}
3339
3340async fn column_value_count(
3343 dataset: &Dataset,
3344 column: &'static str,
3345 fragments: &[lance::table::format::Fragment],
3346) -> Result<usize> {
3347 let count = non_null_scanner(dataset, column, fragments)?
3348 .count_rows()
3349 .await
3350 .with_context(|| format!("non-null count on {column} failed"))?;
3351 Ok(count as usize)
3352}
3353
3354async fn rebuild_index(
3355 dataset: &mut Dataset,
3356 intent: &IndexIntent,
3357 progress: Option<&OptimizeProgressFn>,
3358 table: Table,
3359) -> Result<()> {
3360 if !intent.trigger.should_create(dataset).await? {
3361 return Ok(());
3362 }
3363 let params = intent.params.build(dataset).await?;
3364 dataset
3365 .create_index_builder(
3366 &[intent.column],
3367 intent.params.index_type(),
3368 params.as_ref(),
3369 )
3370 .name(intent.name.to_owned())
3371 .replace(true)
3372 .progress(lance_progress(progress, table, intent.name))
3373 .await
3374 .with_context(|| format!("failed to rebuild index {}", intent.name))?;
3375 Ok(())
3376}
3377
3378async fn index_status(
3379 table: Table,
3380 dataset: &Dataset,
3381 intents: &[IndexIntent],
3382 indexable_only: bool,
3383) -> Result<Vec<IndexStatus>> {
3384 let existing = dataset.load_indices().await?;
3385 let existing_names: std::collections::HashSet<String> =
3386 existing.iter().map(|index| index.name.clone()).collect();
3387 let total_fragments = dataset.get_fragments().len();
3388 let total_rows = dataset.count_rows(None).await?;
3389 let mut statuses = Vec::with_capacity(intents.len());
3390 for intent in intents {
3391 let exists = existing_names.contains(intent.name);
3392 if !exists {
3393 statuses.push(IndexStatus {
3394 table,
3395 intent_name: intent.name.to_owned(),
3396 fragments_covered: 0,
3397 unindexed_fragments: total_fragments,
3398 unindexed_rows: total_rows,
3399 exists,
3400 });
3401 continue;
3402 }
3403 let unindexed = dataset
3404 .unindexed_fragments(intent.name)
3405 .await
3406 .with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
3407 let unindexed_fragments = unindexed.len();
3408 let mut unindexed_rows: usize = unindexed
3409 .iter()
3410 .map(|fragment| fragment.num_rows().unwrap_or(0))
3411 .sum();
3412 if indexable_only
3420 && unindexed_rows > 0
3421 && matches!(
3422 intent.params,
3423 IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. }
3424 )
3425 {
3426 unindexed_rows = column_value_count(dataset, intent.column, &unindexed).await?;
3427 }
3428 statuses.push(IndexStatus {
3429 table,
3430 intent_name: intent.name.to_owned(),
3431 fragments_covered: total_fragments.saturating_sub(unindexed_fragments),
3432 unindexed_fragments,
3433 unindexed_rows,
3434 exists,
3435 });
3436 }
3437 Ok(statuses)
3438}
3439
3440pub mod io_trace {
3458 use lance_io::utils::tracking_store::{IOTracker, IoStats};
3459 use std::sync::{Arc, OnceLock};
3460
3461 static TRACKER: OnceLock<IOTracker> = OnceLock::new();
3462
3463 pub fn enable() {
3466 let _ = TRACKER.set(IOTracker::default());
3467 }
3468
3469 pub(super) fn wrapper() -> Option<Arc<IOTracker>> {
3471 TRACKER.get().map(|tracker| Arc::new(tracker.clone()))
3472 }
3473
3474 pub fn take() -> Option<IoStats> {
3476 TRACKER.get().map(IOTracker::incremental_stats)
3477 }
3478}
3479
3480pub mod index_cache {
3488 use object_store::local::LocalFileSystem;
3489 use object_store::path::Path as ObjPath;
3490 use object_store::{
3491 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
3492 ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result as OsResult,
3493 };
3494 use std::collections::HashMap;
3495 use std::ops::Range;
3496 use std::path::PathBuf;
3497 use std::sync::{Arc, Mutex};
3498
3499 use bytes::Bytes;
3500 use futures::stream::BoxStream;
3501 use lance_io::object_store::WrappingObjectStore;
3502
3503 fn is_index_path(location: &ObjPath) -> bool {
3504 AsRef::<str>::as_ref(location).contains("_indices/")
3505 }
3506
3507 fn local_opts(options: &GetOptions) -> GetOptions {
3510 GetOptions {
3511 range: options.range.clone(),
3512 head: options.head,
3513 ..Default::default()
3514 }
3515 }
3516
3517 #[derive(Debug)]
3520 pub struct IndexDiskCache {
3521 local: Arc<LocalFileSystem>,
3522 inflight: Arc<Mutex<HashMap<ObjPath, Arc<tokio::sync::Mutex<()>>>>>,
3523 }
3524
3525 impl IndexDiskCache {
3526 pub fn new(root: PathBuf) -> std::io::Result<Self> {
3528 std::fs::create_dir_all(&root)?;
3529 Ok(Self {
3530 local: Arc::new(LocalFileSystem::new_with_prefix(&root)?),
3531 inflight: Arc::new(Mutex::new(HashMap::new())),
3532 })
3533 }
3534 }
3535
3536 impl WrappingObjectStore for IndexDiskCache {
3537 fn wrap(&self, _store_prefix: &str, inner: Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
3538 Arc::new(CachingStore {
3539 inner,
3540 local: self.local.clone(),
3541 inflight: self.inflight.clone(),
3542 })
3543 }
3544 }
3545
3546 #[derive(Debug)]
3547 struct CachingStore {
3548 inner: Arc<dyn ObjectStore>,
3549 local: Arc<LocalFileSystem>,
3550 inflight: Arc<Mutex<HashMap<ObjPath, Arc<tokio::sync::Mutex<()>>>>>,
3551 }
3552
3553 impl std::fmt::Display for CachingStore {
3554 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3555 write!(f, "CachingStore({})", self.inner)
3556 }
3557 }
3558
3559 impl CachingStore {
3560 fn flight_lock(&self, location: &ObjPath) -> Arc<tokio::sync::Mutex<()>> {
3561 self.inflight
3562 .lock()
3563 .unwrap_or_else(|poison| poison.into_inner())
3564 .entry(location.clone())
3565 .or_default()
3566 .clone()
3567 }
3568
3569 async fn populate_and_serve(
3575 &self,
3576 location: &ObjPath,
3577 options: GetOptions,
3578 ) -> OsResult<GetResult> {
3579 let lock = self.flight_lock(location);
3580 let _guard = lock.lock().await;
3581 let result = self.fetch_under_flight(location, options).await;
3582 self.inflight
3587 .lock()
3588 .unwrap_or_else(|p| p.into_inner())
3589 .remove(location);
3590 result
3591 }
3592
3593 async fn fetch_under_flight(
3594 &self,
3595 location: &ObjPath,
3596 options: GetOptions,
3597 ) -> OsResult<GetResult> {
3598 if let Ok(result) = self.local.get_opts(location, local_opts(&options)).await {
3599 return Ok(result);
3600 }
3601 let bytes = self.inner.get(location).await?.bytes().await?;
3602 if self
3603 .local
3604 .put(location, PutPayload::from_bytes(bytes))
3605 .await
3606 .is_ok()
3607 && let Ok(result) = self.local.get_opts(location, local_opts(&options)).await
3608 {
3609 return Ok(result);
3610 }
3611 self.inner.get_opts(location, options).await
3613 }
3614 }
3615
3616 #[async_trait::async_trait]
3617 impl ObjectStore for CachingStore {
3618 async fn get_opts(&self, location: &ObjPath, options: GetOptions) -> OsResult<GetResult> {
3619 if !is_index_path(location) {
3620 return self.inner.get_opts(location, options).await;
3621 }
3622 match self.local.get_opts(location, local_opts(&options)).await {
3623 Ok(result) => Ok(result),
3624 Err(object_store::Error::NotFound { .. }) => {
3625 self.populate_and_serve(location, options).await
3626 }
3627 Err(_) => self.inner.get_opts(location, options).await,
3628 }
3629 }
3630
3631 async fn put_opts(
3632 &self,
3633 location: &ObjPath,
3634 payload: PutPayload,
3635 opts: PutOptions,
3636 ) -> OsResult<PutResult> {
3637 self.inner.put_opts(location, payload, opts).await
3638 }
3639
3640 async fn put_multipart_opts(
3641 &self,
3642 location: &ObjPath,
3643 opts: PutMultipartOptions,
3644 ) -> OsResult<Box<dyn MultipartUpload>> {
3645 self.inner.put_multipart_opts(location, opts).await
3646 }
3647
3648 async fn get_ranges(
3649 &self,
3650 location: &ObjPath,
3651 ranges: &[Range<u64>],
3652 ) -> OsResult<Vec<Bytes>> {
3653 if is_index_path(location) {
3654 let mut out = Vec::with_capacity(ranges.len());
3656 for range in ranges {
3657 let opts = GetOptions {
3658 range: Some(range.clone().into()),
3659 ..Default::default()
3660 };
3661 out.push(self.get_opts(location, opts).await?.bytes().await?);
3662 }
3663 return Ok(out);
3664 }
3665 self.inner.get_ranges(location, ranges).await
3666 }
3667
3668 fn delete_stream(
3669 &self,
3670 locations: BoxStream<'static, OsResult<ObjPath>>,
3671 ) -> BoxStream<'static, OsResult<ObjPath>> {
3672 self.inner.delete_stream(locations)
3673 }
3674
3675 fn list(&self, prefix: Option<&ObjPath>) -> BoxStream<'static, OsResult<ObjectMeta>> {
3676 self.inner.list(prefix)
3677 }
3678
3679 fn list_with_offset(
3680 &self,
3681 prefix: Option<&ObjPath>,
3682 offset: &ObjPath,
3683 ) -> BoxStream<'static, OsResult<ObjectMeta>> {
3684 self.inner.list_with_offset(prefix, offset)
3685 }
3686
3687 async fn list_with_delimiter(&self, prefix: Option<&ObjPath>) -> OsResult<ListResult> {
3688 self.inner.list_with_delimiter(prefix).await
3689 }
3690
3691 async fn copy_opts(&self, from: &ObjPath, to: &ObjPath, opts: CopyOptions) -> OsResult<()> {
3692 self.inner.copy_opts(from, to, opts).await
3693 }
3694 }
3695
3696 #[cfg(test)]
3697 mod tests {
3698 #![allow(clippy::unwrap_used)]
3699 use super::*;
3700 use object_store::memory::InMemory;
3701
3702 async fn read(store: &Arc<dyn ObjectStore>, path: &ObjPath) -> Option<Vec<u8>> {
3703 store
3704 .get(path)
3705 .await
3706 .ok()?
3707 .bytes()
3708 .await
3709 .ok()
3710 .map(|b| b.to_vec())
3711 }
3712
3713 #[tokio::test]
3714 async fn caches_index_files_and_passes_data_through() {
3715 let temp = tempfile::tempdir().unwrap();
3716 let inner: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
3717 let index_path = ObjPath::from("d/messages.lance/_indices/uuid1/index.idx");
3718 let data_path = ObjPath::from("d/messages.lance/data/x.lance");
3719 inner
3720 .put(&index_path, PutPayload::from_static(b"INDEX"))
3721 .await
3722 .unwrap();
3723 inner
3724 .put(&data_path, PutPayload::from_static(b"DATA"))
3725 .await
3726 .unwrap();
3727
3728 let cache = IndexDiskCache::new(temp.path().join("indices")).unwrap();
3729 let store = cache.wrap("test", inner.clone());
3730
3731 assert_eq!(
3732 read(&store, &index_path).await.as_deref(),
3733 Some(&b"INDEX"[..])
3734 );
3735 assert_eq!(
3736 read(&store, &data_path).await.as_deref(),
3737 Some(&b"DATA"[..])
3738 );
3739
3740 inner.delete(&index_path).await.unwrap();
3743 inner.delete(&data_path).await.unwrap();
3744 assert_eq!(
3745 read(&store, &index_path).await.as_deref(),
3746 Some(&b"INDEX"[..])
3747 );
3748 assert_eq!(read(&store, &data_path).await, None);
3749
3750 let slice = store.get_range(&index_path, 1..4).await.unwrap();
3752 assert_eq!(slice.as_ref(), b"NDE");
3753 }
3754 }
3755}
3756
3757#[cfg(unix)]
3766pub mod durability {
3767 use std::fs::File;
3768 use std::io::ErrorKind;
3769 use std::ops::Range;
3770 use std::path::Path as FsPath;
3771 use std::sync::Arc;
3772
3773 use bytes::Bytes;
3774 use futures::stream::BoxStream;
3775 use lance_io::object_store::WrappingObjectStore;
3776 use object_store::path::Path as ObjPath;
3777 use object_store::{
3778 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
3779 PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OsResult,
3780 UploadPart,
3781 };
3782
3783 fn durability_error(op: &str, path: &FsPath, source: std::io::Error) -> object_store::Error {
3784 object_store::Error::Generic {
3785 store: "fsync-durability",
3786 source: format!("{op} {}: {source}", path.display()).into(),
3787 }
3788 }
3789
3790 fn sync_file_and_parent(location: &ObjPath) -> OsResult<()> {
3796 let local = lance_io::local::to_local_path(location);
3797 let path = FsPath::new(&local);
3798 let file = File::open(path).map_err(|e| durability_error("open for fsync", path, e))?;
3799 file.sync_all()
3800 .map_err(|e| durability_error("fsync", path, e))?;
3801 if let Some(parent) = path.parent() {
3802 match File::open(parent) {
3804 Ok(dir) => dir
3805 .sync_all()
3806 .map_err(|e| durability_error("fsync dir", parent, e))?,
3807 Err(e) if e.kind() == ErrorKind::NotFound => {}
3808 Err(e) => return Err(durability_error("open dir for fsync", parent, e)),
3809 }
3810 }
3811 Ok(())
3812 }
3813
3814 #[derive(Debug)]
3817 pub struct FsyncOnWrite;
3818
3819 impl WrappingObjectStore for FsyncOnWrite {
3820 fn wrap(&self, _store_prefix: &str, inner: Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
3821 Arc::new(FsyncStore { inner })
3822 }
3823 }
3824
3825 #[derive(Debug)]
3826 struct FsyncStore {
3827 inner: Arc<dyn ObjectStore>,
3828 }
3829
3830 impl std::fmt::Display for FsyncStore {
3831 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3832 write!(f, "FsyncStore({})", self.inner)
3833 }
3834 }
3835
3836 #[async_trait::async_trait]
3837 impl ObjectStore for FsyncStore {
3838 async fn put_opts(
3839 &self,
3840 location: &ObjPath,
3841 payload: PutPayload,
3842 opts: PutOptions,
3843 ) -> OsResult<PutResult> {
3844 let result = self.inner.put_opts(location, payload, opts).await?;
3845 sync_file_and_parent(location)?;
3846 Ok(result)
3847 }
3848
3849 async fn put_multipart_opts(
3850 &self,
3851 location: &ObjPath,
3852 opts: PutMultipartOptions,
3853 ) -> OsResult<Box<dyn MultipartUpload>> {
3854 let upload = self.inner.put_multipart_opts(location, opts).await?;
3855 Ok(Box::new(FsyncUpload {
3856 inner: upload,
3857 location: location.clone(),
3858 }))
3859 }
3860
3861 async fn get_opts(&self, location: &ObjPath, options: GetOptions) -> OsResult<GetResult> {
3862 self.inner.get_opts(location, options).await
3863 }
3864
3865 async fn get_ranges(
3866 &self,
3867 location: &ObjPath,
3868 ranges: &[Range<u64>],
3869 ) -> OsResult<Vec<Bytes>> {
3870 self.inner.get_ranges(location, ranges).await
3871 }
3872
3873 fn delete_stream(
3874 &self,
3875 locations: BoxStream<'static, OsResult<ObjPath>>,
3876 ) -> BoxStream<'static, OsResult<ObjPath>> {
3877 self.inner.delete_stream(locations)
3878 }
3879
3880 fn list(&self, prefix: Option<&ObjPath>) -> BoxStream<'static, OsResult<ObjectMeta>> {
3881 self.inner.list(prefix)
3882 }
3883
3884 fn list_with_offset(
3885 &self,
3886 prefix: Option<&ObjPath>,
3887 offset: &ObjPath,
3888 ) -> BoxStream<'static, OsResult<ObjectMeta>> {
3889 self.inner.list_with_offset(prefix, offset)
3890 }
3891
3892 async fn list_with_delimiter(&self, prefix: Option<&ObjPath>) -> OsResult<ListResult> {
3893 self.inner.list_with_delimiter(prefix).await
3894 }
3895
3896 async fn copy_opts(&self, from: &ObjPath, to: &ObjPath, opts: CopyOptions) -> OsResult<()> {
3897 self.inner.copy_opts(from, to, opts).await?;
3898 sync_file_and_parent(to)?;
3899 Ok(())
3900 }
3901
3902 async fn rename_opts(
3906 &self,
3907 from: &ObjPath,
3908 to: &ObjPath,
3909 opts: RenameOptions,
3910 ) -> OsResult<()> {
3911 self.inner.rename_opts(from, to, opts).await?;
3912 sync_file_and_parent(to)?;
3913 Ok(())
3914 }
3915 }
3916
3917 #[derive(Debug)]
3921 struct FsyncUpload {
3922 inner: Box<dyn MultipartUpload>,
3923 location: ObjPath,
3924 }
3925
3926 #[async_trait::async_trait]
3927 impl MultipartUpload for FsyncUpload {
3928 fn put_part(&mut self, data: PutPayload) -> UploadPart {
3929 self.inner.put_part(data)
3930 }
3931
3932 async fn complete(&mut self) -> OsResult<PutResult> {
3933 let result = self.inner.complete().await?;
3934 sync_file_and_parent(&self.location)?;
3935 Ok(result)
3936 }
3937
3938 async fn abort(&mut self) -> OsResult<()> {
3939 self.inner.abort().await
3940 }
3941 }
3942
3943 #[cfg(test)]
3944 mod tests {
3945 #![allow(clippy::unwrap_used)]
3946 use super::*;
3947 use object_store::ObjectStoreExt;
3948
3949 fn wrapped() -> Arc<dyn ObjectStore> {
3952 let inner: Arc<dyn ObjectStore> = Arc::new(object_store::local::LocalFileSystem::new());
3953 FsyncOnWrite.wrap("test", inner)
3954 }
3955
3956 fn obj_path(root: &FsPath, name: &str) -> ObjPath {
3957 ObjPath::from(root.join(name).to_string_lossy().trim_start_matches('/'))
3959 }
3960
3961 #[tokio::test]
3962 async fn put_through_wrapper_round_trips_and_lands_on_disk() {
3963 let temp = tempfile::tempdir().unwrap();
3964 let store = wrapped();
3965 let path = obj_path(temp.path(), "sub/dir/manifest");
3966 store
3967 .put(&path, PutPayload::from_static(b"DURABLE"))
3968 .await
3969 .unwrap();
3970 let got = store.get(&path).await.unwrap().bytes().await.unwrap();
3972 assert_eq!(got.as_ref(), b"DURABLE");
3973 assert_eq!(
3975 std::fs::read(temp.path().join("sub/dir/manifest")).unwrap(),
3976 b"DURABLE",
3977 );
3978 }
3979
3980 #[tokio::test]
3981 async fn multipart_through_wrapper_completes_and_round_trips() {
3982 let temp = tempfile::tempdir().unwrap();
3983 let store = wrapped();
3984 let path = obj_path(temp.path(), "data/part.lance");
3985 let mut upload = store.put_multipart(&path).await.unwrap();
3986 upload
3987 .put_part(PutPayload::from_static(b"AB"))
3988 .await
3989 .unwrap();
3990 upload
3991 .put_part(PutPayload::from_static(b"CD"))
3992 .await
3993 .unwrap();
3994 upload.complete().await.unwrap();
3995 let got = store.get(&path).await.unwrap().bytes().await.unwrap();
3996 assert_eq!(got.as_ref(), b"ABCD");
3997 }
3998 }
3999}
4000
4001pub fn store_key(location: &Url) -> String {
4006 blake3::hash(location.as_str().as_bytes()).to_hex()[..16].to_owned()
4007}
4008
4009fn prune_stale_uuid_dirs(dir: &std::path::Path, keep: &std::collections::HashSet<String>) {
4013 let Ok(entries) = std::fs::read_dir(dir) else {
4014 return;
4015 };
4016 for entry in entries.flatten() {
4017 let path = entry.path();
4018 if !path.is_dir() {
4019 continue;
4020 }
4021 if entry.file_name() == "_indices" {
4022 let Ok(children) = std::fs::read_dir(&path) else {
4023 continue;
4024 };
4025 for child in children.flatten() {
4026 if child.path().is_dir()
4027 && !keep.contains(child.file_name().to_string_lossy().as_ref())
4028 {
4029 let _ = std::fs::remove_dir_all(child.path());
4030 }
4031 }
4032 } else {
4033 prune_stale_uuid_dirs(&path, keep);
4034 }
4035 }
4036}
4037
4038fn store_wrapper(
4044 location: &Url,
4045 index_cache_dir: Option<&std::path::Path>,
4046) -> Option<Arc<dyn WrappingObjectStore>> {
4047 let mut wrappers: Vec<Arc<dyn WrappingObjectStore>> = Vec::new();
4048 #[cfg(unix)]
4052 if config::is_local(location) {
4053 wrappers.push(Arc::new(durability::FsyncOnWrite));
4054 }
4055 if let Some(dir) = index_cache_dir
4056 && !config::is_local(location)
4057 {
4058 let root = dir.join(store_key(location)).join("indices");
4059 match index_cache::IndexDiskCache::new(root) {
4060 Ok(cache) => wrappers.push(Arc::new(cache)),
4061 Err(error) => tracing::warn!(%error, "index disk cache disabled; reads hit the store"),
4062 }
4063 }
4064 if let Some(tracker) = io_trace::wrapper() {
4065 wrappers.push(tracker);
4066 }
4067 match wrappers.len() {
4068 0 => None,
4069 1 => Some(wrappers.remove(0)),
4070 _ => Some(Arc::new(ChainedWrappingObjectStore::new(wrappers))),
4071 }
4072}
4073
4074async fn open_or_create_via_ns(
4075 nm: &Arc<dyn LanceNamespace>,
4076 nm_ident: &NamespaceIdent,
4077 table_name: &str,
4078 schema: lance::deps::arrow_schema::SchemaRef,
4079 session: &Arc<Session>,
4080 storage_options: &HashMap<String, String>,
4081 wrapper: Option<Arc<dyn WrappingObjectStore>>,
4082) -> Result<Dataset> {
4083 let table_id = nm_ident.as_table_id(table_name);
4084
4085 let request = DescribeTableRequest {
4086 id: Some(table_id.clone()),
4087 ..Default::default()
4088 };
4089 match nm.describe_table(request).await {
4090 Ok(response) => {
4091 let location = response.location.with_context(|| {
4092 format!("namespace returned no location for table {table_name}")
4093 })?;
4094 let builder = apply_open_params(
4095 DatasetBuilder::from_uri(&location).with_session(session.clone()),
4096 &wrapper,
4097 storage_options,
4098 );
4099 let mut dataset = match builder.load().await {
4100 Ok(dataset) => dataset,
4101 Err(load_error) => {
4102 let load_error = anyhow::Error::new(load_error)
4103 .context(format!("failed to open table {table_name}"));
4104 match config::local_path(&uri_to_url(&location)?) {
4109 Some(table_root) => {
4110 heal_local_dataset(
4111 &location,
4112 &table_root,
4113 table_name,
4114 session,
4115 storage_options,
4116 &wrapper,
4117 load_error,
4118 )
4119 .await?
4120 }
4121 None => return Err(load_error),
4122 }
4123 }
4124 };
4125 ensure_current_schema(&mut dataset, schema.as_ref(), table_name).await?;
4126 return Ok(dataset);
4127 }
4128 Err(error) => match &error {
4129 error if is_namespace_error_code(error, ErrorCode::TableNotFound) => {
4130 }
4132 _ => {
4133 return Err(anyhow::Error::from(error))
4134 .with_context(|| format!("failed to describe table {table_name}"));
4135 }
4136 },
4137 }
4138
4139 let mut write_params = sessions::write_params_for_create();
4142 write_params.session = Some(session.clone());
4143 write_params.mode = WriteMode::Create;
4144 if wrapper.is_some() || !storage_options.is_empty() {
4148 write_params.store_params = Some(ObjectStoreParams {
4149 object_store_wrapper: wrapper.clone(),
4150 storage_options_accessor: (!storage_options.is_empty()).then(|| {
4151 Arc::new(StorageOptionsAccessor::with_static_options(
4152 storage_options.clone(),
4153 ))
4154 }),
4155 ..Default::default()
4156 });
4157 }
4158 let reader = sessions::empty_reader(schema)?;
4159 Dataset::write_into_namespace(reader, nm.clone(), table_id, Some(write_params))
4160 .await
4161 .with_context(|| format!("failed to create table {table_name}"))
4162}
4163
4164fn apply_open_params(
4168 builder: DatasetBuilder,
4169 wrapper: &Option<Arc<dyn WrappingObjectStore>>,
4170 storage_options: &HashMap<String, String>,
4171) -> DatasetBuilder {
4172 match wrapper {
4173 Some(wrapper) => builder.with_store_params(ObjectStoreParams {
4174 object_store_wrapper: Some(wrapper.clone()),
4175 storage_options_accessor: (!storage_options.is_empty()).then(|| {
4176 Arc::new(StorageOptionsAccessor::with_static_options(
4177 storage_options.clone(),
4178 ))
4179 }),
4180 ..Default::default()
4181 }),
4182 None if !storage_options.is_empty() => {
4183 builder.with_storage_options(storage_options.clone())
4184 }
4185 None => builder,
4186 }
4187}
4188
4189const VERSIONS_DIR_NAME: &str = "_versions";
4191const HEAL_MAX_PROBES: usize = 32;
4195
4196fn parse_manifest_version(filename: &str) -> Option<u64> {
4202 if filename.starts_with('d') {
4203 return None;
4204 }
4205 let stem = filename.strip_suffix(".manifest")?;
4206 if stem.len() == 20 {
4207 stem.parse::<u64>().ok().map(|inverted| u64::MAX - inverted)
4208 } else {
4209 stem.parse::<u64>().ok()
4210 }
4211}
4212
4213async fn scan_verify_version(
4222 table_uri: &str,
4223 version: u64,
4224 session: &Arc<Session>,
4225 storage_options: &HashMap<String, String>,
4226 wrapper: &Option<Arc<dyn WrappingObjectStore>>,
4227) -> Result<()> {
4228 let builder = apply_open_params(
4229 DatasetBuilder::from_uri(table_uri)
4230 .with_session(session.clone())
4231 .with_version(version),
4232 wrapper,
4233 storage_options,
4234 );
4235 let dataset = builder.load().await?;
4236 let scanner = dataset.scan();
4237 let mut stream = scanner.try_into_stream().await?;
4238 while let Some(batch) = stream.next().await {
4239 batch?;
4240 }
4241 Ok(())
4242}
4243
4244async fn heal_local_dataset(
4251 table_uri: &str,
4252 table_root: &std::path::Path,
4253 table_name: &str,
4254 session: &Arc<Session>,
4255 storage_options: &HashMap<String, String>,
4256 wrapper: &Option<Arc<dyn WrappingObjectStore>>,
4257 load_error: anyhow::Error,
4258) -> Result<Dataset> {
4259 let versions_dir = table_root.join(VERSIONS_DIR_NAME);
4260 let entries = match std::fs::read_dir(&versions_dir) {
4261 Ok(entries) => entries,
4262 Err(_) => {
4263 return Err(enriched_open_error(
4264 table_name,
4265 format!(
4266 "open failed and no {} directory exists at {} - not a crash-damaged manifest",
4267 VERSIONS_DIR_NAME,
4268 versions_dir.display()
4269 ),
4270 load_error,
4271 ));
4272 }
4273 };
4274 let mut manifests: Vec<(u64, PathBuf)> = Vec::new();
4275 for entry in entries {
4276 let entry = entry.with_context(|| format!("listing {}", versions_dir.display()))?;
4277 if let Some(version) = parse_manifest_version(&entry.file_name().to_string_lossy()) {
4278 manifests.push((version, entry.path()));
4279 }
4280 }
4281 manifests.sort_by_key(|(version, _)| std::cmp::Reverse(*version));
4282 let Some((_, newest_path)) = manifests.first().cloned() else {
4283 return Err(enriched_open_error(
4284 table_name,
4285 format!(
4286 "open failed and no manifest files exist under {} - not a crash-damaged manifest",
4287 versions_dir.display()
4288 ),
4289 load_error,
4290 ));
4291 };
4292 let newest_desc = || {
4293 let name = newest_path
4294 .file_name()
4295 .map(|n| n.to_string_lossy().into_owned())
4296 .unwrap_or_default();
4297 let bytes = std::fs::metadata(&newest_path).map(|m| m.len()).ok();
4298 match bytes {
4299 Some(bytes) => format!("newest manifest {name} ({bytes} bytes)"),
4300 None => format!("newest manifest {name}"),
4301 }
4302 };
4303
4304 let mut doomed: Vec<PathBuf> = Vec::new();
4308 let mut landed: Option<u64> = None;
4309 let mut probes = 0usize;
4310 for (version, path) in &manifests {
4311 let len = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
4312 if len < 16 {
4315 doomed.push(path.clone());
4316 continue;
4317 }
4318 if probes >= HEAL_MAX_PROBES {
4319 break;
4320 }
4321 probes += 1;
4322 match scan_verify_version(table_uri, *version, session, storage_options, wrapper).await {
4323 Ok(()) => {
4324 landed = Some(*version);
4325 break;
4326 }
4327 Err(_) => doomed.push(path.clone()),
4328 }
4329 }
4330
4331 let Some(landed_version) = landed else {
4333 return Err(enriched_open_error(
4334 table_name,
4335 format!(
4336 "{} is unreadable (interrupted commit during a hard host stop) and no older version passed a scan-verify probe; nothing was quarantined",
4337 newest_desc()
4338 ),
4339 load_error,
4340 ));
4341 };
4342 if doomed.is_empty() {
4344 return Err(enriched_open_error(
4345 table_name,
4346 format!(
4347 "open failed but the manifest head under {} is readable - not a crash-damaged manifest",
4348 versions_dir.display()
4349 ),
4350 load_error,
4351 ));
4352 }
4353
4354 let mut quarantined: Vec<String> = Vec::new();
4358 for path in &doomed {
4359 let mut corrupt = path.clone().into_os_string();
4360 corrupt.push(".corrupt");
4361 match std::fs::rename(path, &corrupt) {
4362 Ok(()) => {
4363 if let Some(name) = path.file_name() {
4364 quarantined.push(name.to_string_lossy().into_owned());
4365 }
4366 }
4367 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
4368 Err(error) => {
4369 return Err(anyhow::Error::new(error).context(format!(
4370 "table {table_name}: failed to quarantine corrupt manifest {}",
4371 path.display()
4372 )));
4373 }
4374 }
4375 }
4376
4377 let builder = apply_open_params(
4379 DatasetBuilder::from_uri(table_uri).with_session(session.clone()),
4380 wrapper,
4381 storage_options,
4382 );
4383 let dataset = builder.load().await.with_context(|| {
4384 format!(
4385 "table {table_name}: open still failed after quarantining {} corrupt manifest(s); restore from a `pond copy` replica or re-run `pond init`",
4386 quarantined.len()
4387 )
4388 })?;
4389
4390 tracing::warn!(
4393 "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.",
4394 quarantined.len(),
4395 quarantined.join(", "),
4396 );
4397
4398 Ok(dataset)
4399}
4400
4401fn enriched_open_error(
4405 table_name: &str,
4406 finding: String,
4407 load_error: anyhow::Error,
4408) -> anyhow::Error {
4409 load_error.context(format!(
4410 "table {table_name}: {finding}. Restore this store from a `pond copy` replica or re-run `pond init` to re-sync from source histories"
4411 ))
4412}
4413
4414fn is_namespace_error_code(error: &lance::Error, code: ErrorCode) -> bool {
4418 if !matches!(error, lance::Error::Namespace { .. }) {
4419 return false;
4420 }
4421 std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |link| {
4422 link.source()
4423 })
4424 .filter_map(|link| link.downcast_ref::<NamespaceError>())
4425 .any(|inner| inner.code() == code)
4426}
4427
4428fn scanner_with_prefilter(
4429 dataset: &Dataset,
4430 predicate: Option<&Predicate>,
4431) -> Result<lance::dataset::scanner::Scanner> {
4432 let mut scanner = dataset.scan();
4433 scanner.prefilter(true);
4434 if let Some(predicate) = predicate {
4435 let filter = predicate.to_lance();
4436 if !filter.is_empty() {
4437 scanner.filter(&filter)?;
4438 }
4439 }
4440 Ok(scanner)
4441}
4442enum SchemaFit {
4444 Match,
4445 MissingNullable(Vec<lance::deps::arrow_schema::Field>),
4448 UnknownExtra(Vec<String>),
4453}
4454
4455fn classify_schema(
4456 actual: &lance::deps::arrow_schema::Schema,
4457 expected: &lance::deps::arrow_schema::Schema,
4458 table_name: &str,
4459) -> Result<SchemaFit> {
4460 use std::collections::BTreeSet;
4461 let actual_names: BTreeSet<&str> = actual.fields().iter().map(|f| f.name().as_str()).collect();
4462 let expected_names: BTreeSet<&str> = expected
4463 .fields()
4464 .iter()
4465 .map(|f| f.name().as_str())
4466 .collect();
4467 let missing: Vec<_> = expected
4468 .fields()
4469 .iter()
4470 .filter(|f| !actual_names.contains(f.name().as_str()))
4471 .map(|f| f.as_ref().clone())
4472 .collect();
4473 let extra: Vec<String> = actual_names
4474 .difference(&expected_names)
4475 .map(|name| (*name).to_owned())
4476 .collect();
4477 match (missing.is_empty(), extra.is_empty()) {
4478 (true, true) => Ok(SchemaFit::Match),
4479 (false, true) if missing.iter().all(|f| f.is_nullable()) => {
4480 Ok(SchemaFit::MissingNullable(missing))
4481 }
4482 (true, false) => Ok(SchemaFit::UnknownExtra(extra)),
4483 _ => anyhow::bail!(
4484 "table {table_name} has columns {actual_names:?} but this pond build expects \
4485 {expected_names:?}, and the difference is not an additive nullable-column \
4486 change this build can migrate - upgrade pond, or restore the store from a \
4487 `pond copy` snapshot taken by the version that wrote it",
4488 ),
4489 }
4490}
4491
4492async fn ensure_current_schema(
4500 dataset: &mut Dataset,
4501 expected: &lance::deps::arrow_schema::Schema,
4502 table_name: &str,
4503) -> Result<()> {
4504 use lance::deps::arrow_schema::DataType;
4505 const MAX_MIGRATION_ATTEMPTS: usize = 3;
4506 for _ in 0..MAX_MIGRATION_ATTEMPTS {
4507 let actual = lance::deps::arrow_schema::Schema::from(dataset.schema());
4508 match classify_schema(&actual, expected, table_name)? {
4509 SchemaFit::MissingNullable(missing) => {
4510 backfill_missing_columns(dataset, table_name, missing).await?;
4511 continue;
4512 }
4513 SchemaFit::Match => {}
4514 SchemaFit::UnknownExtra(extra) => {
4515 tracing::warn!(
4516 table = table_name,
4517 ?extra,
4518 "store carries columns unknown to this pond build (written by a newer \
4519 version); reads proceed, writes need the newer pond",
4520 );
4521 }
4522 }
4523 for actual_field in actual.fields() {
4528 let Some(expected_field) = expected.field_with_name(actual_field.name()).ok() else {
4529 continue;
4530 };
4531 if let (
4532 DataType::FixedSizeList(_, actual_dim),
4533 DataType::FixedSizeList(_, expected_dim),
4534 ) = (actual_field.data_type(), expected_field.data_type())
4535 && actual_dim != expected_dim
4536 {
4537 tracing::warn!(
4538 table = table_name,
4539 column = actual_field.name(),
4540 actual_dim,
4541 expected_dim,
4542 "embedding dimension differs from config; open proceeds because model swaps are operator-driven",
4543 );
4544 }
4545 }
4546 return Ok(());
4547 }
4548 anyhow::bail!(
4549 "schema migration for table {table_name} did not converge after \
4550 {MAX_MIGRATION_ATTEMPTS} attempts (a concurrent writer kept changing \
4551 the schema); re-run once the other pond process finishes",
4552 )
4553}
4554
4555async fn backfill_missing_columns(
4560 dataset: &mut Dataset,
4561 table_name: &str,
4562 missing: Vec<lance::deps::arrow_schema::Field>,
4563) -> Result<()> {
4564 use lance::dataset::{BatchUDF, NewColumnTransform};
4565 let names: Vec<&str> = missing.iter().map(|f| f.name().as_str()).collect();
4566 let spec = sessions::column_backfill(table_name, &missing)?;
4567 let _ = crate::output::line_err(&format!(
4571 "migrating {table_name}: backfilling {names:?} from stored data (one-time, in place)...",
4572 ));
4573 let started = std::time::Instant::now();
4574 let mapper = spec.mapper;
4575 let migration: std::pin::Pin<
4579 Box<dyn std::future::Future<Output = lance::Result<()>> + Send + '_>,
4580 > = Box::pin(dataset.add_columns(
4581 NewColumnTransform::BatchUDF(BatchUDF {
4582 mapper: Box::new(move |batch| {
4583 mapper(batch).map_err(|error| lance::Error::io(format!("{error:#}")))
4584 }),
4585 output_schema: spec.output_schema,
4586 result_checkpoint: None,
4587 }),
4588 Some(spec.read_columns),
4589 None,
4590 ));
4591 let result = migration.await;
4592 match result {
4593 Ok(()) => {
4594 let _ = crate::output::line_err(&format!(
4595 "migrated {table_name} in {:.1}s",
4596 started.elapsed().as_secs_f64(),
4597 ));
4598 Ok(())
4599 }
4600 Err(error) => {
4601 let error = anyhow::Error::from(error);
4602 if is_commit_conflict(&error) {
4603 dataset.checkout_latest().await?;
4606 Ok(())
4607 } else {
4608 Err(error).with_context(|| {
4609 format!(
4610 "schema backfill failed for {table_name}; the one-time migration \
4611 writes new column files, so it needs write access to the store - \
4612 re-run any pond command with write-capable credentials to complete it",
4613 )
4614 })
4615 }
4616 }
4617 }
4618}
4619fn apply_remote_storage_defaults(options: &mut HashMap<String, String>) {
4626 fn set_default(options: &mut HashMap<String, String>, aliases: &[&str], value: &str) {
4627 if aliases
4628 .iter()
4629 .any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)))
4630 {
4631 return;
4632 }
4633 options.insert(aliases[0].to_owned(), value.to_owned());
4634 }
4635 set_default(options, &["pool_idle_timeout"], "300 seconds");
4636 set_default(options, &["connect_timeout"], "10 seconds");
4637 set_default(options, &["request_timeout"], "60 seconds");
4644 let has_custom_endpoint = ["aws_endpoint", "endpoint"]
4645 .iter()
4646 .any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)));
4647 if has_custom_endpoint {
4648 set_default(
4649 options,
4650 &["aws_unsigned_payload", "unsigned_payload"],
4651 "true",
4652 );
4653 }
4654}
4655
4656fn quoted_string(value: &str) -> String {
4657 format!("'{}'", value.replace('\'', "''"))
4658}
4659fn like_contains(value: &str) -> String {
4660 let escaped = value
4661 .replace('\\', "\\\\")
4662 .replace('%', "\\%")
4663 .replace('_', "\\_")
4664 .replace('\'', "''");
4665 format!("'%{escaped}%'")
4666}
4667
4668#[cfg(test)]
4669mod tests {
4670 #![allow(clippy::expect_used, clippy::unwrap_used)]
4671
4672 use super::*;
4673 use tempfile::TempDir;
4674
4675 #[test]
4676 fn is_index_error_matches_lance_index_class_through_context() {
4677 let index_fault = anyhow::Error::from(lance::Error::index(
4678 "cannot merge inverted index segments with different posting tail codecs",
4679 ))
4680 .context("optimize_indices(merge) failed during index optimize");
4681 assert!(is_index_error(&index_fault));
4682
4683 let io_fault = anyhow::anyhow!("connection reset").context("optimize_indices failed");
4684 assert!(!is_index_error(&io_fault));
4685 }
4686
4687 #[test]
4688 fn prune_keeps_live_uuid_dirs_and_drops_dead_ones() {
4689 let temp = TempDir::new().unwrap();
4690 let indices = temp.path().join("bkt/messages.lance/_indices");
4691 for uuid in ["live", "dead"] {
4692 std::fs::create_dir_all(indices.join(uuid)).unwrap();
4693 std::fs::write(indices.join(uuid).join("index.idx"), b"x").unwrap();
4694 }
4695 let keep = std::collections::HashSet::from(["live".to_owned()]);
4696 prune_stale_uuid_dirs(temp.path(), &keep);
4697 assert!(indices.join("live").exists());
4698 assert!(!indices.join("dead").exists());
4699 }
4700
4701 #[cfg(unix)]
4702 #[test]
4703 fn store_wrapper_present_for_local_absent_for_remote() {
4704 let local = Url::parse("file:///tmp/pond-wrapper-test").unwrap();
4707 assert!(store_wrapper(&local, None).is_some());
4708 for remote in ["memory:///pond-wrapper-test", "s3://bucket/prefix"] {
4709 let url = Url::parse(remote).unwrap();
4710 assert!(
4711 store_wrapper(&url, None).is_none(),
4712 "remote store must not carry the fsync wrapper: {remote}",
4713 );
4714 }
4715 }
4716
4717 fn set(scope: Option<&str>) -> CredsSet {
4718 CredsSet {
4719 scope: scope.map(str::to_owned),
4720 access_key_id: Some("AKIA".to_owned()),
4721 secret_access_key: Some("shh".to_owned()),
4722 ..CredsSet::default()
4723 }
4724 }
4725
4726 fn opts(resolved: &ResolvedStorage, key: &str) -> Option<String> {
4727 resolved.options.get(key).cloned()
4728 }
4729
4730 #[test]
4731 fn storage_url_translation_table() {
4732 let local = StorageUrl::parse("/srv/pond").unwrap();
4735 assert_eq!(local.lance_url().as_str(), "file:///srv/pond/");
4736 assert!(local.is_local());
4737 assert!(local.scheme_options.is_empty());
4738 let aws = StorageUrl::parse("s3://bucket/prefix").unwrap();
4740 assert_eq!(aws.lance_url().as_str(), "s3://bucket/prefix");
4741 assert!(aws.scheme_options.is_empty());
4742 let fat = StorageUrl::parse("s3+https://nbg1.example.com/my-pond/sub").unwrap();
4747 assert_eq!(fat.lance_url().as_str(), "s3://my-pond/sub");
4748 assert_eq!(
4749 fat.scheme_options,
4750 vec![
4751 ("allow_http", "false".to_owned()),
4752 ("virtual_hosted_style_request", "true".to_owned()),
4753 ("region", "us-east-1".to_owned()),
4754 ],
4755 );
4756 let resolved = fat.resolve(&BTreeMap::new()).unwrap();
4757 assert_eq!(
4758 opts(&resolved, "endpoint").as_deref(),
4759 Some("https://my-pond.nbg1.example.com"),
4760 );
4761 assert_eq!(opts(&resolved, "region").as_deref(), Some("us-east-1"));
4762 let plain = StorageUrl::parse("s3+http://127.0.0.1:9000/pond").unwrap();
4765 assert_eq!(plain.lance_url().as_str(), "s3://pond/");
4766 assert_eq!(plain.scheme_options[0], ("allow_http", "true".to_owned()));
4767 assert_eq!(
4768 plain.scheme_options[1],
4769 ("virtual_hosted_style_request", "false".to_owned()),
4770 );
4771 let resolved = plain.resolve(&BTreeMap::new()).unwrap();
4772 assert_eq!(
4773 opts(&resolved, "endpoint").as_deref(),
4774 Some("http://127.0.0.1:9000"),
4775 );
4776 let mut pinned = BTreeMap::new();
4778 pinned.insert(
4779 "default".to_owned(),
4780 CredsSet {
4781 extra: [(
4782 "endpoint".to_owned(),
4783 "https://pinned.example.com".to_owned(),
4784 )]
4785 .into_iter()
4786 .collect(),
4787 ..CredsSet::default()
4788 },
4789 );
4790 let resolved = fat.resolve(&pinned).unwrap();
4791 assert_eq!(
4792 opts(&resolved, "endpoint").as_deref(),
4793 Some("https://pinned.example.com"),
4794 );
4795 let gcs = StorageUrl::parse("gs://bucket/p").unwrap();
4797 assert_eq!(gcs.lance_url().as_str(), "gs://bucket/p");
4798 let azure = StorageUrl::parse("az://acct/container/p").unwrap();
4800 assert_eq!(azure.lance_url().as_str(), "az://container/p");
4801 assert_eq!(
4802 azure.scheme_options,
4803 vec![("account_name", "acct".to_owned())]
4804 );
4805 let shared = StorageUrl::parse("shared-memory://pond-test-x/").unwrap();
4807 assert_eq!(shared.lance_url().as_str(), "shared-memory://pond-test-x/");
4808 }
4809
4810 #[test]
4811 fn storage_url_rejects_bad_shapes() {
4812 let err = StorageUrl::parse("s3+https://user:pass@host/bucket")
4814 .expect_err("userinfo must be rejected")
4815 .to_string();
4816 assert!(
4817 err.contains("creds"),
4818 "error must name the alternative: {err}"
4819 );
4820 assert!(StorageUrl::parse("s3+https://host").is_err());
4822 assert!(StorageUrl::parse("az://acct").is_err());
4823 let err = StorageUrl::parse("ftp://host/x")
4825 .expect_err("ftp")
4826 .to_string();
4827 assert!(err.contains("s3+https"), "got: {err}");
4828 let err = StorageUrl::parse("s3://b/p?regoin=x")
4830 .expect_err("typo")
4831 .to_string();
4832 assert!(err.contains("regoin"), "got: {err}");
4833 let err = StorageUrl::parse("memory://x?creds=y")
4836 .expect_err("memory query")
4837 .to_string();
4838 assert!(err.contains("query params"), "got: {err}");
4839 let err = StorageUrl::parse("file:///x?creds=y")
4840 .expect_err("file query")
4841 .to_string();
4842 assert!(err.contains("query params"), "got: {err}");
4843 assert!(StorageUrl::parse("/tmp/a?b").is_ok());
4845 }
4846
4847 #[test]
4848 fn storage_url_canonicalizes_ports_and_keeps_percent_encoding() {
4849 let with_port = StorageUrl::parse("s3+https://host:443/bucket/p").unwrap();
4851 let without = StorageUrl::parse("s3+https://host/bucket/p").unwrap();
4852 assert_eq!(with_port.canonical(), without.canonical());
4853 let odd = StorageUrl::parse("s3+https://host:8443/bucket").unwrap();
4855 let resolved = odd.resolve(&BTreeMap::new()).unwrap();
4856 assert_eq!(
4857 resolved.options.get("endpoint").map(String::as_str),
4858 Some("https://bucket.host:8443"),
4859 );
4860 let encoded = StorageUrl::parse("s3+https://host/bucket/pre%20fix").unwrap();
4862 assert_eq!(encoded.lance_url().as_str(), "s3://bucket/pre%20fix");
4863 }
4864
4865 #[test]
4866 fn query_params_strip_and_apply_over_set_fields() {
4867 let mut creds = BTreeMap::new();
4868 creds.insert(
4869 "default".to_owned(),
4870 CredsSet {
4871 region: Some("from-set".to_owned()),
4872 virtual_hosted_style_request: Some(false),
4873 ..set(None)
4874 },
4875 );
4876 let url = StorageUrl::parse(
4877 "s3+https://host/bucket/p?region=from-query&virtual_hosted_style_request=true",
4878 )
4879 .unwrap();
4880 assert_eq!(url.lance_url().as_str(), "s3://bucket/p");
4882 assert!(url.canonical().query().is_none());
4883 let resolved = url.resolve(&creds).unwrap();
4884 assert_eq!(opts(&resolved, "region").as_deref(), Some("from-query"));
4886 assert_eq!(
4887 opts(&resolved, "virtual_hosted_style_request").as_deref(),
4888 Some("true"),
4889 );
4890 assert_eq!(
4892 opts(&resolved, "endpoint").as_deref(),
4893 Some("https://bucket.host"),
4894 );
4895 }
4896
4897 #[test]
4898 fn scope_matching_binds_by_longest_prefix_at_segment_boundaries() {
4899 let mut creds = BTreeMap::new();
4900 creds.insert("all".to_owned(), set(None));
4901 creds.insert("bucket".to_owned(), set(Some("s3+https://host/pond/")));
4902 creds.insert("deep".to_owned(), set(Some("s3+https://host/pond/sub")));
4903
4904 let bind = |input: &str| {
4905 StorageUrl::parse(input)
4906 .unwrap()
4907 .resolve(&creds)
4908 .unwrap()
4909 .binding
4910 };
4911 assert_eq!(
4913 bind("s3+https://host/pond/sub/x"),
4914 CredsBinding::Set {
4915 name: "deep".to_owned(),
4916 via: BindVia::Scope
4917 },
4918 );
4919 assert_eq!(
4920 bind("s3+https://host/pond/other"),
4921 CredsBinding::Set {
4922 name: "bucket".to_owned(),
4923 via: BindVia::Scope
4924 },
4925 );
4926 assert_eq!(
4928 bind("s3+https://host/pond-2"),
4929 CredsBinding::Set {
4930 name: "all".to_owned(),
4931 via: BindVia::CatchAll
4932 },
4933 );
4934 assert_eq!(
4936 bind("s3://pond/sub"),
4937 CredsBinding::Set {
4938 name: "all".to_owned(),
4939 via: BindVia::CatchAll
4940 },
4941 );
4942 assert_eq!(
4944 bind("s3+https://host:443/pond/x"),
4945 CredsBinding::Set {
4946 name: "bucket".to_owned(),
4947 via: BindVia::Scope
4948 },
4949 );
4950 assert_eq!(
4952 bind("s3+https://host/pond/sub/x?creds=all"),
4953 CredsBinding::Set {
4954 name: "all".to_owned(),
4955 via: BindVia::Pointer
4956 },
4957 );
4958 let err = StorageUrl::parse("s3://b/p?creds=nope")
4960 .unwrap()
4961 .resolve(&creds)
4962 .expect_err("missing set")
4963 .to_string();
4964 assert!(err.contains("creds=nope"), "got: {err}");
4965
4966 let empty = BTreeMap::new();
4968 assert_eq!(
4969 StorageUrl::parse("s3://b/p")
4970 .unwrap()
4971 .resolve(&empty)
4972 .unwrap()
4973 .binding,
4974 CredsBinding::Ambient,
4975 );
4976 assert_eq!(
4977 StorageUrl::parse("/srv/pond")
4978 .unwrap()
4979 .resolve(&creds)
4980 .unwrap()
4981 .binding,
4982 CredsBinding::NotApplicable,
4983 );
4984 }
4985
4986 #[test]
4987 fn unmatched_sets_are_reported_only_on_remote_invocations() {
4988 let mut creds = BTreeMap::new();
4989 creds.insert("used".to_owned(), set(Some("s3://bucket/")));
4990 creds.insert("idle".to_owned(), set(Some("s3://other/")));
4991
4992 let remote = StorageUrl::parse("s3://bucket/p")
4993 .unwrap()
4994 .resolve(&creds)
4995 .unwrap();
4996 assert_eq!(unmatched_creds_sets(&[&remote], &creds), vec!["idle"]);
4997
4998 let local = StorageUrl::parse("/srv/pond")
5000 .unwrap()
5001 .resolve(&creds)
5002 .unwrap();
5003 assert!(unmatched_creds_sets(&[&local], &creds).is_empty());
5004 }
5005
5006 #[test]
5007 fn secrets_materialize_from_file_and_command() {
5008 let dir = TempDir::new().unwrap();
5009 let key_path = dir.path().join("key");
5010 std::fs::write(&key_path, "from-file\n").unwrap();
5011 let mut creds = BTreeMap::new();
5012 creds.insert(
5013 "default".to_owned(),
5014 CredsSet {
5015 access_key_id_file: Some(key_path),
5016 secret_access_key_command: Some("printf 'from-command\\n\\n'".to_owned()),
5018 ..CredsSet::default()
5019 },
5020 );
5021 let url = StorageUrl::parse("s3://bucket/p").unwrap();
5022 let resolved = url.resolve(&creds).unwrap();
5023 assert_eq!(
5024 opts(&resolved, "access_key_id").as_deref(),
5025 Some("from-file")
5026 );
5027 assert_eq!(
5028 opts(&resolved, "secret_access_key").as_deref(),
5029 Some("from-command\n"),
5030 );
5031
5032 let mut failing = BTreeMap::new();
5034 failing.insert(
5035 "default".to_owned(),
5036 CredsSet {
5037 secret_access_key_command: Some("exit 3".to_owned()),
5038 ..CredsSet::default()
5039 },
5040 );
5041 let err = url
5042 .resolve(&failing)
5043 .expect_err("command must fail")
5044 .to_string();
5045 assert!(err.contains("exit 3"), "got: {err}");
5046
5047 let marker = dir.path().join("runs");
5049 let command = format!("echo run >> {} && echo secret", marker.display());
5050 let mut counted = BTreeMap::new();
5051 counted.insert(
5052 "default".to_owned(),
5053 CredsSet {
5054 secret_access_key_command: Some(command),
5055 ..CredsSet::default()
5056 },
5057 );
5058 url.resolve(&counted).unwrap();
5059 url.resolve(&counted).unwrap();
5060 let runs = std::fs::read_to_string(&marker).unwrap();
5061 assert_eq!(runs.lines().count(), 1, "command must run exactly once");
5062 }
5063
5064 #[test]
5065 fn check_errors_classify_by_kind_and_binding() {
5066 let auth_error = || object_store::Error::Unauthenticated {
5067 path: "k".to_owned(),
5068 source: "denied".into(),
5069 };
5070 let bound = CredsBinding::Set {
5071 name: "work".to_owned(),
5072 via: BindVia::Scope,
5073 };
5074 match classify_check_error(auth_error(), &bound, "put") {
5076 CheckFailure::Auth { set, .. } => assert_eq!(set, "work"),
5077 other => panic!("want Auth, got {other:?}"),
5078 }
5079 assert!(matches!(
5081 classify_check_error(auth_error(), &CredsBinding::Ambient, "put"),
5082 CheckFailure::NoCreds { .. },
5083 ));
5084 let denied = object_store::Error::PermissionDenied {
5085 path: "k".to_owned(),
5086 source: "403".into(),
5087 };
5088 assert!(matches!(
5089 classify_check_error(denied, &bound, "put"),
5090 CheckFailure::Auth { .. },
5091 ));
5092 let missing = object_store::Error::NotFound {
5094 path: "k".to_owned(),
5095 source: "404".into(),
5096 };
5097 assert!(matches!(
5098 classify_check_error(missing, &bound, "get"),
5099 CheckFailure::Io { .. },
5100 ));
5101 let no_creds = || object_store::Error::Generic {
5105 store: "S3",
5106 source: "Failed to get AWS credentials: CredentialsNotLoaded".into(),
5107 };
5108 assert!(matches!(
5109 classify_check_error(no_creds(), &bound, "put"),
5110 CheckFailure::Auth { .. },
5111 ));
5112 assert!(matches!(
5113 classify_check_error(no_creds(), &CredsBinding::Ambient, "put"),
5114 CheckFailure::NoCreds { .. },
5115 ));
5116 }
5117
5118 #[test]
5119 fn concise_cause_strips_upstream_noise_to_one_line() {
5120 let inner = "Encountered internal error. Please file a bug report at \
5123 https://github.com/lance-format/lance/issues. Failed to get AWS \
5124 credentials: CredentialsNotLoaded, <WORKSPACE>/src/object_store/providers/aws.rs:401:21: \
5125 Encountered internal error. Please file a bug report at \
5126 https://github.com/lance-format/lance/issues. Failed to get AWS \
5127 credentials: CredentialsNotLoaded";
5128 let failure = CheckFailure::NoCreds {
5129 source: anyhow!(inner.to_owned()).context("initial conditional put"),
5130 };
5131 let cause = failure.concise_cause().expect("auth-class carries a cause");
5132 assert_eq!(cause, "Failed to get AWS credentials: CredentialsNotLoaded");
5133 assert!(
5135 !failure.to_string().contains("file a bug report"),
5136 "lead must not trail the chain: {failure}"
5137 );
5138 let occ = CheckFailure::OccUnsupported {
5140 detail: "put-if-none-match ignored".to_owned(),
5141 };
5142 assert!(occ.concise_cause().is_none());
5143 let long = CheckFailure::Io {
5146 source: anyhow!(format!("{} dns error: lookup failed", "x".repeat(500))),
5147 };
5148 let cause = long.concise_cause().expect("io carries a cause");
5149 assert!(cause.contains(" ... "), "long causes truncate: {cause}");
5150 assert!(
5151 cause.ends_with("dns error: lookup failed"),
5152 "the tail survives: {cause}"
5153 );
5154 }
5155
5156 #[tokio::test]
5157 async fn storage_check_passes_on_memory_backend() {
5158 let resolved = StorageUrl::parse("memory://check/probe")
5159 .unwrap()
5160 .resolve(&BTreeMap::new())
5161 .unwrap();
5162 storage_check(&resolved).await.expect("memory probe passes");
5163 }
5164
5165 fn stat(bytes: u64) -> FragmentStat {
5166 FragmentStat {
5167 bytes: Some(bytes),
5168 rows: bytes / 1_000,
5169 deleted_rows: 0,
5170 }
5171 }
5172
5173 #[test]
5174 fn compaction_veto_blocks_absorb_keeps_peers() {
5175 let absorb = [stat(665_000_000), stat(1_000_000), stat(2_000_000)];
5177 assert!(!keep_task(&absorb, 64, 0.1));
5178 let peers = [stat(300_000_000), stat(300_000_000)];
5180 assert!(keep_task(&peers, 64, 0.1));
5181 let tiered = [stat(400_000), stat(60_000), stat(40_000)];
5183 assert!(keep_task(&tiered, 64, 0.1));
5184 }
5185
5186 #[test]
5187 fn compaction_veto_passes_deletions_and_cap() {
5188 let mut deleting = stat(665_000_000);
5189 deleting.deleted_rows = deleting.rows / 5;
5190 assert!(keep_task(&[deleting, stat(1_000)], 64, 0.1));
5191
5192 let wide: Vec<FragmentStat> = std::iter::once(stat(665_000_000))
5193 .chain(std::iter::repeat_with(|| stat(1_000)).take(63))
5194 .collect();
5195 assert!(keep_task(&wide, 64, 0.1));
5196 }
5197
5198 #[test]
5199 fn compaction_veto_falls_back_to_rows_on_unknown_sizes() {
5200 let mut unknown = stat(665_000_000);
5201 unknown.bytes = None;
5202 assert!(!keep_task(
5204 &[unknown, stat(1_000_000), stat(2_000_000)],
5205 64,
5206 0.1
5207 ));
5208 }
5209
5210 #[test]
5211 fn parse_manifest_version_handles_all_naming_schemes() {
5212 assert_eq!(parse_manifest_version("5.manifest"), Some(5));
5214 assert_eq!(parse_manifest_version("0.manifest"), Some(0));
5215 assert_eq!(
5217 parse_manifest_version("18446744073709551615.manifest"),
5218 Some(0)
5219 );
5220 assert_eq!(
5221 parse_manifest_version("18446744073709551610.manifest"),
5222 Some(5)
5223 );
5224 assert_eq!(parse_manifest_version("d123.manifest"), None);
5226 assert_eq!(parse_manifest_version("5.manifest.corrupt"), None);
5228 assert_eq!(
5229 parse_manifest_version(".tmp_7.manifest_9c100374-3298-4537-afc6-f5ee7913666d"),
5230 None
5231 );
5232 assert_eq!(parse_manifest_version("data.lance"), None);
5234 assert_eq!(parse_manifest_version("notanumber.manifest"), None);
5235 }
5236
5237 #[tokio::test]
5238 async fn scan_verify_rejects_zeroed_column_add_data_file() {
5239 let temp = tempfile::tempdir().unwrap();
5243 let uri_owned = temp.path().join("t.lance");
5244 let uri = uri_owned.to_str().unwrap();
5245 let schema = Arc::new(lance::deps::arrow_schema::Schema::new(vec![
5246 lance::deps::arrow_schema::Field::new(
5247 "id",
5248 lance::deps::arrow_schema::DataType::Utf8,
5249 false,
5250 ),
5251 ]));
5252 let batch = RecordBatch::try_new(
5253 schema.clone(),
5254 vec![Arc::new(StringArray::from(vec!["a", "b", "c"]))],
5255 )
5256 .unwrap();
5257 let reader = RecordBatchIterator::new([Ok(batch)], schema);
5258 let mut dataset = Dataset::write(reader, uri, None).await.unwrap();
5259
5260 let data_files = || -> std::collections::BTreeSet<PathBuf> {
5261 std::fs::read_dir(uri_owned.join("data"))
5262 .unwrap()
5263 .map(|entry| entry.unwrap().path())
5264 .collect()
5265 };
5266 let before = data_files();
5267 dataset
5268 .add_columns(
5269 lance::dataset::NewColumnTransform::SqlExpressions(vec![(
5270 "extra".to_string(),
5271 "id".to_string(),
5272 )]),
5273 None,
5274 None,
5275 )
5276 .await
5277 .unwrap();
5278 let column_add_file = data_files()
5279 .difference(&before)
5280 .next()
5281 .cloned()
5282 .expect("add_columns writes a new per-fragment data file");
5283 let version = dataset.version().version;
5284 drop(dataset);
5285
5286 let fresh = || Arc::new(Session::new(0, 0, Arc::new(ObjectStoreRegistry::default())));
5288 scan_verify_version(uri, version, &fresh(), &HashMap::new(), &None)
5289 .await
5290 .expect("intact version must pass scan-verify");
5291 std::fs::write(&column_add_file, b"").unwrap();
5292 let verdict = scan_verify_version(uri, version, &fresh(), &HashMap::new(), &None).await;
5293 assert!(
5294 verdict.is_err(),
5295 "zeroed column-add data file must fail scan-verify",
5296 );
5297 }
5298
5299 #[test]
5300 fn cleanup_due_gates_on_version_interval() {
5301 assert!(cleanup_due(0, 1));
5303 assert!(cleanup_due(7, 1));
5304 assert!(cleanup_due(5, 0));
5305 assert!(cleanup_due(0, 16));
5307 assert!(cleanup_due(16, 16));
5308 assert!(cleanup_due(48, 16));
5309 assert!(!cleanup_due(15, 16));
5310 assert!(!cleanup_due(17, 16));
5311 assert!(!cleanup_due(31, 16));
5312 }
5313
5314 #[test]
5315 fn derived_target_rows_tracks_row_size_and_clamps() {
5316 let parts_like = [FragmentStat {
5319 bytes: Some(665_000_000),
5320 rows: 511_000,
5321 deleted_rows: 0,
5322 }];
5323 let target = derived_target_rows(&parts_like);
5324 assert!((80_000..150_000).contains(&target), "{target}");
5325 let unknown = [FragmentStat {
5327 bytes: None,
5328 rows: 511_000,
5329 deleted_rows: 0,
5330 }];
5331 assert_eq!(
5332 derived_target_rows(&unknown),
5333 MAX_TARGET_ROWS_PER_FRAGMENT as usize
5334 );
5335 let tiny = [FragmentStat {
5337 bytes: Some(1_000_000),
5338 rows: 100_000,
5339 deleted_rows: 0,
5340 }];
5341 assert_eq!(
5342 derived_target_rows(&tiny),
5343 MAX_TARGET_ROWS_PER_FRAGMENT as usize
5344 );
5345 let huge = [FragmentStat {
5346 bytes: Some(1_000_000_000),
5347 rows: 100,
5348 deleted_rows: 0,
5349 }];
5350 assert_eq!(
5351 derived_target_rows(&huge),
5352 MIN_TARGET_ROWS_PER_FRAGMENT as usize
5353 );
5354 }
5355
5356 #[test]
5357 fn namespace_error_code_walks_wrapped_chain() {
5358 let direct = lance::Error::namespace_source(Box::new(NamespaceError::TableNotFound {
5359 message: "missing".into(),
5360 }));
5361 assert!(is_namespace_error_code(&direct, ErrorCode::TableNotFound));
5362
5363 let wrapped = lance::Error::namespace_source(Box::new(direct));
5364 assert!(is_namespace_error_code(&wrapped, ErrorCode::TableNotFound));
5365
5366 let other_code =
5367 lance::Error::namespace_source(Box::new(NamespaceError::NamespaceNotFound {
5368 message: "nope".into(),
5369 }));
5370 assert!(!is_namespace_error_code(
5371 &other_code,
5372 ErrorCode::TableNotFound
5373 ));
5374
5375 let not_namespace = lance::Error::internal("unrelated");
5376 assert!(!is_namespace_error_code(
5377 ¬_namespace,
5378 ErrorCode::TableNotFound
5379 ));
5380 }
5381
5382 #[tokio::test]
5386 async fn store_opens_via_namespace_and_scan_works() -> Result<()> {
5387 let temp = TempDir::new()?;
5388 let url = Url::from_directory_path(temp.path())
5389 .map_err(|()| anyhow::anyhow!("temp path is not absolute"))?;
5390 let handle = Handle::open(&url).await?;
5391 let cases: [(Table, &[&str]); 3] = [
5394 (Table::Sessions, &["id"]),
5395 (Table::Messages, &["id"]),
5396 (Table::Parts, &["id"]),
5397 ];
5398 for (table, projection) in cases {
5399 let scanner = handle
5400 .scan(table, ScanOpts::project_only(projection))
5401 .await?;
5402 let batch = scanner.try_into_batch().await?;
5403 assert_eq!(batch.num_rows(), 0, "fresh table should be empty");
5404 }
5405 Ok(())
5406 }
5407}