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 MAX_TARGET_ROWS_PER_FRAGMENT: u64 = 1024 * 1024;
848
849pub const COMPACTION_ABSORB_FACTOR: u64 = 4;
852
853pub fn default_cleanup_older_than() -> chrono::Duration {
859 chrono::Duration::hours(1)
863}
864
865pub const DEFAULT_SYNC_CLEANUP_INTERVAL: u64 = 16;
872
873pub const DEFAULT_SYNC_SCALAR_FOLD_ROWS: usize = 50_000;
882
883pub const DEFAULT_SYNC_INDEX_FOLD_ROWS: usize = 5_000;
893
894#[derive(Debug, Clone, Copy)]
899pub struct MaintenancePolicy {
900 pub compaction_fragment_cap: usize,
902 pub cleanup_older_than: chrono::Duration,
904 pub cleanup_interval: u64,
909 pub scalar_fold_row_threshold: usize,
914 pub index_fold_row_threshold: usize,
920}
921
922impl MaintenancePolicy {
923 pub fn always_compact() -> Self {
925 Self {
926 compaction_fragment_cap: 0,
927 cleanup_older_than: default_cleanup_older_than(),
928 cleanup_interval: 1,
929 scalar_fold_row_threshold: 0,
930 index_fold_row_threshold: 0,
931 }
932 }
933
934 #[must_use]
937 pub fn with_cleanup_interval(mut self, interval: u64) -> Self {
938 self.cleanup_interval = interval.max(1);
939 self
940 }
941
942 #[must_use]
946 pub fn with_scalar_fold_row_threshold(mut self, threshold: usize) -> Self {
947 self.scalar_fold_row_threshold = threshold;
948 self
949 }
950
951 #[must_use]
955 pub fn with_index_fold_row_threshold(mut self, threshold: usize) -> Self {
956 self.index_fold_row_threshold = threshold;
957 self
958 }
959
960 fn fold_thresholds(&self) -> FoldThresholds {
963 FoldThresholds {
964 scalar: self.scalar_fold_row_threshold,
965 index: self.index_fold_row_threshold,
966 }
967 }
968}
969
970#[derive(Debug, Clone, Copy)]
973struct FoldThresholds {
974 scalar: usize,
975 index: usize,
976}
977
978struct FragmentStat {
979 bytes: Option<u64>,
981 rows: u64,
982 deleted_rows: u64,
983}
984
985fn fragment_bytes(fragment: &lance::table::format::Fragment) -> Option<u64> {
988 fragment.files.iter().try_fold(0u64, |total, file| {
989 Some(total + file.file_size_bytes.get()?.get())
990 })
991}
992
993fn fragment_stat(fragment: &lance::table::format::Fragment) -> FragmentStat {
994 FragmentStat {
995 bytes: fragment_bytes(fragment),
996 rows: fragment.physical_rows.unwrap_or(0) as u64,
997 deleted_rows: fragment
998 .deletion_file
999 .as_ref()
1000 .and_then(|deletions| deletions.num_deleted_rows)
1001 .unwrap_or(0) as u64,
1002 }
1003}
1004
1005fn derived_target_rows(stats: &[FragmentStat]) -> usize {
1016 let (mut bytes, mut rows) = (0u64, 0u64);
1017 for stat in stats {
1018 if let Some(fragment_bytes) = stat.bytes
1019 && stat.rows > 0
1020 {
1021 bytes += fragment_bytes;
1022 rows += stat.rows;
1023 }
1024 }
1025 if bytes == 0 || rows == 0 {
1026 return MAX_TARGET_ROWS_PER_FRAGMENT as usize;
1027 }
1028 ((u128::from(TARGET_FRAGMENT_BYTES / 2) * u128::from(rows) / u128::from(bytes))
1029 .clamp(1, u128::from(MAX_TARGET_ROWS_PER_FRAGMENT))) as usize
1030}
1031
1032fn task_veto_reason(
1035 stats: &[FragmentStat],
1036 cap: usize,
1037 deletion_threshold: f32,
1038 target_rows_per_fragment: usize,
1039 max_bytes_per_file: u64,
1040) -> Option<&'static str> {
1041 if stats.iter().any(|stat| {
1042 stat.rows > 0 && (stat.deleted_rows as f32 / stat.rows as f32) > deletion_threshold
1043 }) {
1044 return None;
1045 }
1046
1047 let budget = u128::from(max_bytes_per_file);
1048 if budget == 0 {
1049 return Some("invalid_byte_budget");
1050 }
1051
1052 let (mut total_bytes, mut largest) = (0u128, 0u128);
1053 for stat in stats {
1054 let Some(bytes) = stat.bytes.map(u128::from) else {
1055 return Some("missing_sizes");
1056 };
1057 total_bytes += bytes;
1058 largest = largest.max(bytes);
1059 }
1060
1061 let minimum_outputs = total_bytes.div_ceil(budget).max(1);
1062 if u128::try_from(stats.len()).unwrap_or(u128::MAX) <= minimum_outputs {
1063 return Some("cannot_shrink");
1064 }
1065 if stats.len() >= cap {
1066 return None;
1067 }
1068
1069 if total_bytes > budget {
1070 for stat in stats {
1071 let bytes = u128::from(stat.bytes.unwrap_or(0));
1072 let rows = u128::from(stat.rows);
1073 if rows == 0 || bytes * target_rows_per_fragment as u128 * 2 > rows * budget {
1074 return Some("row_target_unattainable");
1075 }
1076 }
1077 }
1078
1079 if (total_bytes - largest) * u128::from(COMPACTION_ABSORB_FACTOR) < largest {
1080 return Some("absorb_veto");
1081 }
1082 None
1083}
1084
1085#[derive(Debug, Clone)]
1088pub struct IndexIntent {
1089 pub name: &'static str,
1092 pub column: &'static str,
1094 pub trigger: IndexTrigger,
1096 pub params: IndexParamsKind,
1099}
1100
1101#[derive(Debug, Clone)]
1103pub enum IndexTrigger {
1104 OnAnyRows,
1107 OnNonNullCount {
1110 column: &'static str,
1111 threshold: usize,
1112 },
1113}
1114
1115#[derive(Debug, Clone)]
1118pub enum IndexParamsKind {
1119 Scalar(BuiltinIndexType),
1122 InvertedFtsWord,
1128 IvfSqCosine { num_bits: u16, max_iters: usize },
1137}
1138
1139impl IndexTrigger {
1140 async fn should_create(&self, dataset: &Dataset) -> Result<bool> {
1141 match self {
1142 Self::OnAnyRows => Ok(dataset.count_rows(None).await? > 0),
1143 Self::OnNonNullCount { column, threshold } => {
1144 let count = dataset
1145 .count_rows(Some(format!("{column} IS NOT NULL")))
1146 .await?;
1147 Ok(count >= *threshold)
1148 }
1149 }
1150 }
1151}
1152
1153impl IndexParamsKind {
1154 fn index_type(&self) -> IndexType {
1155 match self {
1156 Self::Scalar(BuiltinIndexType::Bitmap) => IndexType::Bitmap,
1157 Self::Scalar(BuiltinIndexType::ZoneMap) => IndexType::ZoneMap,
1158 Self::Scalar(_) => IndexType::BTree,
1159 Self::InvertedFtsWord => IndexType::Inverted,
1160 Self::IvfSqCosine { .. } => IndexType::Vector,
1161 }
1162 }
1163
1164 async fn build(&self, dataset: &Dataset) -> Result<Box<dyn lance::index::IndexParams>> {
1165 match self {
1166 Self::Scalar(kind) => Ok(Box::new(ScalarIndexParams::for_builtin(kind.clone()))),
1167 Self::InvertedFtsWord => Ok(Box::new(
1168 InvertedIndexParams::default()
1169 .base_tokenizer("simple".to_owned())
1170 .stem(true)
1171 .remove_stop_words(false),
1172 )),
1173 Self::IvfSqCosine {
1174 num_bits,
1175 max_iters,
1176 } => {
1177 let count = dataset
1178 .count_rows(Some("vector IS NOT NULL".to_owned()))
1179 .await?;
1180 let partitions = count.checked_div(4096).unwrap_or(0).max(1);
1181 let mut ivf = IvfBuildParams::new(partitions);
1182 ivf.max_iters = *max_iters;
1183 let sq = SQBuildParams {
1184 num_bits: *num_bits,
1185 ..Default::default()
1186 };
1187 Ok(Box::new(VectorIndexParams::with_ivf_sq_params(
1188 MetricType::Cosine,
1189 ivf,
1190 sq,
1191 )))
1192 }
1193 }
1194 }
1195}
1196
1197#[derive(Debug, Clone, PartialEq, Eq)]
1198pub struct IndexStatus {
1199 pub table: Table,
1200 pub intent_name: String,
1201 pub fragments_covered: usize,
1202 pub unindexed_fragments: usize,
1203 pub unindexed_rows: usize,
1204 pub exists: bool,
1205}
1206
1207#[derive(Debug, Clone, Copy)]
1212pub struct ConflictExhausted {
1213 pub attempts: u8,
1214}
1215
1216impl std::fmt::Display for ConflictExhausted {
1217 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1218 write!(
1219 formatter,
1220 "commit conflict exhausted after {} attempt(s)",
1221 self.attempts
1222 )
1223 }
1224}
1225
1226impl std::error::Error for ConflictExhausted {}
1227
1228#[derive(Debug)]
1233pub enum PhaseOutcome {
1234 Ok,
1236 Noop,
1238 SkippedConflict,
1241 Failed(anyhow::Error),
1243 NotAttempted,
1246}
1247
1248impl PhaseOutcome {
1249 pub fn is_failed(&self) -> bool {
1250 matches!(self, Self::Failed(_))
1251 }
1252}
1253
1254#[derive(Debug)]
1256pub struct TableOptimizeOutcome {
1257 pub table: Table,
1258 pub indices: PhaseOutcome,
1259 pub compaction: PhaseOutcome,
1260}
1261
1262#[derive(Debug, Clone)]
1265pub enum OptimizeEvent {
1266 PhaseStart {
1267 table: Table,
1268 phase: OptimizePhase,
1269 detail: Option<String>,
1270 },
1271 PhaseDone {
1272 table: Table,
1273 phase: OptimizePhase,
1274 elapsed_ms: u64,
1275 },
1276 IndexStage {
1281 table: Table,
1282 index: String,
1283 stage: String,
1284 completed: u64,
1285 total: Option<u64>,
1286 unit: String,
1287 },
1288}
1289
1290#[derive(Debug, Clone, Copy)]
1291pub enum OptimizePhase {
1292 Compact,
1293 Cleanup,
1294 IndexCreate,
1295 IndexRebuild,
1296 IndexAppend,
1297}
1298
1299impl OptimizePhase {
1300 pub fn label(self) -> &'static str {
1301 match self {
1302 Self::Compact => "compact",
1303 Self::Cleanup => "cleanup",
1304 Self::IndexCreate => "index-create",
1305 Self::IndexRebuild => "index-rebuild",
1306 Self::IndexAppend => "index-append",
1307 }
1308 }
1309}
1310
1311pub type OptimizeProgressFn = Arc<dyn Fn(OptimizeEvent) + Send + Sync>;
1315
1316fn emit(progress: Option<&OptimizeProgressFn>, event: OptimizeEvent) {
1317 if let Some(callback) = progress {
1318 callback(event);
1319 }
1320}
1321
1322struct PondIndexProgress {
1331 callback: OptimizeProgressFn,
1332 table: Table,
1333 index: String,
1334 state: std::sync::Mutex<PondIndexStageState>,
1335}
1336
1337impl std::fmt::Debug for PondIndexProgress {
1340 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1341 f.debug_struct("PondIndexProgress")
1342 .field("table", &self.table)
1343 .field("index", &self.index)
1344 .finish_non_exhaustive()
1345 }
1346}
1347
1348#[derive(Debug, Default)]
1349struct PondIndexStageState {
1350 total: Option<u64>,
1351 unit: String,
1352 last_emit: Option<Instant>,
1353}
1354
1355impl PondIndexProgress {
1356 fn new(callback: OptimizeProgressFn, table: Table, index: String) -> Arc<Self> {
1357 Arc::new(Self {
1358 callback,
1359 table,
1360 index,
1361 state: std::sync::Mutex::new(PondIndexStageState::default()),
1362 })
1363 }
1364}
1365
1366#[async_trait::async_trait]
1367impl lance_index::progress::IndexBuildProgress for PondIndexProgress {
1368 async fn stage_start(&self, stage: &str, total: Option<u64>, unit: &str) -> lance::Result<()> {
1369 if let Ok(mut state) = self.state.lock() {
1370 state.total = total;
1371 state.unit = unit.to_owned();
1372 state.last_emit = Some(Instant::now());
1373 }
1374 (self.callback)(OptimizeEvent::IndexStage {
1375 table: self.table,
1376 index: self.index.clone(),
1377 stage: stage.to_owned(),
1378 completed: 0,
1379 total,
1380 unit: unit.to_owned(),
1381 });
1382 Ok(())
1383 }
1384
1385 async fn stage_progress(&self, stage: &str, completed: u64) -> lance::Result<()> {
1386 let (total, unit) = {
1387 let Ok(mut state) = self.state.lock() else {
1388 return Ok(());
1389 };
1390 let now = Instant::now();
1391 if let Some(prev) = state.last_emit
1392 && now.duration_since(prev) < Duration::from_millis(100)
1393 {
1394 return Ok(());
1395 }
1396 state.last_emit = Some(now);
1397 (state.total, state.unit.clone())
1398 };
1399 (self.callback)(OptimizeEvent::IndexStage {
1400 table: self.table,
1401 index: self.index.clone(),
1402 stage: stage.to_owned(),
1403 completed,
1404 total,
1405 unit,
1406 });
1407 Ok(())
1408 }
1409
1410 async fn stage_complete(&self, stage: &str) -> lance::Result<()> {
1411 let (total, unit) = {
1412 let Ok(state) = self.state.lock() else {
1413 return Ok(());
1414 };
1415 (state.total, state.unit.clone())
1416 };
1417 (self.callback)(OptimizeEvent::IndexStage {
1418 table: self.table,
1419 index: self.index.clone(),
1420 stage: stage.to_owned(),
1421 completed: total.unwrap_or(0),
1422 total,
1423 unit,
1424 });
1425 Ok(())
1426 }
1427}
1428
1429fn lance_progress(
1430 progress: Option<&OptimizeProgressFn>,
1431 table: Table,
1432 index: &str,
1433) -> Arc<dyn lance_index::progress::IndexBuildProgress> {
1434 match progress {
1435 Some(callback) => PondIndexProgress::new(callback.clone(), table, index.to_owned()),
1436 None => Arc::new(lance_index::progress::NoopIndexBuildProgress),
1437 }
1438}
1439
1440pub fn is_commit_conflict(error: &anyhow::Error) -> bool {
1444 error.downcast_ref::<lance::Error>().is_some_and(|err| {
1445 matches!(
1446 err,
1447 lance::Error::CommitConflict { .. }
1448 | lance::Error::RetryableCommitConflict { .. }
1449 | lance::Error::TooMuchWriteContention { .. }
1450 )
1451 })
1452}
1453
1454fn is_conflict_exhausted(error: &anyhow::Error) -> bool {
1457 error.chain().any(|cause| cause.is::<ConflictExhausted>())
1458}
1459
1460pub fn is_index_error(error: &anyhow::Error) -> bool {
1464 error
1465 .downcast_ref::<lance::Error>()
1466 .is_some_and(|err| matches!(err, lance::Error::Index { .. }))
1467}
1468
1469#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1473pub struct TableSizes {
1474 pub sessions: u64,
1475 pub messages: u64,
1476 pub parts: u64,
1477 pub other: u64,
1478 pub sessions_data: DataLiveness,
1479 pub messages_data: DataLiveness,
1480 pub parts_data: DataLiveness,
1481}
1482
1483#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1486pub struct DataLiveness {
1487 pub on_disk: u64,
1488 pub live: Option<u64>,
1490}
1491
1492impl DataLiveness {
1493 pub fn dead(&self) -> Option<u64> {
1494 self.live.map(|live| self.on_disk.saturating_sub(live))
1495 }
1496}
1497
1498#[derive(Debug, Clone, PartialEq, Eq)]
1499pub enum ScalarValue {
1500 String(String),
1501 Int32(i32),
1502 Raw(String),
1503}
1504impl From<&str> for ScalarValue {
1505 fn from(value: &str) -> Self {
1506 Self::String(value.to_owned())
1507 }
1508}
1509impl From<String> for ScalarValue {
1510 fn from(value: String) -> Self {
1511 Self::String(value)
1512 }
1513}
1514impl From<i32> for ScalarValue {
1515 fn from(value: i32) -> Self {
1516 Self::Int32(value)
1517 }
1518}
1519#[derive(Debug, Clone, PartialEq, Eq)]
1520pub enum Predicate {
1521 Eq(&'static str, ScalarValue),
1522 Ne(&'static str, ScalarValue),
1523 IsNull(&'static str),
1524 IsNotNull(&'static str),
1525 In(&'static str, Vec<ScalarValue>),
1526 LikeContains(&'static str, String),
1527 Regex(&'static str, String),
1532 Gte(&'static str, ScalarValue),
1533 Lte(&'static str, ScalarValue),
1534 And(Vec<Predicate>),
1535 Or(Vec<Predicate>),
1536 Not(Box<Predicate>),
1537}
1538impl Predicate {
1539 pub fn to_lance(&self) -> String {
1540 match self {
1541 Self::Eq(column, value) => format!("{column} = {}", value.to_lance()),
1542 Self::Ne(column, value) => format!("{column} <> {}", value.to_lance()),
1543 Self::IsNull(column) => format!("{column} IS NULL"),
1544 Self::IsNotNull(column) => format!("{column} IS NOT NULL"),
1545 Self::In(column, values) => {
1546 let values = values
1547 .iter()
1548 .map(ScalarValue::to_lance)
1549 .collect::<Vec<_>>()
1550 .join(", ");
1551 format!("{column} IN ({values})")
1552 }
1553 Self::LikeContains(column, value) => {
1554 format!("{column} LIKE {} ESCAPE '\\'", like_contains(value))
1555 }
1556 Self::Regex(column, pattern) => {
1557 format!("regexp_like({column}, {})", quoted_string(pattern))
1558 }
1559 Self::Gte(column, value) => format!("{column} >= {}", value.to_lance()),
1560 Self::Lte(column, value) => format!("{column} <= {}", value.to_lance()),
1561 Self::And(predicates) => predicates
1562 .iter()
1563 .map(Self::to_lance)
1564 .filter(|predicate| !predicate.is_empty())
1565 .collect::<Vec<_>>()
1566 .join(" AND "),
1567 Self::Or(predicates) => {
1568 let body = predicates
1571 .iter()
1572 .map(Self::to_lance)
1573 .filter(|predicate| !predicate.is_empty())
1574 .collect::<Vec<_>>()
1575 .join(" OR ");
1576 if body.is_empty() {
1577 String::new()
1578 } else {
1579 format!("({body})")
1580 }
1581 }
1582 Self::Not(inner) => {
1583 let body = inner.to_lance();
1584 if body.is_empty() {
1585 String::new()
1586 } else {
1587 format!("NOT ({body})")
1588 }
1589 }
1590 }
1591 }
1592}
1593#[derive(Default)]
1596pub struct ScanOpts<'a> {
1597 pub predicate: Option<&'a Predicate>,
1598 pub projection: Option<&'a [&'a str]>,
1599}
1600
1601impl<'a> ScanOpts<'a> {
1602 pub fn project_only(projection: &'a [&'a str]) -> Self {
1603 Self {
1604 predicate: None,
1605 projection: Some(projection),
1606 }
1607 }
1608 pub fn with_predicate_and_projection(
1609 predicate: &'a Predicate,
1610 projection: &'a [&'a str],
1611 ) -> Self {
1612 Self {
1613 predicate: Some(predicate),
1614 projection: Some(projection),
1615 }
1616 }
1617}
1618
1619impl ScalarValue {
1620 fn to_lance(&self) -> String {
1621 match self {
1622 Self::String(value) => quoted_string(value),
1623 Self::Int32(value) => value.to_string(),
1624 Self::Raw(value) => value.clone(),
1625 }
1626 }
1627}
1628#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1632pub struct RuntimeCaps {
1633 pub index_cache_bytes: Option<usize>,
1634 pub metadata_cache_bytes: Option<usize>,
1635}
1636
1637impl RuntimeCaps {
1638 pub fn from_config(config: &crate::config::RuntimeConfig) -> Self {
1639 Self {
1640 index_cache_bytes: config.index_cache_bytes,
1641 metadata_cache_bytes: config.metadata_cache_bytes,
1642 }
1643 }
1644}
1645
1646const LOCAL_INDEX_CACHE_BYTES: usize = 256 * 1024 * 1024;
1650const LOCAL_METADATA_CACHE_BYTES: usize = 128 * 1024 * 1024;
1651const REMOTE_INDEX_CACHE_BYTES: usize = 1024 * 1024 * 1024;
1656const REMOTE_METADATA_CACHE_BYTES: usize = 512 * 1024 * 1024;
1657
1658fn resolve_cache_caps(location: &Url, caps: RuntimeCaps) -> (usize, usize) {
1659 let (index_default, metadata_default) = if config::is_local(location) {
1660 (LOCAL_INDEX_CACHE_BYTES, LOCAL_METADATA_CACHE_BYTES)
1661 } else {
1662 (REMOTE_INDEX_CACHE_BYTES, REMOTE_METADATA_CACHE_BYTES)
1663 };
1664 (
1665 caps.index_cache_bytes.unwrap_or(index_default),
1666 caps.metadata_cache_bytes.unwrap_or(metadata_default),
1667 )
1668}
1669
1670pub struct Handle {
1671 datasets: DatasetSet,
1672 retry: RetryPolicy,
1673 #[allow(dead_code)]
1681 session: Arc<Session>,
1682 nm: Arc<dyn LanceNamespace>,
1686 nm_ident: NamespaceIdent,
1690 storage_options: HashMap<String, String>,
1695 location: Url,
1699 lazy_refresh_after: Duration,
1703 store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
1707}
1708
1709impl std::fmt::Debug for Handle {
1710 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1711 formatter
1712 .debug_struct("Handle")
1713 .field("datasets", &self.datasets)
1714 .field("retry", &self.retry)
1715 .field("nm_ident", &self.nm_ident)
1716 .field("storage_options", &self.storage_options)
1717 .field("location", &self.location)
1718 .finish()
1719 }
1720}
1721
1722#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1723pub enum Table {
1724 Sessions,
1725 Messages,
1726 Parts,
1727}
1728impl Table {
1729 pub fn as_str(self) -> &'static str {
1730 self.label()
1731 }
1732
1733 fn label(self) -> &'static str {
1734 match self {
1735 Self::Sessions => "sessions",
1736 Self::Messages => "messages",
1737 Self::Parts => "parts",
1738 }
1739 }
1740}
1741#[derive(Debug)]
1742struct DatasetSet {
1743 sessions: OnceCell<Mutex<CachedDataset>>,
1747 messages: Mutex<CachedDataset>,
1748 parts: OnceCell<Mutex<CachedDataset>>,
1756}
1757#[derive(Debug)]
1758struct CachedDataset {
1759 dataset: Dataset,
1760 last_refresh: Instant,
1761 refresh_after: Duration,
1762}
1763impl CachedDataset {
1764 fn new(dataset: Dataset, refresh_after: Duration) -> Self {
1765 Self {
1766 dataset,
1767 last_refresh: Instant::now(),
1768 refresh_after,
1769 }
1770 }
1771 async fn latest(&mut self) -> Result<Dataset> {
1772 if self.last_refresh.elapsed() >= self.refresh_after {
1773 self.dataset.checkout_latest().await?;
1774 self.last_refresh = Instant::now();
1775 }
1776 Ok(self.dataset.clone())
1777 }
1778 fn replace(&mut self, dataset: Dataset) {
1779 self.dataset = dataset;
1780 self.last_refresh = Instant::now();
1781 }
1782}
1783
1784#[derive(Debug, Clone, Copy, Default)]
1789pub struct AppendStats {
1790 pub rows: u64,
1791 pub bytes_written: u64,
1792 pub files_written: u64,
1793 pub attempts: u32,
1794}
1795
1796#[derive(Default)]
1802struct WriteAccum {
1803 rows: std::sync::atomic::AtomicU64,
1804 bytes: std::sync::atomic::AtomicU64,
1805 files: std::sync::atomic::AtomicU64,
1806}
1807
1808impl WriteAccum {
1809 fn observe(&self, stats: &WriteStats) {
1810 use std::sync::atomic::Ordering::Relaxed;
1811 self.rows.fetch_max(stats.rows_written, Relaxed);
1812 self.bytes.fetch_max(stats.bytes_written, Relaxed);
1813 self.files.fetch_max(stats.files_written as u64, Relaxed);
1814 }
1815 fn rows(&self) -> u64 {
1816 self.rows.load(std::sync::atomic::Ordering::Relaxed)
1817 }
1818 fn bytes(&self) -> u64 {
1819 self.bytes.load(std::sync::atomic::Ordering::Relaxed)
1820 }
1821 fn files(&self) -> u64 {
1822 self.files.load(std::sync::atomic::Ordering::Relaxed)
1823 }
1824}
1825
1826fn append_write_params() -> WriteParams {
1831 let mut params = sessions::write_params_for_create();
1832 params.mode = WriteMode::Append;
1833 params.max_bytes_per_file = TARGET_FRAGMENT_BYTES as usize;
1834 params
1835}
1836
1837impl Handle {
1838 pub async fn open(location: &Url) -> Result<Self> {
1841 Self::open_with_options(location, HashMap::new(), RuntimeCaps::default()).await
1842 }
1843
1844 pub fn lance_cache_bytes(&self) -> u64 {
1847 self.session.size_bytes()
1848 }
1849
1850 pub async fn open_with_options(
1857 location: &Url,
1858 storage_options: HashMap<String, String>,
1859 caps: RuntimeCaps,
1860 ) -> Result<Self> {
1861 Self::open_with_options_cached(location, storage_options, caps, None).await
1862 }
1863
1864 pub async fn open_with_options_cached(
1868 location: &Url,
1869 mut storage_options: HashMap<String, String>,
1870 caps: RuntimeCaps,
1871 index_cache_dir: Option<PathBuf>,
1872 ) -> Result<Self> {
1873 if let Some(path) = config::local_path(location) {
1874 tokio::fs::create_dir_all(&path).await.with_context(|| {
1875 format!(
1876 "failed to create data dir {}; fix the storage destination ([storage].path in config) or re-run `pond init`",
1877 path.display()
1878 )
1879 })?;
1880 } else {
1881 apply_remote_storage_defaults(&mut storage_options);
1882 }
1883 let (index_cache_bytes, metadata_cache_bytes) = resolve_cache_caps(location, caps);
1889 let session = Arc::new(Session::new(
1890 index_cache_bytes,
1891 metadata_cache_bytes,
1892 Arc::new(ObjectStoreRegistry::default()),
1893 ));
1894 let root = location.as_str().trim_end_matches('/').to_string();
1900 let mut connect = ConnectBuilder::new("dir")
1901 .property("root", root)
1902 .session(session.clone());
1903 for (key, value) in &storage_options {
1907 connect = connect.property(format!("storage.{key}"), value.clone());
1908 }
1909 let nm: Arc<dyn LanceNamespace> = connect
1910 .connect()
1911 .await
1912 .context("failed to connect lance Directory namespace")?;
1913 let nm_ident = NamespaceIdent::root();
1914 let refresh_after = if config::is_local(location) {
1920 Duration::ZERO
1921 } else {
1922 Duration::from_secs(5)
1923 };
1924 let wrapper = store_wrapper(location, index_cache_dir.as_deref());
1925 let handle = Self {
1926 datasets: DatasetSet {
1927 sessions: OnceCell::new(),
1928 messages: Mutex::new(CachedDataset::new(
1929 open_or_create_via_ns(
1930 &nm,
1931 &nm_ident,
1932 sessions::MESSAGES,
1933 sessions::message_schema(),
1934 &session,
1935 &storage_options,
1936 wrapper.clone(),
1937 )
1938 .await?,
1939 refresh_after,
1940 )),
1941 parts: OnceCell::new(),
1942 },
1943 retry: RetryPolicy::default(),
1944 session,
1945 nm,
1946 nm_ident,
1947 storage_options,
1948 location: location.clone(),
1949 lazy_refresh_after: refresh_after,
1950 store_wrapper: wrapper,
1951 };
1952 Ok(handle)
1953 }
1954
1955 pub fn location(&self) -> &Url {
1956 &self.location
1957 }
1958
1959 pub fn storage_options(&self) -> &HashMap<String, String> {
1963 &self.storage_options
1964 }
1965
1966 fn export_uri(&self, name: &str) -> String {
1972 format!(
1973 "{}/exports/{name}",
1974 self.location.as_str().trim_end_matches('/')
1975 )
1976 }
1977
1978 fn object_store_params(&self) -> ObjectStoreParams {
1982 ObjectStoreParams {
1983 storage_options_accessor: (!self.storage_options.is_empty()).then(|| {
1984 Arc::new(StorageOptionsAccessor::with_static_options(
1985 self.storage_options.clone(),
1986 ))
1987 }),
1988 ..Default::default()
1989 }
1990 }
1991
1992 pub(crate) async fn export_write(&self, name: &str, bytes: &[u8]) -> Result<()> {
1995 let uri = self.export_uri(name);
1996 let registry = Arc::new(ObjectStoreRegistry::default());
1997 let (store, path) =
1998 ObjectStore::from_uri_and_params(registry, &uri, &self.object_store_params())
1999 .await
2000 .with_context(|| format!("failed to open object store for {uri}"))?;
2001 store
2002 .put(&path, bytes)
2003 .await
2004 .with_context(|| format!("failed to write export {uri}"))?;
2005 Ok(())
2006 }
2007
2008 pub(crate) async fn export_read(&self, name: &str) -> Result<Vec<u8>> {
2011 let uri = self.export_uri(name);
2012 let registry = Arc::new(ObjectStoreRegistry::default());
2013 let (store, path) =
2014 ObjectStore::from_uri_and_params(registry, &uri, &self.object_store_params())
2015 .await
2016 .with_context(|| format!("failed to open object store for {uri}"))?;
2017 let bytes = store
2018 .read_one_all(&path)
2019 .await
2020 .with_context(|| format!("failed to read export {uri}"))?;
2021 Ok(bytes.to_vec())
2022 }
2023
2024 pub(crate) fn export_local_path(&self, name: &str) -> Option<std::path::PathBuf> {
2029 if self.location.scheme() != "file" {
2030 return None;
2031 }
2032 let dir = self.location.to_file_path().ok()?;
2033 Some(dir.join("exports").join(name))
2034 }
2035
2036 pub async fn row_counts(&self) -> Result<(usize, usize, usize)> {
2037 Ok((
2038 self.count_rows(Table::Sessions).await?,
2039 self.count_rows(Table::Messages).await?,
2040 self.count_rows(Table::Parts).await?,
2041 ))
2042 }
2043
2044 pub(crate) async fn merge_insert(
2048 &self,
2049 table: Table,
2050 batch: RecordBatch,
2051 row_count: usize,
2052 ) -> Result<u64> {
2053 self.merge_insert_stats(table, batch, row_count)
2054 .await
2055 .map(|stats| stats.num_inserted_rows + stats.num_updated_rows)
2056 }
2057
2058 pub(crate) async fn merge_insert_stats(
2063 &self,
2064 table: Table,
2065 batch: RecordBatch,
2066 row_count: usize,
2067 ) -> Result<MergeStats> {
2068 self.merge(
2069 table,
2070 batch,
2071 row_count,
2072 "merge_insert",
2073 WhenMatched::DoNothing,
2074 WhenNotMatched::InsertAll,
2075 )
2076 .await
2077 }
2078
2079 pub(crate) async fn merge_update(
2082 &self,
2083 table: Table,
2084 batch: RecordBatch,
2085 row_count: usize,
2086 ) -> Result<u64> {
2087 self.merge(
2088 table,
2089 batch,
2090 row_count,
2091 "merge_update",
2092 WhenMatched::UpdateAll,
2093 WhenNotMatched::DoNothing,
2094 )
2095 .await
2096 .map(|stats| stats.num_inserted_rows + stats.num_updated_rows)
2097 }
2098
2099 async fn write_committed<E, Fut, P>(&self, table: Table, execute: E) -> Result<P>
2108 where
2109 E: Fn(Arc<Dataset>) -> Fut,
2110 Fut: std::future::Future<Output = Result<(Dataset, P)>>,
2111 {
2112 self.write_committed_with(table, |_| true, execute).await
2113 }
2114
2115 async fn write_committed_with<E, Fut, P, R>(
2121 &self,
2122 table: Table,
2123 should_retry: R,
2124 execute: E,
2125 ) -> Result<P>
2126 where
2127 E: Fn(Arc<Dataset>) -> Fut,
2128 Fut: std::future::Future<Output = Result<(Dataset, P)>>,
2129 R: Fn(&anyhow::Error) -> bool,
2130 {
2131 self.retry_lance_filtered(table.label(), should_retry, || {
2132 let execute = &execute;
2133 async move {
2134 let mut cached = self.cached(table).await?.lock().await;
2135 let existing = cached.latest().await?;
2136 let (dataset, payload) = execute(Arc::new(existing)).await?;
2137 cached.replace(dataset);
2138 Ok(payload)
2139 }
2140 })
2141 .await
2142 }
2143
2144 async fn merge(
2150 &self,
2151 table: Table,
2152 batch: RecordBatch,
2153 row_count: usize,
2154 op: &'static str,
2155 when_matched: WhenMatched,
2156 when_not_matched: WhenNotMatched,
2157 ) -> Result<MergeStats> {
2158 if row_count == 0 {
2159 return Ok(MergeStats::default());
2160 }
2161 let started = Instant::now();
2162 let result = self
2163 .write_committed(table, |existing| {
2164 let batch = batch.clone();
2165 let when_matched = when_matched.clone();
2166 let when_not_matched = when_not_matched.clone();
2167 async move {
2168 let schema = batch.schema();
2169 let reader = RecordBatchIterator::new([Ok(batch)], schema);
2170 let mut builder = MergeInsertBuilder::try_new(existing, Vec::new())?;
2171 builder.when_matched(when_matched);
2172 builder.when_not_matched(when_not_matched);
2173 builder.source_dedupe_behavior(SourceDedupeBehavior::FirstSeen);
2176 builder.skip_auto_cleanup(true);
2180 let (dataset, stats) = builder
2181 .try_build()?
2182 .execute_reader(Box::new(reader))
2183 .await?;
2184 Ok((dataset.as_ref().clone(), stats))
2185 }
2186 })
2187 .await;
2188 let skipped = result
2189 .as_ref()
2190 .map(|s| s.num_skipped_duplicates)
2191 .unwrap_or(0);
2192 tracing::info!(
2193 target: "pond::perf",
2194 op,
2195 table = %table.label(),
2196 rows = row_count,
2197 elapsed_ms = started.elapsed().as_millis() as u64,
2198 skipped,
2199 "merge",
2200 );
2201 result
2202 }
2203
2204 pub(crate) async fn append_stream<F, Fut>(
2225 &self,
2226 table: Table,
2227 make_source: F,
2228 ) -> Result<AppendStats>
2229 where
2230 F: Fn() -> Fut,
2231 Fut: std::future::Future<Output = Result<SendableRecordBatchStream>>,
2232 {
2233 let cum = Arc::new(WriteAccum::default());
2234 let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
2235 let started = Instant::now();
2236 self.write_committed(table, |existing| {
2237 let make_source = &make_source;
2238 let cum = cum.clone();
2239 let attempts = attempts.clone();
2240 async move {
2241 attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2242 let stream = make_source().await?;
2243 let dataset = InsertBuilder::new(existing)
2244 .with_params(&append_write_params())
2245 .progress(move |stats| cum.observe(&stats))
2246 .execute_stream(stream)
2247 .await?;
2248 Ok((dataset, ()))
2249 }
2250 })
2251 .await?;
2252
2253 let attempts = attempts.load(std::sync::atomic::Ordering::Relaxed);
2254 let stats = AppendStats {
2255 rows: cum.rows(),
2256 bytes_written: cum.bytes(),
2257 files_written: cum.files(),
2258 attempts,
2259 };
2260 tracing::info!(
2261 target: "pond::perf",
2262 op = "append",
2263 table = %table.label(),
2264 rows = stats.rows,
2265 files = stats.files_written,
2266 attempts,
2267 elapsed_ms = started.elapsed().as_millis() as u64,
2268 "append",
2269 );
2270 Ok(stats)
2271 }
2272
2273 pub(crate) async fn append_batches(
2284 &self,
2285 table: Table,
2286 batches: Vec<RecordBatch>,
2287 ) -> Result<AppendStats> {
2288 let total_rows: u64 = batches.iter().map(|batch| batch.num_rows() as u64).sum();
2289 if total_rows == 0 {
2290 return Ok(AppendStats::default());
2291 }
2292 let cum = Arc::new(WriteAccum::default());
2293 let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
2294 let started = Instant::now();
2295 self.write_committed_with(table, is_commit_conflict, |existing| {
2296 let cum = cum.clone();
2297 let attempts = attempts.clone();
2298 let batches = batches.clone();
2299 async move {
2300 attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2301 let dataset = InsertBuilder::new(existing)
2302 .with_params(&append_write_params())
2303 .progress(move |stats| cum.observe(&stats))
2304 .execute(batches)
2305 .await?;
2306 Ok((dataset, ()))
2307 }
2308 })
2309 .await?;
2310
2311 let attempts = attempts.load(std::sync::atomic::Ordering::Relaxed);
2312 let stats = AppendStats {
2313 rows: total_rows,
2314 bytes_written: cum.bytes(),
2315 files_written: cum.files(),
2316 attempts,
2317 };
2318 tracing::info!(
2319 target: "pond::perf",
2320 op = "append_batches",
2321 table = %table.label(),
2322 rows = stats.rows,
2323 files = stats.files_written,
2324 attempts,
2325 elapsed_ms = started.elapsed().as_millis() as u64,
2326 "append",
2327 );
2328 Ok(stats)
2329 }
2330
2331 pub async fn optimize_table(
2340 &self,
2341 table: Table,
2342 intents: &[IndexIntent],
2343 progress: Option<&OptimizeProgressFn>,
2344 policy: &MaintenancePolicy,
2345 ) -> TableOptimizeOutcome {
2346 let compaction = self
2347 .run_optimize_compact_phase(table, progress, policy)
2348 .await;
2349 let indices = self
2350 .run_optimize_indices_phase(table, intents, progress, policy.fold_thresholds())
2351 .await;
2352 TableOptimizeOutcome {
2353 table,
2354 indices,
2355 compaction,
2356 }
2357 }
2358
2359 pub async fn optimize_table_indices_only(
2364 &self,
2365 table: Table,
2366 intents: &[IndexIntent],
2367 progress: Option<&OptimizeProgressFn>,
2368 ) -> PhaseOutcome {
2369 self.run_optimize_indices_phase(
2373 table,
2374 intents,
2375 progress,
2376 FoldThresholds {
2377 scalar: 0,
2378 index: 0,
2379 },
2380 )
2381 .await
2382 }
2383
2384 async fn run_optimize_indices_phase(
2385 &self,
2386 table: Table,
2387 intents: &[IndexIntent],
2388 progress: Option<&OptimizeProgressFn>,
2389 folds: FoldThresholds,
2390 ) -> PhaseOutcome {
2391 if intents.is_empty() {
2392 return PhaseOutcome::Noop;
2393 }
2394 let result = self
2395 .retry_lance(table.label(), || async {
2396 let mut guard = self.cached(table).await?.lock().await;
2397 let mut dataset = guard.latest().await?;
2398 let did_work =
2399 optimize_table_indices(&mut dataset, intents, table, progress, folds).await?;
2400 guard.replace(dataset);
2401 Ok::<_, anyhow::Error>(did_work)
2402 })
2403 .await;
2404 match result {
2405 Ok(true) => PhaseOutcome::Ok,
2406 Ok(false) => PhaseOutcome::Noop,
2407 Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
2408 Err(error) => PhaseOutcome::Failed(error),
2409 }
2410 }
2411
2412 async fn run_optimize_compact_phase(
2413 &self,
2414 table: Table,
2415 progress: Option<&OptimizeProgressFn>,
2416 policy: &MaintenancePolicy,
2417 ) -> PhaseOutcome {
2418 let result = self
2419 .retry_lance(table.label(), || async {
2420 let mut guard = self.cached(table).await?.lock().await;
2421 let mut dataset = guard.latest().await?;
2422 optimize_table_compact(&mut dataset, table, progress, policy).await?;
2423 guard.replace(dataset);
2424 Ok::<_, anyhow::Error>(())
2425 })
2426 .await;
2427 match result {
2428 Ok(()) => PhaseOutcome::Ok,
2429 Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
2430 Err(error) => PhaseOutcome::Failed(error),
2431 }
2432 }
2433
2434 pub async fn rebuild_index(
2435 &self,
2436 table: Table,
2437 intent: &IndexIntent,
2438 progress: Option<&OptimizeProgressFn>,
2439 ) -> Result<()> {
2440 emit(
2441 progress,
2442 OptimizeEvent::PhaseStart {
2443 table,
2444 phase: OptimizePhase::IndexRebuild,
2445 detail: Some(intent.name.to_owned()),
2446 },
2447 );
2448 let started = Instant::now();
2449 let result = self
2450 .retry_lance(table.label(), || async {
2451 let mut guard = self.cached(table).await?.lock().await;
2452 let mut dataset = guard.latest().await?;
2453 rebuild_index(&mut dataset, intent, progress, table).await?;
2454 guard.replace(dataset);
2455 Ok(())
2456 })
2457 .await;
2458 emit(
2459 progress,
2460 OptimizeEvent::PhaseDone {
2461 table,
2462 phase: OptimizePhase::IndexRebuild,
2463 elapsed_ms: started.elapsed().as_millis() as u64,
2464 },
2465 );
2466 result
2467 }
2468
2469 pub async fn cleanup_table_versions(
2473 &self,
2474 table: Table,
2475 older_than: chrono::Duration,
2476 ) -> Result<()> {
2477 let mut guard = self.cached(table).await?.lock().await;
2478 let dataset = guard.latest().await?;
2479 dataset
2480 .cleanup_old_versions(older_than, Some(false), Some(false))
2481 .await
2482 .with_context(|| format!("cleanup_old_versions failed for {}", table.label()))?;
2483 Ok(())
2484 }
2485
2486 pub async fn index_status(
2487 &self,
2488 table: Table,
2489 intents: &[IndexIntent],
2490 indexable_only: bool,
2491 ) -> Result<Vec<IndexStatus>> {
2492 let dataset = self.dataset(table).await?;
2493 index_status(table, &dataset, intents, indexable_only).await
2494 }
2495
2496 pub(crate) async fn dataset(&self, table: Table) -> Result<Dataset> {
2497 let mut cached = self.cached(table).await?.lock().await;
2498 cached.latest().await
2499 }
2500 pub(crate) async fn scanner(
2505 &self,
2506 table: Table,
2507 predicate: Option<&Predicate>,
2508 ) -> Result<lance::dataset::scanner::Scanner> {
2509 let dataset = self.dataset(table).await?;
2510 scanner_with_prefilter(&dataset, predicate)
2511 }
2512 pub async fn scan(
2515 &self,
2516 table: Table,
2517 opts: ScanOpts<'_>,
2518 ) -> Result<lance::dataset::scanner::Scanner> {
2519 let mut scanner = self.scanner(table, opts.predicate).await?;
2520 if let Some(projection) = opts.projection {
2521 scanner.project(projection)?;
2522 }
2523 Ok(scanner)
2524 }
2525 pub(crate) async fn scan_batch(
2526 &self,
2527 table: Table,
2528 predicate: Option<&Predicate>,
2529 projection: &[&str],
2530 ) -> Result<RecordBatch> {
2531 let opts = ScanOpts {
2532 predicate,
2533 projection: (!projection.is_empty()).then_some(projection),
2534 };
2535 self.scan(table, opts)
2536 .await?
2537 .try_into_batch()
2538 .await
2539 .context("scan failed")
2540 }
2541 pub async fn count_rows(&self, table: Table) -> Result<usize> {
2542 self.dataset(table)
2543 .await?
2544 .count_rows(None)
2545 .await
2546 .map_err(Into::into)
2547 }
2548 pub async fn collect_ids(&self, table: Table) -> Result<std::collections::HashSet<String>> {
2554 let batch = self.scan_batch(table, None, &["id"]).await?;
2555 let ids = batch
2556 .column_by_name("id")
2557 .context("scan projection dropped the id column")?
2558 .as_any()
2559 .downcast_ref::<StringArray>()
2560 .context("id column is not Utf8")?;
2561 Ok(ids.iter().flatten().map(str::to_owned).collect())
2562 }
2563 #[cfg(test)]
2565 pub(crate) async fn messages_index_names(&self) -> Result<Vec<String>> {
2566 let dataset = self.dataset(Table::Messages).await?;
2567 let indices = dataset.load_indices().await?;
2568 Ok(indices.iter().map(|index| index.name.clone()).collect())
2569 }
2570
2571 pub(crate) async fn messages_has_index(&self, name: &str) -> Result<bool> {
2577 let dataset = self.dataset(Table::Messages).await?;
2578 let indices = dataset.load_indices().await?;
2579 Ok(indices.iter().any(|index| index.name == name))
2580 }
2581
2582 pub(crate) async fn messages_fast_search_ready(&self, name: &str) -> Result<bool> {
2589 let dataset = self.dataset(Table::Messages).await?;
2590 if !dataset
2591 .load_indices()
2592 .await?
2593 .iter()
2594 .any(|index| index.name == name)
2595 {
2596 return Ok(false);
2597 }
2598 let unindexed = dataset
2599 .unindexed_fragments(name)
2600 .await
2601 .with_context(|| format!("unindexed_fragments failed for {name}"))?;
2602 Ok(unindexed.is_empty())
2603 }
2604
2605 pub(crate) async fn prune_index_cache(&self, cache_dir: &std::path::Path) {
2610 if config::is_local(&self.location) {
2611 return;
2612 }
2613 let root = cache_dir.join(store_key(&self.location)).join("indices");
2614 if !root.exists() {
2615 return;
2616 }
2617 let mut keep = std::collections::HashSet::new();
2618 for table in [Table::Sessions, Table::Messages, Table::Parts] {
2619 let Ok(dataset) = self.dataset(table).await else {
2620 return;
2621 };
2622 let Ok(indices) = dataset.load_indices().await else {
2623 return;
2624 };
2625 keep.extend(indices.iter().map(|index| index.uuid.to_string()));
2626 }
2627 prune_stale_uuid_dirs(&root, &keep);
2628 }
2629
2630 pub(crate) async fn unindexed_row_count(
2633 &self,
2634 table: Table,
2635 index_name: &str,
2636 ) -> Result<usize> {
2637 let dataset = self.dataset(table).await?;
2638 let fragments = dataset
2639 .unindexed_fragments(index_name)
2640 .await
2641 .with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
2642 Ok(fragments
2643 .iter()
2644 .map(|fragment| fragment.num_rows().unwrap_or(0))
2645 .sum())
2646 }
2647
2648 pub(crate) async fn find_index_owner(&self, name: &str) -> Result<Option<Table>> {
2655 let list = |table: Table| async move {
2656 let dataset = self.dataset(table).await?;
2657 let names: Vec<String> = dataset
2658 .load_indices()
2659 .await
2660 .with_context(|| format!("load_indices failed for {}", table.label()))?
2661 .iter()
2662 .map(|index| index.name.clone())
2663 .collect();
2664 Ok::<_, anyhow::Error>(names)
2665 };
2666 let (sessions, messages, parts) = tokio::try_join!(
2667 list(Table::Sessions),
2668 list(Table::Messages),
2669 list(Table::Parts),
2670 )?;
2671 for (table, names) in [
2672 (Table::Sessions, sessions),
2673 (Table::Messages, messages),
2674 (Table::Parts, parts),
2675 ] {
2676 if names.iter().any(|n| n == name) {
2677 return Ok(Some(table));
2678 }
2679 }
2680 Ok(None)
2681 }
2682
2683 pub(crate) async fn drop_index(&self, table: Table, name: &str) -> Result<()> {
2689 let mut guard = self.cached(table).await?.lock().await;
2690 let mut dataset = guard.latest().await?;
2691 dataset
2692 .drop_index(name)
2693 .await
2694 .with_context(|| format!("drop_index({name}) failed for {}", table.label()))?;
2695 guard.replace(dataset);
2696 Ok(())
2697 }
2698
2699 async fn table_location(&self, table_name: &str) -> Result<String> {
2702 let request = DescribeTableRequest {
2703 id: Some(self.nm_ident.as_table_id(table_name)),
2704 ..Default::default()
2705 };
2706 let response = self
2707 .nm
2708 .describe_table(request)
2709 .await
2710 .with_context(|| format!("failed to describe table {table_name}"))?;
2711 response
2712 .location
2713 .with_context(|| format!("namespace returned no location for table {table_name}"))
2714 }
2715
2716 pub async fn initialized(&self) -> Result<bool> {
2722 let request = DescribeTableRequest {
2723 id: Some(self.nm_ident.as_table_id(sessions::PARTS)),
2724 ..Default::default()
2725 };
2726 match self.nm.describe_table(request).await {
2727 Ok(_) => Ok(true),
2728 Err(error) if is_namespace_error_code(&error, ErrorCode::TableNotFound) => Ok(false),
2729 Err(error) => {
2730 Err(anyhow::Error::from(error)).context("failed to probe table existence")
2731 }
2732 }
2733 }
2734
2735 pub async fn table_sizes(&self) -> Result<TableSizes> {
2739 let registry = Arc::new(ObjectStoreRegistry::default());
2740 let params = self.object_store_params();
2741
2742 let sessions = self
2743 .listed_size(
2744 ®istry,
2745 ¶ms,
2746 &self.table_location(sessions::SESSIONS).await?,
2747 )
2748 .await?;
2749 let messages = self
2750 .listed_size(
2751 ®istry,
2752 ¶ms,
2753 &self.table_location(sessions::MESSAGES).await?,
2754 )
2755 .await?;
2756 let parts = self
2757 .listed_size(
2758 ®istry,
2759 ¶ms,
2760 &self.table_location(sessions::PARTS).await?,
2761 )
2762 .await?;
2763 let root_total = self
2766 .listed_size(®istry, ¶ms, self.location.as_str())
2767 .await?;
2768 let other = root_total.saturating_sub(sessions + messages + parts);
2769 let sessions_data = self
2770 .data_liveness(®istry, ¶ms, Table::Sessions, sessions::SESSIONS)
2771 .await?;
2772 let messages_data = self
2773 .data_liveness(®istry, ¶ms, Table::Messages, sessions::MESSAGES)
2774 .await?;
2775 let parts_data = self
2776 .data_liveness(®istry, ¶ms, Table::Parts, sessions::PARTS)
2777 .await?;
2778 Ok(TableSizes {
2779 sessions,
2780 messages,
2781 parts,
2782 other,
2783 sessions_data,
2784 messages_data,
2785 parts_data,
2786 })
2787 }
2788
2789 async fn data_liveness(
2790 &self,
2791 registry: &Arc<ObjectStoreRegistry>,
2792 params: &ObjectStoreParams,
2793 table: Table,
2794 table_name: &str,
2795 ) -> Result<DataLiveness> {
2796 let location = self.table_location(table_name).await?;
2797 let data_dir = format!("{}/data", location.trim_end_matches('/'));
2798 let on_disk = self.listed_size(registry, params, &data_dir).await?;
2799 let dataset = self.dataset(table).await?;
2800 let live = dataset
2801 .get_fragments()
2802 .iter()
2803 .try_fold(0u64, |total, fragment| {
2804 Some(total + fragment_bytes(fragment.metadata())?)
2805 });
2806 Ok(DataLiveness { on_disk, live })
2807 }
2808
2809 async fn listed_size(
2811 &self,
2812 registry: &Arc<ObjectStoreRegistry>,
2813 params: &ObjectStoreParams,
2814 uri: &str,
2815 ) -> Result<u64> {
2816 let (store, base) = ObjectStore::from_uri_and_params(registry.clone(), uri, params)
2817 .await
2818 .with_context(|| format!("failed to open object store for {uri}"))?;
2819 let mut listing = store.list(Some(base));
2820 let mut total = 0u64;
2821 while let Some(meta) = listing.next().await {
2822 let meta = meta.with_context(|| format!("listing {uri} failed"))?;
2823 total += meta.size;
2824 }
2825 Ok(total)
2826 }
2827 async fn cached(&self, table: Table) -> Result<&Mutex<CachedDataset>> {
2828 match table {
2829 Table::Sessions => self.sessions_cached().await,
2830 Table::Messages => Ok(&self.datasets.messages),
2831 Table::Parts => self.parts_cached().await,
2832 }
2833 }
2834
2835 async fn sessions_cached(&self) -> Result<&Mutex<CachedDataset>> {
2840 self.lazy_cached(
2841 &self.datasets.sessions,
2842 sessions::SESSIONS,
2843 sessions::session_schema,
2844 )
2845 .await
2846 }
2847
2848 async fn parts_cached(&self) -> Result<&Mutex<CachedDataset>> {
2851 self.lazy_cached(&self.datasets.parts, sessions::PARTS, sessions::part_schema)
2852 .await
2853 }
2854
2855 async fn lazy_cached<'a>(
2859 &self,
2860 cell: &'a OnceCell<Mutex<CachedDataset>>,
2861 table_name: &str,
2862 schema: fn() -> lance::deps::arrow_schema::SchemaRef,
2863 ) -> Result<&'a Mutex<CachedDataset>> {
2864 cell.get_or_try_init(|| async {
2865 let dataset = open_or_create_via_ns(
2866 &self.nm,
2867 &self.nm_ident,
2868 table_name,
2869 schema(),
2870 &self.session,
2871 &self.storage_options,
2872 self.store_wrapper.clone(),
2873 )
2874 .await?;
2875 Ok::<_, anyhow::Error>(Mutex::new(CachedDataset::new(
2876 dataset,
2877 self.lazy_refresh_after,
2878 )))
2879 })
2880 .await
2881 }
2882 async fn retry_lance<T, Fut, Op>(&self, label: &str, operation: Op) -> Result<T>
2883 where
2884 Fut: std::future::Future<Output = Result<T>>,
2885 Op: FnMut() -> Fut,
2886 {
2887 self.retry_lance_filtered(label, |_| true, operation).await
2889 }
2890
2891 async fn retry_lance_filtered<T, Fut, Op, R>(
2901 &self,
2902 label: &str,
2903 should_retry: R,
2904 mut operation: Op,
2905 ) -> Result<T>
2906 where
2907 Fut: std::future::Future<Output = Result<T>>,
2908 Op: FnMut() -> Fut,
2909 R: Fn(&anyhow::Error) -> bool,
2910 {
2911 let mut attempt = 0u8;
2912 loop {
2913 attempt = attempt.saturating_add(1);
2914 match operation().await {
2915 Ok(value) => return Ok(value),
2916 Err(error) if attempt < self.retry.attempts && should_retry(&error) => {
2917 let backoff = self.backoff(attempt);
2918 let error_chain = format!("{error:#}");
2921 tracing::warn!(
2922 label,
2923 attempt,
2924 ?backoff,
2925 error = %error_chain,
2926 "retrying Lance operation"
2927 );
2928 tokio::time::sleep(backoff).await;
2929 }
2930 Err(error) => {
2931 let error_chain = format!("{error:#}");
2932 tracing::warn!(
2933 label,
2934 attempt,
2935 error = %error_chain,
2936 "Lance operation exhausted retries"
2937 );
2938 if is_commit_conflict(&error) {
2945 return Err(error.context(ConflictExhausted { attempts: attempt }));
2946 }
2947 return Err(error);
2948 }
2949 }
2950 }
2951 }
2952 fn backoff(&self, attempt: u8) -> Duration {
2953 let shift = u32::from(attempt.saturating_sub(1));
2954 let multiplier = 1u32.checked_shl(shift).unwrap_or(u32::MAX);
2955 let base = self.retry.initial_backoff.saturating_mul(multiplier);
2956 let factor = (1.0 + self.retry.jitter * (fastrand::f64() * 2.0 - 1.0)).max(0.0);
2959 base.mul_f64(factor).min(self.retry.max_backoff)
2960 }
2961}
2962async fn optimize_table_compact(
2983 dataset: &mut Dataset,
2984 table: Table,
2985 progress: Option<&OptimizeProgressFn>,
2986 policy: &MaintenancePolicy,
2987) -> Result<()> {
2988 let stats: Vec<FragmentStat> = dataset
2989 .get_fragments()
2990 .iter()
2991 .map(|fragment| fragment_stat(fragment.metadata()))
2992 .collect();
2993 let compaction = CompactionOptions {
2994 target_rows_per_fragment: derived_target_rows(&stats),
2995 max_bytes_per_file: Some(TARGET_FRAGMENT_BYTES as usize),
2996 defer_index_remap: false,
2997 compaction_mode: Some(CompactionMode::TryBinaryCopy),
3002 ..CompactionOptions::default()
3003 };
3004
3005 let mut plan = plan_compaction(dataset, &compaction).await?;
3006 if policy.compaction_fragment_cap > 0 {
3007 let max_bytes_per_file = compaction
3008 .max_bytes_per_file
3009 .and_then(|bytes| u64::try_from(bytes).ok())
3010 .unwrap_or_default();
3011 plan.tasks.retain(|task| {
3012 let task_stats: Vec<FragmentStat> = task.fragments.iter().map(fragment_stat).collect();
3013 let reason = task_veto_reason(
3014 &task_stats,
3015 policy.compaction_fragment_cap,
3016 compaction.materialize_deletions_threshold,
3017 compaction.target_rows_per_fragment,
3018 max_bytes_per_file,
3019 );
3020 if let Some(reason) = reason {
3021 tracing::debug!(
3022 target: "pond::perf",
3023 table = table.as_str(),
3024 fragments = task_stats.len(),
3025 reason,
3026 "compaction task vetoed",
3027 );
3028 }
3029 reason.is_none()
3030 });
3031 }
3032 if plan.tasks.is_empty() {
3033 tracing::debug!(
3034 target: "pond::perf",
3035 table = table.as_str(),
3036 "compaction skipped: no task to run",
3037 );
3038 } else {
3039 emit(
3040 progress,
3041 OptimizeEvent::PhaseStart {
3042 table,
3043 phase: OptimizePhase::Compact,
3044 detail: None,
3045 },
3046 );
3047 let started = Instant::now();
3048 let mut completed = Vec::with_capacity(plan.tasks.len());
3049 for task in plan.compaction_tasks() {
3050 completed.push(task.execute(dataset).await?);
3051 }
3052 commit_compaction(
3053 dataset,
3054 completed,
3055 Arc::new(DatasetIndexRemapperOptions::default()),
3056 &compaction,
3057 )
3058 .await?;
3059 emit(
3060 progress,
3061 OptimizeEvent::PhaseDone {
3062 table,
3063 phase: OptimizePhase::Compact,
3064 elapsed_ms: started.elapsed().as_millis() as u64,
3065 },
3066 );
3067 }
3068
3069 if cleanup_due(dataset.version_id(), policy.cleanup_interval) {
3080 emit(
3081 progress,
3082 OptimizeEvent::PhaseStart {
3083 table,
3084 phase: OptimizePhase::Cleanup,
3085 detail: None,
3086 },
3087 );
3088 let started = Instant::now();
3089 dataset
3098 .cleanup_old_versions(policy.cleanup_older_than, Some(false), Some(false))
3099 .await
3100 .context("cleanup_old_versions failed during index optimize")?;
3101 emit(
3102 progress,
3103 OptimizeEvent::PhaseDone {
3104 table,
3105 phase: OptimizePhase::Cleanup,
3106 elapsed_ms: started.elapsed().as_millis() as u64,
3107 },
3108 );
3109 }
3110
3111 Ok(())
3112}
3113
3114fn cleanup_due(version: u64, interval: u64) -> bool {
3120 interval <= 1 || version.is_multiple_of(interval)
3121}
3122
3123async fn optimize_table_indices(
3128 dataset: &mut Dataset,
3129 intents: &[IndexIntent],
3130 table: Table,
3131 progress: Option<&OptimizeProgressFn>,
3132 folds: FoldThresholds,
3133) -> Result<bool> {
3134 let existing = dataset.load_indices().await?;
3135 let existing_names: std::collections::HashSet<String> =
3136 existing.iter().map(|index| index.name.clone()).collect();
3137
3138 let mut append_indices: Vec<String> = Vec::new();
3139 let mut did_work = false;
3140
3141 for intent in intents {
3142 let exists = existing_names.contains(intent.name);
3143
3144 if !exists {
3145 if !intent.trigger.should_create(dataset).await? {
3146 continue;
3147 }
3148 let params = intent.params.build(dataset).await?;
3149 let index_type = intent.params.index_type();
3150 tracing::info!(
3151 index = intent.name,
3152 column = intent.column,
3153 "creating Lance index (trigger fired)",
3154 );
3155 emit(
3156 progress,
3157 OptimizeEvent::PhaseStart {
3158 table,
3159 phase: OptimizePhase::IndexCreate,
3160 detail: Some(intent.name.to_owned()),
3161 },
3162 );
3163 let started = Instant::now();
3164 dataset
3165 .create_index_builder(&[intent.column], index_type, params.as_ref())
3166 .name(intent.name.to_owned())
3167 .replace(false)
3168 .progress(lance_progress(progress, table, intent.name))
3169 .await
3170 .with_context(|| format!("failed to create index {}", intent.name))?;
3171 emit(
3172 progress,
3173 OptimizeEvent::PhaseDone {
3174 table,
3175 phase: OptimizePhase::IndexCreate,
3176 elapsed_ms: started.elapsed().as_millis() as u64,
3177 },
3178 );
3179 did_work = true;
3180 continue;
3181 }
3182
3183 let unindexed = dataset.unindexed_fragments(intent.name).await?;
3189 if unindexed.is_empty() {
3190 continue;
3191 }
3192 let tail_rows: usize = unindexed
3193 .iter()
3194 .map(|fragment| fragment.num_rows().unwrap_or(0))
3195 .sum();
3196 let fold_threshold = match intent.params {
3205 IndexParamsKind::Scalar(_) => folds.scalar,
3206 IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. } => folds.index,
3207 };
3208 if fold_threshold > 0 && tail_rows < fold_threshold {
3209 tracing::debug!(
3210 target: "pond::perf",
3211 index = intent.name,
3212 tail_rows,
3213 threshold = fold_threshold,
3214 "deferring index fold (unindexed tail below threshold)",
3215 );
3216 continue;
3217 }
3218 if matches!(intent.params, IndexParamsKind::InvertedFtsWord)
3226 && !column_has_values(dataset, intent.column, &unindexed).await?
3227 {
3228 tracing::debug!(
3229 target: "pond::perf",
3230 index = intent.name,
3231 tail_rows,
3232 "skipping FTS fold (tail has no indexable values)",
3233 );
3234 continue;
3235 }
3236 append_indices.push(intent.name.to_owned());
3242 }
3243
3244 if !append_indices.is_empty() {
3245 let segment_count = |name: &str| {
3251 existing
3252 .iter()
3253 .filter(|index| index.name.as_str() == name)
3254 .count()
3255 };
3256 let (consolidate, to_append): (Vec<String>, Vec<String>) = append_indices
3257 .iter()
3258 .cloned()
3259 .partition(|name| segment_count(name) >= DELTA_MERGE_THRESHOLD);
3260 let mut fts_rebuilds: Vec<&IndexIntent> = Vec::new();
3269 let mut to_merge: Vec<String> = Vec::new();
3270 for name in consolidate {
3271 let fts_intent = intents.iter().find(|intent| {
3272 intent.name == name && matches!(intent.params, IndexParamsKind::InvertedFtsWord)
3273 });
3274 match fts_intent {
3275 Some(intent) => fts_rebuilds.push(intent),
3276 None => to_merge.push(name),
3277 }
3278 }
3279
3280 emit(
3281 progress,
3282 OptimizeEvent::PhaseStart {
3283 table,
3284 phase: OptimizePhase::IndexAppend,
3285 detail: Some(append_indices.join(", ")),
3286 },
3287 );
3288 let started = Instant::now();
3289 if !to_append.is_empty() {
3290 dataset
3291 .optimize_indices(&OptimizeOptions::append().index_names(to_append))
3292 .await
3293 .context("optimize_indices(append) failed during index optimize")?;
3294 }
3295 if !to_merge.is_empty() {
3296 dataset
3297 .optimize_indices(
3298 &OptimizeOptions::merge(DELTA_MERGE_THRESHOLD).index_names(to_merge),
3299 )
3300 .await
3301 .context("optimize_indices(merge) failed during index optimize")?;
3302 }
3303 emit(
3304 progress,
3305 OptimizeEvent::PhaseDone {
3306 table,
3307 phase: OptimizePhase::IndexAppend,
3308 elapsed_ms: started.elapsed().as_millis() as u64,
3309 },
3310 );
3311 for intent in &fts_rebuilds {
3312 emit(
3313 progress,
3314 OptimizeEvent::PhaseStart {
3315 table,
3316 phase: OptimizePhase::IndexRebuild,
3317 detail: Some(intent.name.to_owned()),
3318 },
3319 );
3320 let rebuild_started = Instant::now();
3321 rebuild_index(dataset, intent, progress, table).await?;
3322 emit(
3323 progress,
3324 OptimizeEvent::PhaseDone {
3325 table,
3326 phase: OptimizePhase::IndexRebuild,
3327 elapsed_ms: rebuild_started.elapsed().as_millis() as u64,
3328 },
3329 );
3330 }
3331 tracing::debug!(
3332 target: "pond::perf",
3333 indices = ?append_indices,
3334 rebuilt = ?fts_rebuilds,
3335 "folded trailing fragments into indices",
3336 );
3337 did_work = true;
3338 }
3339
3340 Ok(did_work)
3341}
3342
3343fn non_null_scanner(
3346 dataset: &Dataset,
3347 column: &'static str,
3348 fragments: &[lance::table::format::Fragment],
3349) -> Result<lance::dataset::scanner::Scanner> {
3350 let mut scanner = dataset.scan();
3351 scanner.with_fragments(fragments.to_vec());
3352 scanner.filter(&Predicate::IsNotNull(column).to_lance())?;
3353 Ok(scanner)
3354}
3355
3356async fn column_has_values(
3361 dataset: &Dataset,
3362 column: &'static str,
3363 fragments: &[lance::table::format::Fragment],
3364) -> Result<bool> {
3365 let mut scanner = non_null_scanner(dataset, column, fragments)?;
3366 scanner.project(&[column])?;
3367 scanner.limit(Some(1), None)?;
3368 let batch = scanner
3369 .try_into_batch()
3370 .await
3371 .with_context(|| format!("non-null probe on {column} failed"))?;
3372 Ok(batch.num_rows() > 0)
3373}
3374
3375async fn column_value_count(
3378 dataset: &Dataset,
3379 column: &'static str,
3380 fragments: &[lance::table::format::Fragment],
3381) -> Result<usize> {
3382 let count = non_null_scanner(dataset, column, fragments)?
3383 .count_rows()
3384 .await
3385 .with_context(|| format!("non-null count on {column} failed"))?;
3386 Ok(count as usize)
3387}
3388
3389async fn rebuild_index(
3390 dataset: &mut Dataset,
3391 intent: &IndexIntent,
3392 progress: Option<&OptimizeProgressFn>,
3393 table: Table,
3394) -> Result<()> {
3395 if !intent.trigger.should_create(dataset).await? {
3396 return Ok(());
3397 }
3398 let params = intent.params.build(dataset).await?;
3399 dataset
3400 .create_index_builder(
3401 &[intent.column],
3402 intent.params.index_type(),
3403 params.as_ref(),
3404 )
3405 .name(intent.name.to_owned())
3406 .replace(true)
3407 .progress(lance_progress(progress, table, intent.name))
3408 .await
3409 .with_context(|| format!("failed to rebuild index {}", intent.name))?;
3410 Ok(())
3411}
3412
3413async fn index_status(
3414 table: Table,
3415 dataset: &Dataset,
3416 intents: &[IndexIntent],
3417 indexable_only: bool,
3418) -> Result<Vec<IndexStatus>> {
3419 let existing = dataset.load_indices().await?;
3420 let existing_names: std::collections::HashSet<String> =
3421 existing.iter().map(|index| index.name.clone()).collect();
3422 let total_fragments = dataset.get_fragments().len();
3423 let total_rows = dataset.count_rows(None).await?;
3424 let mut statuses = Vec::with_capacity(intents.len());
3425 for intent in intents {
3426 let exists = existing_names.contains(intent.name);
3427 if !exists {
3428 statuses.push(IndexStatus {
3429 table,
3430 intent_name: intent.name.to_owned(),
3431 fragments_covered: 0,
3432 unindexed_fragments: total_fragments,
3433 unindexed_rows: total_rows,
3434 exists,
3435 });
3436 continue;
3437 }
3438 let unindexed = dataset
3439 .unindexed_fragments(intent.name)
3440 .await
3441 .with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
3442 let unindexed_fragments = unindexed.len();
3443 let mut unindexed_rows: usize = unindexed
3444 .iter()
3445 .map(|fragment| fragment.num_rows().unwrap_or(0))
3446 .sum();
3447 if indexable_only
3455 && unindexed_rows > 0
3456 && matches!(
3457 intent.params,
3458 IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. }
3459 )
3460 {
3461 unindexed_rows = column_value_count(dataset, intent.column, &unindexed).await?;
3462 }
3463 statuses.push(IndexStatus {
3464 table,
3465 intent_name: intent.name.to_owned(),
3466 fragments_covered: total_fragments.saturating_sub(unindexed_fragments),
3467 unindexed_fragments,
3468 unindexed_rows,
3469 exists,
3470 });
3471 }
3472 Ok(statuses)
3473}
3474
3475pub mod io_trace {
3493 use lance_io::utils::tracking_store::{IOTracker, IoStats};
3494 use std::sync::{Arc, OnceLock};
3495
3496 static TRACKER: OnceLock<IOTracker> = OnceLock::new();
3497
3498 pub fn enable() {
3501 let _ = TRACKER.set(IOTracker::default());
3502 }
3503
3504 pub(super) fn wrapper() -> Option<Arc<IOTracker>> {
3506 TRACKER.get().map(|tracker| Arc::new(tracker.clone()))
3507 }
3508
3509 pub fn take() -> Option<IoStats> {
3511 TRACKER.get().map(IOTracker::incremental_stats)
3512 }
3513}
3514
3515pub mod index_cache {
3523 use object_store::local::LocalFileSystem;
3524 use object_store::path::Path as ObjPath;
3525 use object_store::{
3526 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
3527 ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result as OsResult,
3528 };
3529 use std::collections::HashMap;
3530 use std::ops::Range;
3531 use std::path::PathBuf;
3532 use std::sync::{Arc, Mutex};
3533
3534 use bytes::Bytes;
3535 use futures::stream::BoxStream;
3536 use lance_io::object_store::WrappingObjectStore;
3537
3538 fn is_index_path(location: &ObjPath) -> bool {
3539 AsRef::<str>::as_ref(location).contains("_indices/")
3540 }
3541
3542 fn local_opts(options: &GetOptions) -> GetOptions {
3545 GetOptions {
3546 range: options.range.clone(),
3547 head: options.head,
3548 ..Default::default()
3549 }
3550 }
3551
3552 #[derive(Debug)]
3555 pub struct IndexDiskCache {
3556 local: Arc<LocalFileSystem>,
3557 inflight: Arc<Mutex<HashMap<ObjPath, Arc<tokio::sync::Mutex<()>>>>>,
3558 }
3559
3560 impl IndexDiskCache {
3561 pub fn new(root: PathBuf) -> std::io::Result<Self> {
3563 std::fs::create_dir_all(&root)?;
3564 Ok(Self {
3565 local: Arc::new(LocalFileSystem::new_with_prefix(&root)?),
3566 inflight: Arc::new(Mutex::new(HashMap::new())),
3567 })
3568 }
3569 }
3570
3571 impl WrappingObjectStore for IndexDiskCache {
3572 fn wrap(&self, _store_prefix: &str, inner: Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
3573 Arc::new(CachingStore {
3574 inner,
3575 local: self.local.clone(),
3576 inflight: self.inflight.clone(),
3577 })
3578 }
3579 }
3580
3581 #[derive(Debug)]
3582 struct CachingStore {
3583 inner: Arc<dyn ObjectStore>,
3584 local: Arc<LocalFileSystem>,
3585 inflight: Arc<Mutex<HashMap<ObjPath, Arc<tokio::sync::Mutex<()>>>>>,
3586 }
3587
3588 impl std::fmt::Display for CachingStore {
3589 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3590 write!(f, "CachingStore({})", self.inner)
3591 }
3592 }
3593
3594 impl CachingStore {
3595 fn flight_lock(&self, location: &ObjPath) -> Arc<tokio::sync::Mutex<()>> {
3596 self.inflight
3597 .lock()
3598 .unwrap_or_else(|poison| poison.into_inner())
3599 .entry(location.clone())
3600 .or_default()
3601 .clone()
3602 }
3603
3604 async fn populate_and_serve(
3610 &self,
3611 location: &ObjPath,
3612 options: GetOptions,
3613 ) -> OsResult<GetResult> {
3614 let lock = self.flight_lock(location);
3615 let _guard = lock.lock().await;
3616 let result = self.fetch_under_flight(location, options).await;
3617 self.inflight
3622 .lock()
3623 .unwrap_or_else(|p| p.into_inner())
3624 .remove(location);
3625 result
3626 }
3627
3628 async fn fetch_under_flight(
3629 &self,
3630 location: &ObjPath,
3631 options: GetOptions,
3632 ) -> OsResult<GetResult> {
3633 if let Ok(result) = self.local.get_opts(location, local_opts(&options)).await {
3634 return Ok(result);
3635 }
3636 let bytes = self.inner.get(location).await?.bytes().await?;
3637 if self
3638 .local
3639 .put(location, PutPayload::from_bytes(bytes))
3640 .await
3641 .is_ok()
3642 && let Ok(result) = self.local.get_opts(location, local_opts(&options)).await
3643 {
3644 return Ok(result);
3645 }
3646 self.inner.get_opts(location, options).await
3648 }
3649 }
3650
3651 #[async_trait::async_trait]
3652 impl ObjectStore for CachingStore {
3653 async fn get_opts(&self, location: &ObjPath, options: GetOptions) -> OsResult<GetResult> {
3654 if !is_index_path(location) {
3655 return self.inner.get_opts(location, options).await;
3656 }
3657 match self.local.get_opts(location, local_opts(&options)).await {
3658 Ok(result) => Ok(result),
3659 Err(object_store::Error::NotFound { .. }) => {
3660 self.populate_and_serve(location, options).await
3661 }
3662 Err(_) => self.inner.get_opts(location, options).await,
3663 }
3664 }
3665
3666 async fn put_opts(
3667 &self,
3668 location: &ObjPath,
3669 payload: PutPayload,
3670 opts: PutOptions,
3671 ) -> OsResult<PutResult> {
3672 self.inner.put_opts(location, payload, opts).await
3673 }
3674
3675 async fn put_multipart_opts(
3676 &self,
3677 location: &ObjPath,
3678 opts: PutMultipartOptions,
3679 ) -> OsResult<Box<dyn MultipartUpload>> {
3680 self.inner.put_multipart_opts(location, opts).await
3681 }
3682
3683 async fn get_ranges(
3684 &self,
3685 location: &ObjPath,
3686 ranges: &[Range<u64>],
3687 ) -> OsResult<Vec<Bytes>> {
3688 if is_index_path(location) {
3689 let mut out = Vec::with_capacity(ranges.len());
3691 for range in ranges {
3692 let opts = GetOptions {
3693 range: Some(range.clone().into()),
3694 ..Default::default()
3695 };
3696 out.push(self.get_opts(location, opts).await?.bytes().await?);
3697 }
3698 return Ok(out);
3699 }
3700 self.inner.get_ranges(location, ranges).await
3701 }
3702
3703 fn delete_stream(
3704 &self,
3705 locations: BoxStream<'static, OsResult<ObjPath>>,
3706 ) -> BoxStream<'static, OsResult<ObjPath>> {
3707 self.inner.delete_stream(locations)
3708 }
3709
3710 fn list(&self, prefix: Option<&ObjPath>) -> BoxStream<'static, OsResult<ObjectMeta>> {
3711 self.inner.list(prefix)
3712 }
3713
3714 fn list_with_offset(
3715 &self,
3716 prefix: Option<&ObjPath>,
3717 offset: &ObjPath,
3718 ) -> BoxStream<'static, OsResult<ObjectMeta>> {
3719 self.inner.list_with_offset(prefix, offset)
3720 }
3721
3722 async fn list_with_delimiter(&self, prefix: Option<&ObjPath>) -> OsResult<ListResult> {
3723 self.inner.list_with_delimiter(prefix).await
3724 }
3725
3726 async fn copy_opts(&self, from: &ObjPath, to: &ObjPath, opts: CopyOptions) -> OsResult<()> {
3727 self.inner.copy_opts(from, to, opts).await
3728 }
3729 }
3730
3731 #[cfg(test)]
3732 mod tests {
3733 #![allow(clippy::unwrap_used)]
3734 use super::*;
3735 use object_store::memory::InMemory;
3736
3737 async fn read(store: &Arc<dyn ObjectStore>, path: &ObjPath) -> Option<Vec<u8>> {
3738 store
3739 .get(path)
3740 .await
3741 .ok()?
3742 .bytes()
3743 .await
3744 .ok()
3745 .map(|b| b.to_vec())
3746 }
3747
3748 #[tokio::test]
3749 async fn caches_index_files_and_passes_data_through() {
3750 let temp = tempfile::tempdir().unwrap();
3751 let inner: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
3752 let index_path = ObjPath::from("d/messages.lance/_indices/uuid1/index.idx");
3753 let data_path = ObjPath::from("d/messages.lance/data/x.lance");
3754 inner
3755 .put(&index_path, PutPayload::from_static(b"INDEX"))
3756 .await
3757 .unwrap();
3758 inner
3759 .put(&data_path, PutPayload::from_static(b"DATA"))
3760 .await
3761 .unwrap();
3762
3763 let cache = IndexDiskCache::new(temp.path().join("indices")).unwrap();
3764 let store = cache.wrap("test", inner.clone());
3765
3766 assert_eq!(
3767 read(&store, &index_path).await.as_deref(),
3768 Some(&b"INDEX"[..])
3769 );
3770 assert_eq!(
3771 read(&store, &data_path).await.as_deref(),
3772 Some(&b"DATA"[..])
3773 );
3774
3775 inner.delete(&index_path).await.unwrap();
3778 inner.delete(&data_path).await.unwrap();
3779 assert_eq!(
3780 read(&store, &index_path).await.as_deref(),
3781 Some(&b"INDEX"[..])
3782 );
3783 assert_eq!(read(&store, &data_path).await, None);
3784
3785 let slice = store.get_range(&index_path, 1..4).await.unwrap();
3787 assert_eq!(slice.as_ref(), b"NDE");
3788 }
3789 }
3790}
3791
3792#[cfg(unix)]
3801pub mod durability {
3802 use std::fs::File;
3803 use std::io::ErrorKind;
3804 use std::ops::Range;
3805 use std::path::Path as FsPath;
3806 use std::sync::Arc;
3807
3808 use bytes::Bytes;
3809 use futures::stream::BoxStream;
3810 use lance_io::object_store::WrappingObjectStore;
3811 use object_store::path::Path as ObjPath;
3812 use object_store::{
3813 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
3814 PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OsResult,
3815 UploadPart,
3816 };
3817
3818 fn durability_error(op: &str, path: &FsPath, source: std::io::Error) -> object_store::Error {
3819 object_store::Error::Generic {
3820 store: "fsync-durability",
3821 source: format!("{op} {}: {source}", path.display()).into(),
3822 }
3823 }
3824
3825 fn sync_file_and_parent(location: &ObjPath) -> OsResult<()> {
3831 let local = lance_io::local::to_local_path(location);
3832 let path = FsPath::new(&local);
3833 let file = File::open(path).map_err(|e| durability_error("open for fsync", path, e))?;
3834 file.sync_all()
3835 .map_err(|e| durability_error("fsync", path, e))?;
3836 if let Some(parent) = path.parent() {
3837 match File::open(parent) {
3839 Ok(dir) => dir
3840 .sync_all()
3841 .map_err(|e| durability_error("fsync dir", parent, e))?,
3842 Err(e) if e.kind() == ErrorKind::NotFound => {}
3843 Err(e) => return Err(durability_error("open dir for fsync", parent, e)),
3844 }
3845 }
3846 Ok(())
3847 }
3848
3849 #[derive(Debug)]
3852 pub struct FsyncOnWrite;
3853
3854 impl WrappingObjectStore for FsyncOnWrite {
3855 fn wrap(&self, _store_prefix: &str, inner: Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
3856 Arc::new(FsyncStore { inner })
3857 }
3858 }
3859
3860 #[derive(Debug)]
3861 struct FsyncStore {
3862 inner: Arc<dyn ObjectStore>,
3863 }
3864
3865 impl std::fmt::Display for FsyncStore {
3866 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3867 write!(f, "FsyncStore({})", self.inner)
3868 }
3869 }
3870
3871 #[async_trait::async_trait]
3872 impl ObjectStore for FsyncStore {
3873 async fn put_opts(
3874 &self,
3875 location: &ObjPath,
3876 payload: PutPayload,
3877 opts: PutOptions,
3878 ) -> OsResult<PutResult> {
3879 let result = self.inner.put_opts(location, payload, opts).await?;
3880 sync_file_and_parent(location)?;
3881 Ok(result)
3882 }
3883
3884 async fn put_multipart_opts(
3885 &self,
3886 location: &ObjPath,
3887 opts: PutMultipartOptions,
3888 ) -> OsResult<Box<dyn MultipartUpload>> {
3889 let upload = self.inner.put_multipart_opts(location, opts).await?;
3890 Ok(Box::new(FsyncUpload {
3891 inner: upload,
3892 location: location.clone(),
3893 }))
3894 }
3895
3896 async fn get_opts(&self, location: &ObjPath, options: GetOptions) -> OsResult<GetResult> {
3897 self.inner.get_opts(location, options).await
3898 }
3899
3900 async fn get_ranges(
3901 &self,
3902 location: &ObjPath,
3903 ranges: &[Range<u64>],
3904 ) -> OsResult<Vec<Bytes>> {
3905 self.inner.get_ranges(location, ranges).await
3906 }
3907
3908 fn delete_stream(
3909 &self,
3910 locations: BoxStream<'static, OsResult<ObjPath>>,
3911 ) -> BoxStream<'static, OsResult<ObjPath>> {
3912 self.inner.delete_stream(locations)
3913 }
3914
3915 fn list(&self, prefix: Option<&ObjPath>) -> BoxStream<'static, OsResult<ObjectMeta>> {
3916 self.inner.list(prefix)
3917 }
3918
3919 fn list_with_offset(
3920 &self,
3921 prefix: Option<&ObjPath>,
3922 offset: &ObjPath,
3923 ) -> BoxStream<'static, OsResult<ObjectMeta>> {
3924 self.inner.list_with_offset(prefix, offset)
3925 }
3926
3927 async fn list_with_delimiter(&self, prefix: Option<&ObjPath>) -> OsResult<ListResult> {
3928 self.inner.list_with_delimiter(prefix).await
3929 }
3930
3931 async fn copy_opts(&self, from: &ObjPath, to: &ObjPath, opts: CopyOptions) -> OsResult<()> {
3932 self.inner.copy_opts(from, to, opts).await?;
3933 sync_file_and_parent(to)?;
3934 Ok(())
3935 }
3936
3937 async fn rename_opts(
3941 &self,
3942 from: &ObjPath,
3943 to: &ObjPath,
3944 opts: RenameOptions,
3945 ) -> OsResult<()> {
3946 self.inner.rename_opts(from, to, opts).await?;
3947 sync_file_and_parent(to)?;
3948 Ok(())
3949 }
3950 }
3951
3952 #[derive(Debug)]
3956 struct FsyncUpload {
3957 inner: Box<dyn MultipartUpload>,
3958 location: ObjPath,
3959 }
3960
3961 #[async_trait::async_trait]
3962 impl MultipartUpload for FsyncUpload {
3963 fn put_part(&mut self, data: PutPayload) -> UploadPart {
3964 self.inner.put_part(data)
3965 }
3966
3967 async fn complete(&mut self) -> OsResult<PutResult> {
3968 let result = self.inner.complete().await?;
3969 sync_file_and_parent(&self.location)?;
3970 Ok(result)
3971 }
3972
3973 async fn abort(&mut self) -> OsResult<()> {
3974 self.inner.abort().await
3975 }
3976 }
3977
3978 #[cfg(test)]
3979 mod tests {
3980 #![allow(clippy::unwrap_used)]
3981 use super::*;
3982 use object_store::ObjectStoreExt;
3983
3984 fn wrapped() -> Arc<dyn ObjectStore> {
3987 let inner: Arc<dyn ObjectStore> = Arc::new(object_store::local::LocalFileSystem::new());
3988 FsyncOnWrite.wrap("test", inner)
3989 }
3990
3991 fn obj_path(root: &FsPath, name: &str) -> ObjPath {
3992 ObjPath::from(root.join(name).to_string_lossy().trim_start_matches('/'))
3994 }
3995
3996 #[tokio::test]
3997 async fn put_through_wrapper_round_trips_and_lands_on_disk() {
3998 let temp = tempfile::tempdir().unwrap();
3999 let store = wrapped();
4000 let path = obj_path(temp.path(), "sub/dir/manifest");
4001 store
4002 .put(&path, PutPayload::from_static(b"DURABLE"))
4003 .await
4004 .unwrap();
4005 let got = store.get(&path).await.unwrap().bytes().await.unwrap();
4007 assert_eq!(got.as_ref(), b"DURABLE");
4008 assert_eq!(
4010 std::fs::read(temp.path().join("sub/dir/manifest")).unwrap(),
4011 b"DURABLE",
4012 );
4013 }
4014
4015 #[tokio::test]
4016 async fn multipart_through_wrapper_completes_and_round_trips() {
4017 let temp = tempfile::tempdir().unwrap();
4018 let store = wrapped();
4019 let path = obj_path(temp.path(), "data/part.lance");
4020 let mut upload = store.put_multipart(&path).await.unwrap();
4021 upload
4022 .put_part(PutPayload::from_static(b"AB"))
4023 .await
4024 .unwrap();
4025 upload
4026 .put_part(PutPayload::from_static(b"CD"))
4027 .await
4028 .unwrap();
4029 upload.complete().await.unwrap();
4030 let got = store.get(&path).await.unwrap().bytes().await.unwrap();
4031 assert_eq!(got.as_ref(), b"ABCD");
4032 }
4033 }
4034}
4035
4036pub fn store_key(location: &Url) -> String {
4041 blake3::hash(location.as_str().as_bytes()).to_hex()[..16].to_owned()
4042}
4043
4044fn prune_stale_uuid_dirs(dir: &std::path::Path, keep: &std::collections::HashSet<String>) {
4048 let Ok(entries) = std::fs::read_dir(dir) else {
4049 return;
4050 };
4051 for entry in entries.flatten() {
4052 let path = entry.path();
4053 if !path.is_dir() {
4054 continue;
4055 }
4056 if entry.file_name() == "_indices" {
4057 let Ok(children) = std::fs::read_dir(&path) else {
4058 continue;
4059 };
4060 for child in children.flatten() {
4061 if child.path().is_dir()
4062 && !keep.contains(child.file_name().to_string_lossy().as_ref())
4063 {
4064 let _ = std::fs::remove_dir_all(child.path());
4065 }
4066 }
4067 } else {
4068 prune_stale_uuid_dirs(&path, keep);
4069 }
4070 }
4071}
4072
4073fn store_wrapper(
4079 location: &Url,
4080 index_cache_dir: Option<&std::path::Path>,
4081) -> Option<Arc<dyn WrappingObjectStore>> {
4082 let mut wrappers: Vec<Arc<dyn WrappingObjectStore>> = Vec::new();
4083 #[cfg(unix)]
4087 if config::is_local(location) {
4088 wrappers.push(Arc::new(durability::FsyncOnWrite));
4089 }
4090 if let Some(dir) = index_cache_dir
4091 && !config::is_local(location)
4092 {
4093 let root = dir.join(store_key(location)).join("indices");
4094 match index_cache::IndexDiskCache::new(root) {
4095 Ok(cache) => wrappers.push(Arc::new(cache)),
4096 Err(error) => tracing::warn!(%error, "index disk cache disabled; reads hit the store"),
4097 }
4098 }
4099 if let Some(tracker) = io_trace::wrapper() {
4100 wrappers.push(tracker);
4101 }
4102 match wrappers.len() {
4103 0 => None,
4104 1 => Some(wrappers.remove(0)),
4105 _ => Some(Arc::new(ChainedWrappingObjectStore::new(wrappers))),
4106 }
4107}
4108
4109async fn open_or_create_via_ns(
4110 nm: &Arc<dyn LanceNamespace>,
4111 nm_ident: &NamespaceIdent,
4112 table_name: &str,
4113 schema: lance::deps::arrow_schema::SchemaRef,
4114 session: &Arc<Session>,
4115 storage_options: &HashMap<String, String>,
4116 wrapper: Option<Arc<dyn WrappingObjectStore>>,
4117) -> Result<Dataset> {
4118 let table_id = nm_ident.as_table_id(table_name);
4119
4120 let request = DescribeTableRequest {
4121 id: Some(table_id.clone()),
4122 ..Default::default()
4123 };
4124 match nm.describe_table(request).await {
4125 Ok(response) => {
4126 let location = response.location.with_context(|| {
4127 format!("namespace returned no location for table {table_name}")
4128 })?;
4129 let builder = apply_open_params(
4130 DatasetBuilder::from_uri(&location).with_session(session.clone()),
4131 &wrapper,
4132 storage_options,
4133 );
4134 let mut dataset = match builder.load().await {
4135 Ok(dataset) => dataset,
4136 Err(load_error) => {
4137 let load_error = anyhow::Error::new(load_error)
4138 .context(format!("failed to open table {table_name}"));
4139 match config::local_path(&uri_to_url(&location)?) {
4144 Some(table_root) => {
4145 heal_local_dataset(
4146 &location,
4147 &table_root,
4148 table_name,
4149 session,
4150 storage_options,
4151 &wrapper,
4152 load_error,
4153 )
4154 .await?
4155 }
4156 None => return Err(load_error),
4157 }
4158 }
4159 };
4160 ensure_current_schema(&mut dataset, schema.as_ref(), table_name).await?;
4161 return Ok(dataset);
4162 }
4163 Err(error) => match &error {
4164 error if is_namespace_error_code(error, ErrorCode::TableNotFound) => {
4165 }
4167 _ => {
4168 return Err(anyhow::Error::from(error))
4169 .with_context(|| format!("failed to describe table {table_name}"));
4170 }
4171 },
4172 }
4173
4174 let mut write_params = sessions::write_params_for_create();
4177 write_params.session = Some(session.clone());
4178 write_params.mode = WriteMode::Create;
4179 if wrapper.is_some() || !storage_options.is_empty() {
4183 write_params.store_params = Some(ObjectStoreParams {
4184 object_store_wrapper: wrapper.clone(),
4185 storage_options_accessor: (!storage_options.is_empty()).then(|| {
4186 Arc::new(StorageOptionsAccessor::with_static_options(
4187 storage_options.clone(),
4188 ))
4189 }),
4190 ..Default::default()
4191 });
4192 }
4193 let reader = sessions::empty_reader(schema)?;
4194 Dataset::write_into_namespace(reader, nm.clone(), table_id, Some(write_params))
4195 .await
4196 .with_context(|| format!("failed to create table {table_name}"))
4197}
4198
4199fn apply_open_params(
4203 builder: DatasetBuilder,
4204 wrapper: &Option<Arc<dyn WrappingObjectStore>>,
4205 storage_options: &HashMap<String, String>,
4206) -> DatasetBuilder {
4207 match wrapper {
4208 Some(wrapper) => builder.with_store_params(ObjectStoreParams {
4209 object_store_wrapper: Some(wrapper.clone()),
4210 storage_options_accessor: (!storage_options.is_empty()).then(|| {
4211 Arc::new(StorageOptionsAccessor::with_static_options(
4212 storage_options.clone(),
4213 ))
4214 }),
4215 ..Default::default()
4216 }),
4217 None if !storage_options.is_empty() => {
4218 builder.with_storage_options(storage_options.clone())
4219 }
4220 None => builder,
4221 }
4222}
4223
4224const VERSIONS_DIR_NAME: &str = "_versions";
4226const HEAL_MAX_PROBES: usize = 32;
4230
4231fn parse_manifest_version(filename: &str) -> Option<u64> {
4237 if filename.starts_with('d') {
4238 return None;
4239 }
4240 let stem = filename.strip_suffix(".manifest")?;
4241 if stem.len() == 20 {
4242 stem.parse::<u64>().ok().map(|inverted| u64::MAX - inverted)
4243 } else {
4244 stem.parse::<u64>().ok()
4245 }
4246}
4247
4248async fn scan_verify_version(
4257 table_uri: &str,
4258 version: u64,
4259 session: &Arc<Session>,
4260 storage_options: &HashMap<String, String>,
4261 wrapper: &Option<Arc<dyn WrappingObjectStore>>,
4262) -> Result<()> {
4263 let builder = apply_open_params(
4264 DatasetBuilder::from_uri(table_uri)
4265 .with_session(session.clone())
4266 .with_version(version),
4267 wrapper,
4268 storage_options,
4269 );
4270 let dataset = builder.load().await?;
4271 let scanner = dataset.scan();
4272 let mut stream = scanner.try_into_stream().await?;
4273 while let Some(batch) = stream.next().await {
4274 batch?;
4275 }
4276 Ok(())
4277}
4278
4279async fn heal_local_dataset(
4286 table_uri: &str,
4287 table_root: &std::path::Path,
4288 table_name: &str,
4289 session: &Arc<Session>,
4290 storage_options: &HashMap<String, String>,
4291 wrapper: &Option<Arc<dyn WrappingObjectStore>>,
4292 load_error: anyhow::Error,
4293) -> Result<Dataset> {
4294 let versions_dir = table_root.join(VERSIONS_DIR_NAME);
4295 let entries = match std::fs::read_dir(&versions_dir) {
4296 Ok(entries) => entries,
4297 Err(_) => {
4298 return Err(enriched_open_error(
4299 table_name,
4300 format!(
4301 "open failed and no {} directory exists at {} - not a crash-damaged manifest",
4302 VERSIONS_DIR_NAME,
4303 versions_dir.display()
4304 ),
4305 load_error,
4306 ));
4307 }
4308 };
4309 let mut manifests: Vec<(u64, PathBuf)> = Vec::new();
4310 for entry in entries {
4311 let entry = entry.with_context(|| format!("listing {}", versions_dir.display()))?;
4312 if let Some(version) = parse_manifest_version(&entry.file_name().to_string_lossy()) {
4313 manifests.push((version, entry.path()));
4314 }
4315 }
4316 manifests.sort_by_key(|(version, _)| std::cmp::Reverse(*version));
4317 let Some((_, newest_path)) = manifests.first().cloned() else {
4318 return Err(enriched_open_error(
4319 table_name,
4320 format!(
4321 "open failed and no manifest files exist under {} - not a crash-damaged manifest",
4322 versions_dir.display()
4323 ),
4324 load_error,
4325 ));
4326 };
4327 let newest_desc = || {
4328 let name = newest_path
4329 .file_name()
4330 .map(|n| n.to_string_lossy().into_owned())
4331 .unwrap_or_default();
4332 let bytes = std::fs::metadata(&newest_path).map(|m| m.len()).ok();
4333 match bytes {
4334 Some(bytes) => format!("newest manifest {name} ({bytes} bytes)"),
4335 None => format!("newest manifest {name}"),
4336 }
4337 };
4338
4339 let mut doomed: Vec<PathBuf> = Vec::new();
4343 let mut landed: Option<u64> = None;
4344 let mut probes = 0usize;
4345 for (version, path) in &manifests {
4346 let len = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
4347 if len < 16 {
4350 doomed.push(path.clone());
4351 continue;
4352 }
4353 if probes >= HEAL_MAX_PROBES {
4354 break;
4355 }
4356 probes += 1;
4357 match scan_verify_version(table_uri, *version, session, storage_options, wrapper).await {
4358 Ok(()) => {
4359 landed = Some(*version);
4360 break;
4361 }
4362 Err(_) => doomed.push(path.clone()),
4363 }
4364 }
4365
4366 let Some(landed_version) = landed else {
4368 return Err(enriched_open_error(
4369 table_name,
4370 format!(
4371 "{} is unreadable (interrupted commit during a hard host stop) and no older version passed a scan-verify probe; nothing was quarantined",
4372 newest_desc()
4373 ),
4374 load_error,
4375 ));
4376 };
4377 if doomed.is_empty() {
4379 return Err(enriched_open_error(
4380 table_name,
4381 format!(
4382 "open failed but the manifest head under {} is readable - not a crash-damaged manifest",
4383 versions_dir.display()
4384 ),
4385 load_error,
4386 ));
4387 }
4388
4389 let mut quarantined: Vec<String> = Vec::new();
4393 for path in &doomed {
4394 let mut corrupt = path.clone().into_os_string();
4395 corrupt.push(".corrupt");
4396 match std::fs::rename(path, &corrupt) {
4397 Ok(()) => {
4398 if let Some(name) = path.file_name() {
4399 quarantined.push(name.to_string_lossy().into_owned());
4400 }
4401 }
4402 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
4403 Err(error) => {
4404 return Err(anyhow::Error::new(error).context(format!(
4405 "table {table_name}: failed to quarantine corrupt manifest {}",
4406 path.display()
4407 )));
4408 }
4409 }
4410 }
4411
4412 let builder = apply_open_params(
4414 DatasetBuilder::from_uri(table_uri).with_session(session.clone()),
4415 wrapper,
4416 storage_options,
4417 );
4418 let dataset = builder.load().await.with_context(|| {
4419 format!(
4420 "table {table_name}: open still failed after quarantining {} corrupt manifest(s); restore from a `pond copy` replica or re-run `pond init`",
4421 quarantined.len()
4422 )
4423 })?;
4424
4425 tracing::warn!(
4428 "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.",
4429 quarantined.len(),
4430 quarantined.join(", "),
4431 );
4432
4433 Ok(dataset)
4434}
4435
4436fn enriched_open_error(
4440 table_name: &str,
4441 finding: String,
4442 load_error: anyhow::Error,
4443) -> anyhow::Error {
4444 load_error.context(format!(
4445 "table {table_name}: {finding}. Restore this store from a `pond copy` replica or re-run `pond init` to re-sync from source histories"
4446 ))
4447}
4448
4449fn is_namespace_error_code(error: &lance::Error, code: ErrorCode) -> bool {
4453 if !matches!(error, lance::Error::Namespace { .. }) {
4454 return false;
4455 }
4456 std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |link| {
4457 link.source()
4458 })
4459 .filter_map(|link| link.downcast_ref::<NamespaceError>())
4460 .any(|inner| inner.code() == code)
4461}
4462
4463fn scanner_with_prefilter(
4464 dataset: &Dataset,
4465 predicate: Option<&Predicate>,
4466) -> Result<lance::dataset::scanner::Scanner> {
4467 let mut scanner = dataset.scan();
4468 scanner.prefilter(true);
4469 if let Some(predicate) = predicate {
4470 let filter = predicate.to_lance();
4471 if !filter.is_empty() {
4472 scanner.filter(&filter)?;
4473 }
4474 }
4475 Ok(scanner)
4476}
4477enum SchemaFit {
4479 Match,
4480 MissingNullable(Vec<lance::deps::arrow_schema::Field>),
4483 UnknownExtra(Vec<String>),
4488}
4489
4490fn classify_schema(
4491 actual: &lance::deps::arrow_schema::Schema,
4492 expected: &lance::deps::arrow_schema::Schema,
4493 table_name: &str,
4494) -> Result<SchemaFit> {
4495 use std::collections::BTreeSet;
4496 let actual_names: BTreeSet<&str> = actual.fields().iter().map(|f| f.name().as_str()).collect();
4497 let expected_names: BTreeSet<&str> = expected
4498 .fields()
4499 .iter()
4500 .map(|f| f.name().as_str())
4501 .collect();
4502 let missing: Vec<_> = expected
4503 .fields()
4504 .iter()
4505 .filter(|f| !actual_names.contains(f.name().as_str()))
4506 .map(|f| f.as_ref().clone())
4507 .collect();
4508 let extra: Vec<String> = actual_names
4509 .difference(&expected_names)
4510 .map(|name| (*name).to_owned())
4511 .collect();
4512 match (missing.is_empty(), extra.is_empty()) {
4513 (true, true) => Ok(SchemaFit::Match),
4514 (false, true) if missing.iter().all(|f| f.is_nullable()) => {
4515 Ok(SchemaFit::MissingNullable(missing))
4516 }
4517 (true, false) => Ok(SchemaFit::UnknownExtra(extra)),
4518 _ => anyhow::bail!(
4519 "table {table_name} has columns {actual_names:?} but this pond build expects \
4520 {expected_names:?}, and the difference is not an additive nullable-column \
4521 change this build can migrate - upgrade pond, or restore the store from a \
4522 `pond copy` snapshot taken by the version that wrote it",
4523 ),
4524 }
4525}
4526
4527async fn ensure_current_schema(
4535 dataset: &mut Dataset,
4536 expected: &lance::deps::arrow_schema::Schema,
4537 table_name: &str,
4538) -> Result<()> {
4539 use lance::deps::arrow_schema::DataType;
4540 const MAX_MIGRATION_ATTEMPTS: usize = 3;
4541 for _ in 0..MAX_MIGRATION_ATTEMPTS {
4542 let actual = lance::deps::arrow_schema::Schema::from(dataset.schema());
4543 match classify_schema(&actual, expected, table_name)? {
4544 SchemaFit::MissingNullable(missing) => {
4545 backfill_missing_columns(dataset, table_name, missing).await?;
4546 continue;
4547 }
4548 SchemaFit::Match => {}
4549 SchemaFit::UnknownExtra(extra) => {
4550 tracing::warn!(
4551 table = table_name,
4552 ?extra,
4553 "store carries columns unknown to this pond build (written by a newer \
4554 version); reads proceed, writes need the newer pond",
4555 );
4556 }
4557 }
4558 for actual_field in actual.fields() {
4563 let Some(expected_field) = expected.field_with_name(actual_field.name()).ok() else {
4564 continue;
4565 };
4566 if let (
4567 DataType::FixedSizeList(_, actual_dim),
4568 DataType::FixedSizeList(_, expected_dim),
4569 ) = (actual_field.data_type(), expected_field.data_type())
4570 && actual_dim != expected_dim
4571 {
4572 tracing::warn!(
4573 table = table_name,
4574 column = actual_field.name(),
4575 actual_dim,
4576 expected_dim,
4577 "embedding dimension differs from config; open proceeds because model swaps are operator-driven",
4578 );
4579 }
4580 }
4581 return Ok(());
4582 }
4583 anyhow::bail!(
4584 "schema migration for table {table_name} did not converge after \
4585 {MAX_MIGRATION_ATTEMPTS} attempts (a concurrent writer kept changing \
4586 the schema); re-run once the other pond process finishes",
4587 )
4588}
4589
4590async fn backfill_missing_columns(
4595 dataset: &mut Dataset,
4596 table_name: &str,
4597 missing: Vec<lance::deps::arrow_schema::Field>,
4598) -> Result<()> {
4599 use lance::dataset::{BatchUDF, NewColumnTransform};
4600 let names: Vec<&str> = missing.iter().map(|f| f.name().as_str()).collect();
4601 let spec = sessions::column_backfill(table_name, &missing)?;
4602 let _ = crate::output::line_err(&format!(
4606 "migrating {table_name}: backfilling {names:?} from stored data (one-time, in place)...",
4607 ));
4608 let started = std::time::Instant::now();
4609 let mapper = spec.mapper;
4610 let migration: std::pin::Pin<
4614 Box<dyn std::future::Future<Output = lance::Result<()>> + Send + '_>,
4615 > = Box::pin(dataset.add_columns(
4616 NewColumnTransform::BatchUDF(BatchUDF {
4617 mapper: Box::new(move |batch| {
4618 mapper(batch).map_err(|error| lance::Error::io(format!("{error:#}")))
4619 }),
4620 output_schema: spec.output_schema,
4621 result_checkpoint: None,
4622 }),
4623 Some(spec.read_columns),
4624 None,
4625 ));
4626 let result = migration.await;
4627 match result {
4628 Ok(()) => {
4629 let _ = crate::output::line_err(&format!(
4630 "migrated {table_name} in {:.1}s",
4631 started.elapsed().as_secs_f64(),
4632 ));
4633 Ok(())
4634 }
4635 Err(error) => {
4636 let error = anyhow::Error::from(error);
4637 if is_commit_conflict(&error) {
4638 dataset.checkout_latest().await?;
4641 Ok(())
4642 } else {
4643 Err(error).with_context(|| {
4644 format!(
4645 "schema backfill failed for {table_name}; the one-time migration \
4646 writes new column files, so it needs write access to the store - \
4647 re-run any pond command with write-capable credentials to complete it",
4648 )
4649 })
4650 }
4651 }
4652 }
4653}
4654fn apply_remote_storage_defaults(options: &mut HashMap<String, String>) {
4661 fn set_default(options: &mut HashMap<String, String>, aliases: &[&str], value: &str) {
4662 if aliases
4663 .iter()
4664 .any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)))
4665 {
4666 return;
4667 }
4668 options.insert(aliases[0].to_owned(), value.to_owned());
4669 }
4670 set_default(options, &["pool_idle_timeout"], "300 seconds");
4671 set_default(options, &["connect_timeout"], "10 seconds");
4672 set_default(options, &["request_timeout"], "60 seconds");
4679 let has_custom_endpoint = ["aws_endpoint", "endpoint"]
4680 .iter()
4681 .any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)));
4682 if has_custom_endpoint {
4683 set_default(
4684 options,
4685 &["aws_unsigned_payload", "unsigned_payload"],
4686 "true",
4687 );
4688 }
4689}
4690
4691fn quoted_string(value: &str) -> String {
4692 format!("'{}'", value.replace('\'', "''"))
4693}
4694fn like_contains(value: &str) -> String {
4695 let escaped = value
4696 .replace('\\', "\\\\")
4697 .replace('%', "\\%")
4698 .replace('_', "\\_")
4699 .replace('\'', "''");
4700 format!("'%{escaped}%'")
4701}
4702
4703#[cfg(test)]
4704mod tests {
4705 #![allow(clippy::expect_used, clippy::unwrap_used)]
4706
4707 use super::*;
4708 use tempfile::TempDir;
4709
4710 #[test]
4711 fn is_index_error_matches_lance_index_class_through_context() {
4712 let index_fault = anyhow::Error::from(lance::Error::index(
4713 "cannot merge inverted index segments with different posting tail codecs",
4714 ))
4715 .context("optimize_indices(merge) failed during index optimize");
4716 assert!(is_index_error(&index_fault));
4717
4718 let io_fault = anyhow::anyhow!("connection reset").context("optimize_indices failed");
4719 assert!(!is_index_error(&io_fault));
4720 }
4721
4722 #[test]
4723 fn prune_keeps_live_uuid_dirs_and_drops_dead_ones() {
4724 let temp = TempDir::new().unwrap();
4725 let indices = temp.path().join("bkt/messages.lance/_indices");
4726 for uuid in ["live", "dead"] {
4727 std::fs::create_dir_all(indices.join(uuid)).unwrap();
4728 std::fs::write(indices.join(uuid).join("index.idx"), b"x").unwrap();
4729 }
4730 let keep = std::collections::HashSet::from(["live".to_owned()]);
4731 prune_stale_uuid_dirs(temp.path(), &keep);
4732 assert!(indices.join("live").exists());
4733 assert!(!indices.join("dead").exists());
4734 }
4735
4736 #[cfg(unix)]
4737 #[test]
4738 fn store_wrapper_present_for_local_absent_for_remote() {
4739 let local = Url::parse("file:///tmp/pond-wrapper-test").unwrap();
4742 assert!(store_wrapper(&local, None).is_some());
4743 for remote in ["memory:///pond-wrapper-test", "s3://bucket/prefix"] {
4744 let url = Url::parse(remote).unwrap();
4745 assert!(
4746 store_wrapper(&url, None).is_none(),
4747 "remote store must not carry the fsync wrapper: {remote}",
4748 );
4749 }
4750 }
4751
4752 fn set(scope: Option<&str>) -> CredsSet {
4753 CredsSet {
4754 scope: scope.map(str::to_owned),
4755 access_key_id: Some("AKIA".to_owned()),
4756 secret_access_key: Some("shh".to_owned()),
4757 ..CredsSet::default()
4758 }
4759 }
4760
4761 fn opts(resolved: &ResolvedStorage, key: &str) -> Option<String> {
4762 resolved.options.get(key).cloned()
4763 }
4764
4765 #[test]
4766 fn storage_url_translation_table() {
4767 let local = StorageUrl::parse("/srv/pond").unwrap();
4770 assert_eq!(local.lance_url().as_str(), "file:///srv/pond/");
4771 assert!(local.is_local());
4772 assert!(local.scheme_options.is_empty());
4773 let aws = StorageUrl::parse("s3://bucket/prefix").unwrap();
4775 assert_eq!(aws.lance_url().as_str(), "s3://bucket/prefix");
4776 assert!(aws.scheme_options.is_empty());
4777 let fat = StorageUrl::parse("s3+https://nbg1.example.com/my-pond/sub").unwrap();
4782 assert_eq!(fat.lance_url().as_str(), "s3://my-pond/sub");
4783 assert_eq!(
4784 fat.scheme_options,
4785 vec![
4786 ("allow_http", "false".to_owned()),
4787 ("virtual_hosted_style_request", "true".to_owned()),
4788 ("region", "us-east-1".to_owned()),
4789 ],
4790 );
4791 let resolved = fat.resolve(&BTreeMap::new()).unwrap();
4792 assert_eq!(
4793 opts(&resolved, "endpoint").as_deref(),
4794 Some("https://my-pond.nbg1.example.com"),
4795 );
4796 assert_eq!(opts(&resolved, "region").as_deref(), Some("us-east-1"));
4797 let plain = StorageUrl::parse("s3+http://127.0.0.1:9000/pond").unwrap();
4800 assert_eq!(plain.lance_url().as_str(), "s3://pond/");
4801 assert_eq!(plain.scheme_options[0], ("allow_http", "true".to_owned()));
4802 assert_eq!(
4803 plain.scheme_options[1],
4804 ("virtual_hosted_style_request", "false".to_owned()),
4805 );
4806 let resolved = plain.resolve(&BTreeMap::new()).unwrap();
4807 assert_eq!(
4808 opts(&resolved, "endpoint").as_deref(),
4809 Some("http://127.0.0.1:9000"),
4810 );
4811 let mut pinned = BTreeMap::new();
4813 pinned.insert(
4814 "default".to_owned(),
4815 CredsSet {
4816 extra: [(
4817 "endpoint".to_owned(),
4818 "https://pinned.example.com".to_owned(),
4819 )]
4820 .into_iter()
4821 .collect(),
4822 ..CredsSet::default()
4823 },
4824 );
4825 let resolved = fat.resolve(&pinned).unwrap();
4826 assert_eq!(
4827 opts(&resolved, "endpoint").as_deref(),
4828 Some("https://pinned.example.com"),
4829 );
4830 let gcs = StorageUrl::parse("gs://bucket/p").unwrap();
4832 assert_eq!(gcs.lance_url().as_str(), "gs://bucket/p");
4833 let azure = StorageUrl::parse("az://acct/container/p").unwrap();
4835 assert_eq!(azure.lance_url().as_str(), "az://container/p");
4836 assert_eq!(
4837 azure.scheme_options,
4838 vec![("account_name", "acct".to_owned())]
4839 );
4840 let shared = StorageUrl::parse("shared-memory://pond-test-x/").unwrap();
4842 assert_eq!(shared.lance_url().as_str(), "shared-memory://pond-test-x/");
4843 }
4844
4845 #[test]
4846 fn storage_url_rejects_bad_shapes() {
4847 let err = StorageUrl::parse("s3+https://user:pass@host/bucket")
4849 .expect_err("userinfo must be rejected")
4850 .to_string();
4851 assert!(
4852 err.contains("creds"),
4853 "error must name the alternative: {err}"
4854 );
4855 assert!(StorageUrl::parse("s3+https://host").is_err());
4857 assert!(StorageUrl::parse("az://acct").is_err());
4858 let err = StorageUrl::parse("ftp://host/x")
4860 .expect_err("ftp")
4861 .to_string();
4862 assert!(err.contains("s3+https"), "got: {err}");
4863 let err = StorageUrl::parse("s3://b/p?regoin=x")
4865 .expect_err("typo")
4866 .to_string();
4867 assert!(err.contains("regoin"), "got: {err}");
4868 let err = StorageUrl::parse("memory://x?creds=y")
4871 .expect_err("memory query")
4872 .to_string();
4873 assert!(err.contains("query params"), "got: {err}");
4874 let err = StorageUrl::parse("file:///x?creds=y")
4875 .expect_err("file query")
4876 .to_string();
4877 assert!(err.contains("query params"), "got: {err}");
4878 assert!(StorageUrl::parse("/tmp/a?b").is_ok());
4880 }
4881
4882 #[test]
4883 fn storage_url_canonicalizes_ports_and_keeps_percent_encoding() {
4884 let with_port = StorageUrl::parse("s3+https://host:443/bucket/p").unwrap();
4886 let without = StorageUrl::parse("s3+https://host/bucket/p").unwrap();
4887 assert_eq!(with_port.canonical(), without.canonical());
4888 let odd = StorageUrl::parse("s3+https://host:8443/bucket").unwrap();
4890 let resolved = odd.resolve(&BTreeMap::new()).unwrap();
4891 assert_eq!(
4892 resolved.options.get("endpoint").map(String::as_str),
4893 Some("https://bucket.host:8443"),
4894 );
4895 let encoded = StorageUrl::parse("s3+https://host/bucket/pre%20fix").unwrap();
4897 assert_eq!(encoded.lance_url().as_str(), "s3://bucket/pre%20fix");
4898 }
4899
4900 #[test]
4901 fn query_params_strip_and_apply_over_set_fields() {
4902 let mut creds = BTreeMap::new();
4903 creds.insert(
4904 "default".to_owned(),
4905 CredsSet {
4906 region: Some("from-set".to_owned()),
4907 virtual_hosted_style_request: Some(false),
4908 ..set(None)
4909 },
4910 );
4911 let url = StorageUrl::parse(
4912 "s3+https://host/bucket/p?region=from-query&virtual_hosted_style_request=true",
4913 )
4914 .unwrap();
4915 assert_eq!(url.lance_url().as_str(), "s3://bucket/p");
4917 assert!(url.canonical().query().is_none());
4918 let resolved = url.resolve(&creds).unwrap();
4919 assert_eq!(opts(&resolved, "region").as_deref(), Some("from-query"));
4921 assert_eq!(
4922 opts(&resolved, "virtual_hosted_style_request").as_deref(),
4923 Some("true"),
4924 );
4925 assert_eq!(
4927 opts(&resolved, "endpoint").as_deref(),
4928 Some("https://bucket.host"),
4929 );
4930 }
4931
4932 #[test]
4933 fn scope_matching_binds_by_longest_prefix_at_segment_boundaries() {
4934 let mut creds = BTreeMap::new();
4935 creds.insert("all".to_owned(), set(None));
4936 creds.insert("bucket".to_owned(), set(Some("s3+https://host/pond/")));
4937 creds.insert("deep".to_owned(), set(Some("s3+https://host/pond/sub")));
4938
4939 let bind = |input: &str| {
4940 StorageUrl::parse(input)
4941 .unwrap()
4942 .resolve(&creds)
4943 .unwrap()
4944 .binding
4945 };
4946 assert_eq!(
4948 bind("s3+https://host/pond/sub/x"),
4949 CredsBinding::Set {
4950 name: "deep".to_owned(),
4951 via: BindVia::Scope
4952 },
4953 );
4954 assert_eq!(
4955 bind("s3+https://host/pond/other"),
4956 CredsBinding::Set {
4957 name: "bucket".to_owned(),
4958 via: BindVia::Scope
4959 },
4960 );
4961 assert_eq!(
4963 bind("s3+https://host/pond-2"),
4964 CredsBinding::Set {
4965 name: "all".to_owned(),
4966 via: BindVia::CatchAll
4967 },
4968 );
4969 assert_eq!(
4971 bind("s3://pond/sub"),
4972 CredsBinding::Set {
4973 name: "all".to_owned(),
4974 via: BindVia::CatchAll
4975 },
4976 );
4977 assert_eq!(
4979 bind("s3+https://host:443/pond/x"),
4980 CredsBinding::Set {
4981 name: "bucket".to_owned(),
4982 via: BindVia::Scope
4983 },
4984 );
4985 assert_eq!(
4987 bind("s3+https://host/pond/sub/x?creds=all"),
4988 CredsBinding::Set {
4989 name: "all".to_owned(),
4990 via: BindVia::Pointer
4991 },
4992 );
4993 let err = StorageUrl::parse("s3://b/p?creds=nope")
4995 .unwrap()
4996 .resolve(&creds)
4997 .expect_err("missing set")
4998 .to_string();
4999 assert!(err.contains("creds=nope"), "got: {err}");
5000
5001 let empty = BTreeMap::new();
5003 assert_eq!(
5004 StorageUrl::parse("s3://b/p")
5005 .unwrap()
5006 .resolve(&empty)
5007 .unwrap()
5008 .binding,
5009 CredsBinding::Ambient,
5010 );
5011 assert_eq!(
5012 StorageUrl::parse("/srv/pond")
5013 .unwrap()
5014 .resolve(&creds)
5015 .unwrap()
5016 .binding,
5017 CredsBinding::NotApplicable,
5018 );
5019 }
5020
5021 #[test]
5022 fn unmatched_sets_are_reported_only_on_remote_invocations() {
5023 let mut creds = BTreeMap::new();
5024 creds.insert("used".to_owned(), set(Some("s3://bucket/")));
5025 creds.insert("idle".to_owned(), set(Some("s3://other/")));
5026
5027 let remote = StorageUrl::parse("s3://bucket/p")
5028 .unwrap()
5029 .resolve(&creds)
5030 .unwrap();
5031 assert_eq!(unmatched_creds_sets(&[&remote], &creds), vec!["idle"]);
5032
5033 let local = StorageUrl::parse("/srv/pond")
5035 .unwrap()
5036 .resolve(&creds)
5037 .unwrap();
5038 assert!(unmatched_creds_sets(&[&local], &creds).is_empty());
5039 }
5040
5041 #[test]
5042 fn secrets_materialize_from_file_and_command() {
5043 let dir = TempDir::new().unwrap();
5044 let key_path = dir.path().join("key");
5045 std::fs::write(&key_path, "from-file\n").unwrap();
5046 let mut creds = BTreeMap::new();
5047 creds.insert(
5048 "default".to_owned(),
5049 CredsSet {
5050 access_key_id_file: Some(key_path),
5051 secret_access_key_command: Some("printf 'from-command\\n\\n'".to_owned()),
5053 ..CredsSet::default()
5054 },
5055 );
5056 let url = StorageUrl::parse("s3://bucket/p").unwrap();
5057 let resolved = url.resolve(&creds).unwrap();
5058 assert_eq!(
5059 opts(&resolved, "access_key_id").as_deref(),
5060 Some("from-file")
5061 );
5062 assert_eq!(
5063 opts(&resolved, "secret_access_key").as_deref(),
5064 Some("from-command\n"),
5065 );
5066
5067 let mut failing = BTreeMap::new();
5069 failing.insert(
5070 "default".to_owned(),
5071 CredsSet {
5072 secret_access_key_command: Some("exit 3".to_owned()),
5073 ..CredsSet::default()
5074 },
5075 );
5076 let err = url
5077 .resolve(&failing)
5078 .expect_err("command must fail")
5079 .to_string();
5080 assert!(err.contains("exit 3"), "got: {err}");
5081
5082 let marker = dir.path().join("runs");
5084 let command = format!("echo run >> {} && echo secret", marker.display());
5085 let mut counted = BTreeMap::new();
5086 counted.insert(
5087 "default".to_owned(),
5088 CredsSet {
5089 secret_access_key_command: Some(command),
5090 ..CredsSet::default()
5091 },
5092 );
5093 url.resolve(&counted).unwrap();
5094 url.resolve(&counted).unwrap();
5095 let runs = std::fs::read_to_string(&marker).unwrap();
5096 assert_eq!(runs.lines().count(), 1, "command must run exactly once");
5097 }
5098
5099 #[test]
5100 fn check_errors_classify_by_kind_and_binding() {
5101 let auth_error = || object_store::Error::Unauthenticated {
5102 path: "k".to_owned(),
5103 source: "denied".into(),
5104 };
5105 let bound = CredsBinding::Set {
5106 name: "work".to_owned(),
5107 via: BindVia::Scope,
5108 };
5109 match classify_check_error(auth_error(), &bound, "put") {
5111 CheckFailure::Auth { set, .. } => assert_eq!(set, "work"),
5112 other => panic!("want Auth, got {other:?}"),
5113 }
5114 assert!(matches!(
5116 classify_check_error(auth_error(), &CredsBinding::Ambient, "put"),
5117 CheckFailure::NoCreds { .. },
5118 ));
5119 let denied = object_store::Error::PermissionDenied {
5120 path: "k".to_owned(),
5121 source: "403".into(),
5122 };
5123 assert!(matches!(
5124 classify_check_error(denied, &bound, "put"),
5125 CheckFailure::Auth { .. },
5126 ));
5127 let missing = object_store::Error::NotFound {
5129 path: "k".to_owned(),
5130 source: "404".into(),
5131 };
5132 assert!(matches!(
5133 classify_check_error(missing, &bound, "get"),
5134 CheckFailure::Io { .. },
5135 ));
5136 let no_creds = || object_store::Error::Generic {
5140 store: "S3",
5141 source: "Failed to get AWS credentials: CredentialsNotLoaded".into(),
5142 };
5143 assert!(matches!(
5144 classify_check_error(no_creds(), &bound, "put"),
5145 CheckFailure::Auth { .. },
5146 ));
5147 assert!(matches!(
5148 classify_check_error(no_creds(), &CredsBinding::Ambient, "put"),
5149 CheckFailure::NoCreds { .. },
5150 ));
5151 }
5152
5153 #[test]
5154 fn concise_cause_strips_upstream_noise_to_one_line() {
5155 let inner = "Encountered internal error. Please file a bug report at \
5158 https://github.com/lance-format/lance/issues. Failed to get AWS \
5159 credentials: CredentialsNotLoaded, <WORKSPACE>/src/object_store/providers/aws.rs:401:21: \
5160 Encountered internal error. Please file a bug report at \
5161 https://github.com/lance-format/lance/issues. Failed to get AWS \
5162 credentials: CredentialsNotLoaded";
5163 let failure = CheckFailure::NoCreds {
5164 source: anyhow!(inner.to_owned()).context("initial conditional put"),
5165 };
5166 let cause = failure.concise_cause().expect("auth-class carries a cause");
5167 assert_eq!(cause, "Failed to get AWS credentials: CredentialsNotLoaded");
5168 assert!(
5170 !failure.to_string().contains("file a bug report"),
5171 "lead must not trail the chain: {failure}"
5172 );
5173 let occ = CheckFailure::OccUnsupported {
5175 detail: "put-if-none-match ignored".to_owned(),
5176 };
5177 assert!(occ.concise_cause().is_none());
5178 let long = CheckFailure::Io {
5181 source: anyhow!(format!("{} dns error: lookup failed", "x".repeat(500))),
5182 };
5183 let cause = long.concise_cause().expect("io carries a cause");
5184 assert!(cause.contains(" ... "), "long causes truncate: {cause}");
5185 assert!(
5186 cause.ends_with("dns error: lookup failed"),
5187 "the tail survives: {cause}"
5188 );
5189 }
5190
5191 #[tokio::test]
5192 async fn storage_check_passes_on_memory_backend() {
5193 let resolved = StorageUrl::parse("memory://check/probe")
5194 .unwrap()
5195 .resolve(&BTreeMap::new())
5196 .unwrap();
5197 storage_check(&resolved).await.expect("memory probe passes");
5198 }
5199
5200 fn fragment(bytes: u64, rows: u64, deleted_rows: u64) -> FragmentStat {
5201 FragmentStat {
5202 bytes: Some(bytes),
5203 rows,
5204 deleted_rows,
5205 }
5206 }
5207
5208 fn stat(bytes: u64) -> FragmentStat {
5209 fragment(bytes, bytes / 1_000, 0)
5210 }
5211
5212 fn task_is_kept(stats: &[FragmentStat], target_rows_per_fragment: usize) -> bool {
5213 task_veto_reason(
5214 stats,
5215 64,
5216 0.1,
5217 target_rows_per_fragment,
5218 TARGET_FRAGMENT_BYTES,
5219 )
5220 .is_none()
5221 }
5222
5223 #[test]
5224 fn compaction_veto_blocks_absorb_keeps_peers() {
5225 let absorb = [stat(100_000_000), stat(1_000_000), stat(2_000_000)];
5227 assert!(!task_is_kept(&absorb, derived_target_rows(&absorb)));
5228 let peers = [stat(100_000_000), stat(100_000_000)];
5230 assert!(task_is_kept(&peers, derived_target_rows(&peers)));
5231 let tiered = [stat(400_000), stat(60_000), stat(40_000)];
5233 assert!(task_is_kept(&tiered, derived_target_rows(&tiered)));
5234 }
5235
5236 #[test]
5237 fn compaction_veto_passes_deletions_and_cap() {
5238 let mut deleting = stat(665_000_000);
5239 deleting.deleted_rows = deleting.rows / 5;
5240 let deleting_task = [deleting, stat(1_000)];
5241 assert!(task_is_kept(
5242 &deleting_task,
5243 derived_target_rows(&deleting_task),
5244 ));
5245
5246 let wide: Vec<FragmentStat> = std::iter::once(stat(100_000_000))
5247 .chain(std::iter::repeat_with(|| stat(100_000)).take(63))
5248 .collect();
5249 assert!(task_is_kept(&wide, derived_target_rows(&wide)));
5250 }
5251
5252 #[test]
5253 fn compaction_veto_fails_closed_on_unknown_sizes() {
5254 let mut unknown = stat(665_000_000);
5255 unknown.bytes = None;
5256 let task = [unknown, stat(665_000_000)];
5257 assert_eq!(
5258 task_veto_reason(
5259 &task,
5260 64,
5261 0.1,
5262 derived_target_rows(&task),
5263 TARGET_FRAGMENT_BYTES
5264 ),
5265 Some("missing_sizes"),
5266 );
5267 }
5268
5269 #[test]
5270 fn compaction_veto_uses_physical_rows_after_deletions() {
5271 let partially_deleted = || fragment(100_000_000, 100_000, 9_000);
5272 let task: Vec<FragmentStat> = std::iter::repeat_with(partially_deleted).take(3).collect();
5273 assert!(task_is_kept(&task, derived_target_rows(&task)));
5274 }
5275
5276 #[test]
5277 fn compaction_filter_keeps_above_budget_off_boundary_task() {
5278 let table = [
5279 fragment(100_000_000, 150_000, 0),
5280 fragment(100_000_000, 150_000, 0),
5281 fragment(100_000_000, 150_000, 0),
5282 fragment(100_000_000, 50_000, 0),
5283 ];
5284 let task = &table[..3];
5285 let target = derived_target_rows(&table);
5286 let total_bytes = task.iter().map(|stat| stat.bytes.unwrap()).sum::<u64>();
5287
5288 assert!(total_bytes > TARGET_FRAGMENT_BYTES);
5289 assert!(target < derived_target_rows(task));
5290 assert!(task_is_kept(task, target));
5291 }
5292
5293 #[test]
5294 fn compaction_filter_keeps_real_mixed_width_tasks_below_budget() {
5295 let four_fragment_task = [
5296 fragment(26_612_870, 9_281, 0),
5297 fragment(14_242_314, 4_400, 0),
5298 fragment(54_111_122, 20_988, 0),
5299 fragment(517_923, 166, 0),
5300 ];
5301 assert!(task_is_kept(&four_fragment_task, 58_468));
5302
5303 let nine_fragment_task = [
5304 fragment(13_547_709, 3_946, 0),
5305 fragment(344_320, 155, 0),
5306 fragment(134_209, 34, 0),
5307 fragment(54_624_719, 15_759, 0),
5308 fragment(1_364_162, 292, 0),
5309 fragment(110_225_801, 32_826, 0),
5310 fragment(8_151_118, 1_840, 0),
5311 fragment(728_590, 79, 0),
5312 fragment(685_184, 128, 0),
5313 ];
5314 assert!(task_is_kept(&nine_fragment_task, 58_468));
5315 }
5316
5317 #[test]
5318 fn compaction_veto_rejects_byte_capped_second_cycle() {
5319 let wide_peer = || fragment(200_000_000, 2_000, 0);
5321 let first_cycle: Vec<FragmentStat> = std::iter::repeat_with(wide_peer).take(5).collect();
5322 let expected_outputs_by_bytes = 1_000_000_000u64.div_ceil(TARGET_FRAGMENT_BYTES) as usize;
5323 assert_eq!(expected_outputs_by_bytes, 4);
5324 assert!(expected_outputs_by_bytes < first_cycle.len());
5325 let cap_sized_task: Vec<FragmentStat> =
5326 std::iter::repeat_with(wide_peer).take(64).collect();
5327
5328 let second_cycle = [
5329 fragment(TARGET_FRAGMENT_BYTES, 2_684, 0),
5330 fragment(TARGET_FRAGMENT_BYTES, 2_684, 0),
5331 fragment(TARGET_FRAGMENT_BYTES, 2_684, 0),
5332 fragment(194_693_632, 1_948, 0),
5333 ];
5334
5335 assert_eq!(
5336 task_veto_reason(&first_cycle, 64, 0.1, 66_000, TARGET_FRAGMENT_BYTES),
5337 Some("row_target_unattainable"),
5338 );
5339 assert!(task_is_kept(&cap_sized_task, 66_000));
5340 assert_eq!(
5341 task_veto_reason(&second_cycle, 64, 0.1, 66_000, TARGET_FRAGMENT_BYTES),
5342 Some("cannot_shrink"),
5343 );
5344 }
5345
5346 #[test]
5347 fn parse_manifest_version_handles_all_naming_schemes() {
5348 assert_eq!(parse_manifest_version("5.manifest"), Some(5));
5350 assert_eq!(parse_manifest_version("0.manifest"), Some(0));
5351 assert_eq!(
5353 parse_manifest_version("18446744073709551615.manifest"),
5354 Some(0)
5355 );
5356 assert_eq!(
5357 parse_manifest_version("18446744073709551610.manifest"),
5358 Some(5)
5359 );
5360 assert_eq!(parse_manifest_version("d123.manifest"), None);
5362 assert_eq!(parse_manifest_version("5.manifest.corrupt"), None);
5364 assert_eq!(
5365 parse_manifest_version(".tmp_7.manifest_9c100374-3298-4537-afc6-f5ee7913666d"),
5366 None
5367 );
5368 assert_eq!(parse_manifest_version("data.lance"), None);
5370 assert_eq!(parse_manifest_version("notanumber.manifest"), None);
5371 }
5372
5373 #[tokio::test]
5374 async fn scan_verify_rejects_zeroed_column_add_data_file() {
5375 let temp = tempfile::tempdir().unwrap();
5379 let uri_owned = temp.path().join("t.lance");
5380 let uri = uri_owned.to_str().unwrap();
5381 let schema = Arc::new(lance::deps::arrow_schema::Schema::new(vec![
5382 lance::deps::arrow_schema::Field::new(
5383 "id",
5384 lance::deps::arrow_schema::DataType::Utf8,
5385 false,
5386 ),
5387 ]));
5388 let batch = RecordBatch::try_new(
5389 schema.clone(),
5390 vec![Arc::new(StringArray::from(vec!["a", "b", "c"]))],
5391 )
5392 .unwrap();
5393 let reader = RecordBatchIterator::new([Ok(batch)], schema);
5394 let mut dataset = Dataset::write(reader, uri, None).await.unwrap();
5395
5396 let data_files = || -> std::collections::BTreeSet<PathBuf> {
5397 std::fs::read_dir(uri_owned.join("data"))
5398 .unwrap()
5399 .map(|entry| entry.unwrap().path())
5400 .collect()
5401 };
5402 let before = data_files();
5403 dataset
5404 .add_columns(
5405 lance::dataset::NewColumnTransform::SqlExpressions(vec![(
5406 "extra".to_string(),
5407 "id".to_string(),
5408 )]),
5409 None,
5410 None,
5411 )
5412 .await
5413 .unwrap();
5414 let column_add_file = data_files()
5415 .difference(&before)
5416 .next()
5417 .cloned()
5418 .expect("add_columns writes a new per-fragment data file");
5419 let version = dataset.version().version;
5420 drop(dataset);
5421
5422 let fresh = || Arc::new(Session::new(0, 0, Arc::new(ObjectStoreRegistry::default())));
5424 scan_verify_version(uri, version, &fresh(), &HashMap::new(), &None)
5425 .await
5426 .expect("intact version must pass scan-verify");
5427 std::fs::write(&column_add_file, b"").unwrap();
5428 let verdict = scan_verify_version(uri, version, &fresh(), &HashMap::new(), &None).await;
5429 assert!(
5430 verdict.is_err(),
5431 "zeroed column-add data file must fail scan-verify",
5432 );
5433 }
5434
5435 #[test]
5436 fn cleanup_due_gates_on_version_interval() {
5437 assert!(cleanup_due(0, 1));
5439 assert!(cleanup_due(7, 1));
5440 assert!(cleanup_due(5, 0));
5441 assert!(cleanup_due(0, 16));
5443 assert!(cleanup_due(16, 16));
5444 assert!(cleanup_due(48, 16));
5445 assert!(!cleanup_due(15, 16));
5446 assert!(!cleanup_due(17, 16));
5447 assert!(!cleanup_due(31, 16));
5448 }
5449
5450 #[test]
5451 fn derived_target_rows_tracks_row_size_and_byte_cap() {
5452 let parts_like = [FragmentStat {
5455 bytes: Some(665_000_000),
5456 rows: 511_000,
5457 deleted_rows: 0,
5458 }];
5459 let target = derived_target_rows(&parts_like);
5460 assert!((80_000..150_000).contains(&target), "{target}");
5461 let unknown = [FragmentStat {
5463 bytes: None,
5464 rows: 511_000,
5465 deleted_rows: 0,
5466 }];
5467 assert_eq!(
5468 derived_target_rows(&unknown),
5469 MAX_TARGET_ROWS_PER_FRAGMENT as usize
5470 );
5471 let tiny = [FragmentStat {
5473 bytes: Some(1_000_000),
5474 rows: 100_000,
5475 deleted_rows: 0,
5476 }];
5477 assert_eq!(
5478 derived_target_rows(&tiny),
5479 MAX_TARGET_ROWS_PER_FRAGMENT as usize
5480 );
5481
5482 let incident_parts = [FragmentStat {
5483 bytes: Some(1_027_449_798),
5484 rows: 105_087,
5485 deleted_rows: 0,
5486 }];
5487 let incident_target = derived_target_rows(&incident_parts);
5488 assert_eq!(incident_target, 13_727);
5489 assert!(
5490 u128::from(incident_parts[0].bytes.unwrap()) * incident_target as u128 * 2
5491 <= u128::from(incident_parts[0].rows) * u128::from(TARGET_FRAGMENT_BYTES)
5492 );
5493
5494 let huge = [FragmentStat {
5496 bytes: Some(1_000_000_000),
5497 rows: 1,
5498 deleted_rows: 0,
5499 }];
5500 assert_eq!(derived_target_rows(&huge), 1);
5501 }
5502
5503 #[test]
5504 fn namespace_error_code_walks_wrapped_chain() {
5505 let direct = lance::Error::namespace_source(Box::new(NamespaceError::TableNotFound {
5506 message: "missing".into(),
5507 }));
5508 assert!(is_namespace_error_code(&direct, ErrorCode::TableNotFound));
5509
5510 let wrapped = lance::Error::namespace_source(Box::new(direct));
5511 assert!(is_namespace_error_code(&wrapped, ErrorCode::TableNotFound));
5512
5513 let other_code =
5514 lance::Error::namespace_source(Box::new(NamespaceError::NamespaceNotFound {
5515 message: "nope".into(),
5516 }));
5517 assert!(!is_namespace_error_code(
5518 &other_code,
5519 ErrorCode::TableNotFound
5520 ));
5521
5522 let not_namespace = lance::Error::internal("unrelated");
5523 assert!(!is_namespace_error_code(
5524 ¬_namespace,
5525 ErrorCode::TableNotFound
5526 ));
5527 }
5528
5529 #[tokio::test]
5533 async fn store_opens_via_namespace_and_scan_works() -> Result<()> {
5534 let temp = TempDir::new()?;
5535 let url = Url::from_directory_path(temp.path())
5536 .map_err(|()| anyhow::anyhow!("temp path is not absolute"))?;
5537 let handle = Handle::open(&url).await?;
5538 let cases: [(Table, &[&str]); 3] = [
5541 (Table::Sessions, &["id"]),
5542 (Table::Messages, &["id"]),
5543 (Table::Parts, &["id"]),
5544 ];
5545 for (table, projection) in cases {
5546 let scanner = handle
5547 .scan(table, ScanOpts::project_only(projection))
5548 .await?;
5549 let batch = scanner.try_into_batch().await?;
5550 assert_eq!(batch.num_rows(), 0, "fresh table should be empty");
5551 }
5552 Ok(())
5553 }
5554}