Skip to main content

pond/
substrate.rs

1//! The storage substrate (spec.md#substrate): pond's one seam to Lance,
2//! generic over consumers.
3
4use 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;
50/// Embedded-row count at which pond builds the IVF_SQ vector index on
51/// `messages.vector` (spec.md#search). Below it, vector search runs a
52/// brute-force flat scan - exact and fast at small and medium scale, and
53/// IVF_SQ cannot train well on fewer vectors anyway.
54pub const VECTOR_INDEX_ACTIVATION_ROWS: usize = 100_000;
55
56/// Segment count at which an incremental index fold consolidates instead of
57/// appending. Each `optimize_indices(append)` writes a new same-name segment
58/// (lance `num_indices_to_merge=0`), and every vector/FTS query reads the
59/// probed partition or token postings from *every* segment - so unbounded
60/// delta growth multiplies per-query object-store round-trips. At this many
61/// segments pond folds with `merge` to collapse them back into one.
62pub const DELTA_MERGE_THRESHOLD: usize = 4;
63
64// ---------------------------------------------------------------------------
65// Storage addresses (spec.md#storage-url-grammar)
66// ---------------------------------------------------------------------------
67
68/// A parsed pond storage address. The fat-URL grammar
69/// (`s3+https://host/bucket/prefix`) folds the endpoint into the address so
70/// it can never desync from the bucket (the litestream out-of-band-endpoint
71/// failure class); parsing splits it back into the URL Lance opens plus the
72/// `object_store` options the endpoint implies.
73#[derive(Debug, Clone, PartialEq)]
74pub struct StorageUrl {
75    /// The address as written, canonicalized (scheme/host lowercased by
76    /// `url`, default port stripped, recognized query params removed). Scope
77    /// matching (spec.md#creds-scope-match) and display use this form.
78    canonical: Url,
79    /// The URL handed to Lance.
80    lance: Url,
81    /// Options implied by the scheme - lowest precedence in assembly.
82    scheme_options: Vec<(&'static str, String)>,
83    /// Recognized `?key=value` params - highest precedence.
84    query_options: Vec<(&'static str, String)>,
85    /// `?creds=<name>`: explicit set binding, beats scope matching.
86    creds_pointer: Option<String>,
87    /// Endpoint pieces for the `s3+` schemes. The final endpoint URL depends
88    /// on the resolved `virtual_hosted_style_request` value (object_store
89    /// wants the bucket inside the endpoint host under virtual-hosted
90    /// addressing), so it is assembled at resolve time, not parse time.
91    endpoint: Option<S3Endpoint>,
92}
93
94#[derive(Debug, Clone, PartialEq)]
95struct S3Endpoint {
96    scheme: &'static str,
97    /// host[:port]
98    authority: String,
99    bucket: String,
100}
101
102/// Query params pond recognizes (and strips before the URL reaches Lance).
103/// Anything else is a hard error - a typoed param must not silently reach
104/// the object store as part of the path.
105const RECOGNIZED_QUERY_PARAMS: [&str; 3] = ["creds", "region", "virtual_hosted_style_request"];
106
107impl StorageUrl {
108    /// Parse a storage address (spec.md#storage-url-grammar): bare/`~` paths,
109    /// `file://`, `s3://`, `s3+https://` / `s3+http://`, `gs://`, `az://`,
110    /// and the test-only `memory://` / `shared-memory://`.
111    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        // Bare paths, `~/...`, and `file://` go through Lance's own
117        // `uri_to_url` so pond accepts exactly what Lance accepts.
118        if !trimmed.contains("://") || trimmed.starts_with("file://") {
119            let url =
120                uri_to_url(trimmed).with_context(|| format!("invalid storage path {trimmed:?}"))?;
121            // Bare paths percent-encode `?` (a legal filename character), so
122            // only an explicit `file://...?x=y` parses a query here. No local
123            // scheme takes one; reject like the remote schemes do instead of
124            // silently carrying it into the path Lance opens.
125            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        // RFC 3986 deprecates userinfo; argv/history/ps/logs leak it. Never.
133        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                // `url` treats non-special schemes' default ports as
165                // explicit; strip them so scope matching can't split on
166                // `:443` vs nothing.
167                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                // Virtual-hosted is the Hetzner / R2 / B2 default, but an IP
191                // host can't carry a bucket subdomain (`bucket.127.0.0.1`
192                // does not resolve), so MinIO-style IP endpoints flip to
193                // path-style. Override either way via the creds-set field or
194                // `?virtual_hosted_style_request=`. Note: `url` keeps IPv4
195                // hosts as `Host::Domain` on non-special schemes, hence the
196                // explicit IpAddr parse; IPv6 brackets still need the Host
197                // match.
198                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                    // S3-compatible stores ignore the SigV4 region, so a
204                    // deterministic default (the DuckDB / litestream
205                    // convention) beats Lance's env-chain fallback, where a
206                    // stray AWS_REGION changes behavior. Real AWS (`s3://`,
207                    // no endpoint) auto-resolves the bucket region inside
208                    // Lance instead. Override: creds-set field or ?region=.
209                    ("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    /// A scheme with no creds machinery: canonical == lance, no options.
256    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    /// The URL Lance opens (endpoint folded into options, not the URL).
268    pub fn lance_url(&self) -> &Url {
269        &self.lance
270    }
271
272    /// The canonical as-written address - what scope matching compares
273    /// against and what display surfaces show (it carries the endpoint).
274    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    /// Render for human output: local URLs as plain paths, remote verbatim.
283    pub fn display(&self) -> String {
284        config::display(&self.canonical)
285    }
286
287    /// Whether this scheme authenticates at all. `file`, `memory`, and
288    /// `shared-memory` take no credentials; resolution skips them entirely.
289    fn takes_credentials(&self) -> bool {
290        !matches!(
291            self.canonical.scheme(),
292            "file" | "file+uring" | "memory" | "shared-memory"
293        )
294    }
295
296    /// Resolve this address against the configured creds sets
297    /// (spec.md#creds-scope-match): `?creds=` pointer > longest scoped
298    /// prefix match > the scope-less catch-all > none (object_store's
299    /// ambient SDK chain). Option assembly, later wins: scheme-derived ->
300    /// matched set (non-secret fields + `extra`, then materialized secrets)
301    /// -> URL query params.
302    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        // The endpoint is assembled last: under virtual-hosted addressing
391        // object_store expects the bucket inside the endpoint host, so the
392        // URL depends on the final virtual_hosted_style_request value. An
393        // explicit endpoint in `extra` wins (the escape hatch).
394        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
420/// (canonical URL, recognized query options, `?creds=` pointer).
421type StrippedQuery = (Url, Vec<(&'static str, String)>, Option<String>);
422
423/// Pull recognized query params off the URL; reject unrecognized ones.
424fn 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
445/// Parse a `[creds.*] scope` URL prefix into the same canonical form
446/// `StorageUrl::parse` produces, so comparison is exact.
447pub(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
464/// spec.md#creds-scope-match: scheme, host, and port equal; path matches at
465/// `/` segment boundaries only (`.../pond` does not match `.../pond-2`). No
466/// cross-scheme normalization: a `s3+https://host/bucket/` scope does not
467/// match a `s3://bucket/` URL.
468fn 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/// How a creds set got bound to a URL - surfaced in binding lines so a wrong
484/// match is visible before any auth error.
485#[derive(Debug, Clone, Copy, PartialEq, Eq)]
486pub enum BindVia {
487    /// `?creds=<name>` pointer on the URL.
488    Pointer,
489    /// Longest-prefix `scope` match.
490    Scope,
491    /// The scope-less catch-all set.
492    CatchAll,
493}
494
495#[derive(Debug, Clone, PartialEq)]
496pub enum CredsBinding {
497    /// A `[creds.<name>]` set bound to this URL.
498    Set { name: String, via: BindVia },
499    /// No set matched; object_store's ambient SDK chain applies (AWS_* env,
500    /// shared credentials file, IMDS/container metadata). A documented
501    /// invariant, not an accident - instance profiles and OIDC work with
502    /// zero pond config.
503    Ambient,
504    /// Local / in-memory scheme; credentials don't apply.
505    NotApplicable,
506}
507
508impl CredsBinding {
509    /// One-line human rendering for binding lines and `pond config show`.
510    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/// A storage address with its options assembled and secrets materialized -
527/// everything `Store::open_with_options` needs, plus the binding for
528/// display.
529#[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
546/// Names of defined creds sets that bound to none of this invocation's URLs
547/// (spec.md#creds-scope-match: misbinding must never be silent). Empty when
548/// the invocation touched no credential-taking URL - a local-only command
549/// must not nag about sets kept for remote work.
550pub 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
571/// Materialize one logical secret from its inline / `_file` / `_command`
572/// variant (validation guarantees at most one is set).
573fn 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
598/// Run a `*_command` secret source. Output is cached per command text per
599/// process, so N URLs resolving through one set cost one subprocess.
600fn 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
634/// Strip exactly one trailing newline (the one `echo` / `op read` append);
635/// anything beyond that is part of the secret.
636fn 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/// `pond storage check` failure classes, each with its own exit code at the
647/// CLI so cron and CI can branch on them. Display carries only the
648/// fix-naming lead; the underlying error is exposed separately through
649/// [`CheckFailure::concise_cause`] so surfaces stay one readable line
650/// instead of trailing the upstream chain (Lance flattens its inner errors
651/// into each level's Display, so the raw chain prints the same failure
652/// several times over).
653#[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    /// The root cause, condensed to one operator-readable line: the deepest
671    /// error in the chain with upstream noise stripped - Lance's bug-report
672    /// boilerplate, internal `<WORKSPACE>` source locations, and the repeated
673    /// wrapper text that follows them. `None` for `OccUnsupported`, whose
674    /// `detail` is already curated into its Display.
675    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
684/// One-line root cause for a probe error. Takes the deepest chain entry
685/// (each outer Lance/object_store layer re-prints its inner error, so the
686/// deepest is the least redundant), cuts at the first internal source
687/// location (everything after it is upstream re-printing), strips Lance's
688/// bug-report boilerplate, and middle-truncates - the tail is kept because
689/// wrapped transport errors put the root (DNS, connect) at the end.
690fn 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
716/// Probe a resolved storage destination end-to-end (spec.md#substrate): a
717/// conditional `PutMode::Create` pair proving the `If-None-Match` -> 412 OCC
718/// primitive Lance's commit handler relies on, then read-back and delete of
719/// the synthetic key.
720pub 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, &params)
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    // The probe key exists from here on: run the remaining steps, then
754    // best-effort delete it whatever they returned - a failed probe must
755    // not leave litter behind.
756    let outcome = async {
757        // The second create MUST lose: this is the `If-None-Match: *` -> 412
758        // primitive multi-writer OCC stands on. A backend that lets it
759        // through (or rejects the header) silently overwrites concurrent
760        // commits.
761        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
801/// Map an `object_store` error onto the check's failure classes: an auth
802/// error is attributed to the bound creds set when one matched, and to the
803/// (empty) ambient chain when none did; everything else is I/O.
804fn classify_check_error(
805    error: object_store::Error,
806    binding: &CredsBinding,
807    step: &str,
808) -> CheckFailure {
809    use object_store::Error as OsError;
810    // Lance erases a missing-credentials failure into a `Generic` error - the
811    // typed `Unauthenticated` never surfaces for an empty provider chain - so
812    // also match the AWS SDK's rendered `CredentialsNotLoaded` signal. Both
813    // are auth-class: attributed to the bound set, else the empty ambient chain.
814    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
836/// Per-task fragment-count backstop: tasks this wide bypass the width and
837/// amplification checks once the merge can shrink the fragment count, bounding
838/// manifest growth. As policy cap, 0 disables all task filtering (tests).
839pub const DEFAULT_COMPACTION_FRAGMENT_CAP: usize = 64;
840
841/// Fragments are sized by bytes, not Lance's 1M-row default: kilobyte-average
842/// rows make a row target tolerate multi-GiB fragments that compaction
843/// re-rewrites wholesale to absorb tiny appends (~190 GiB/day of churn).
844pub const TARGET_FRAGMENT_BYTES: u64 = 256 * 1024 * 1024;
845
846/// Ceiling = Lance's own default.
847const MAX_TARGET_ROWS_PER_FRAGMENT: u64 = 1024 * 1024;
848
849/// Keep a task only when the merged-in remainder is >= largest/this:
850/// size-tiered amortization, O(log n) lifetime rewrites per row.
851pub const COMPACTION_ABSORB_FACTOR: u64 = 4;
852
853/// Default manifest-retention window for the safe cleanup pass. Matches
854/// LanceDB's recommended OSS-operator practice (lancedb docs: performance.mdx,
855/// tables/update.mdx). With `delete_unverified=false`, Lance's 7-day
856/// in-progress guard still protects unverified files regardless of this value
857/// (`UNVERIFIED_THRESHOLD_DAYS` in lance/dataset/cleanup.rs).
858pub fn default_cleanup_older_than() -> chrono::Duration {
859    // Toward Lance's 1 h floor: fewer retained manifest versions = cheaper
860    // remote open (spec.md#search). The append fast-path already curbs the
861    // version churn that earlier forced a wider window.
862    chrono::Duration::hours(1)
863}
864
865/// `pond sync` runs every few minutes; reclaiming old manifest versions on
866/// every run pays the full version-log walk over S3 (~9 s measured on the real
867/// corpus) to free roughly one version. Amortize by cleaning only when a
868/// table's manifest version is a multiple of this many commits. Explicit
869/// `pond optimize` and the one-shot `pond copy` keep interval 1 (clean every
870/// run) so maintenance and durability moves are never skipped.
871pub const DEFAULT_SYNC_CLEANUP_INTERVAL: u64 = 16;
872
873/// `pond sync` defers a scalar (BTree/bitmap) index fold until its unindexed
874/// tail reaches this many rows. Lance 7.0.0 ignores `OptimizeOptions::append()`
875/// for scalar indexes and rewrites the whole index file on every fold
876/// (O(index size), not O(delta)), so folding on every tiny sync pays a full
877/// rewrite for a handful of new rows. Batching amortizes that rewrite; the
878/// deferred tail stays correct for get/count/sql (they scan it) with scan cost
879/// bounded by this cap, and vector/FTS still fold every run so search recall is
880/// unaffected. `pond optimize`/`pond copy` fold every run (threshold `0`).
881pub const DEFAULT_SYNC_SCALAR_FOLD_ROWS: usize = 50_000;
882
883/// Defer the FTS + vector (IVF) index fold until the unindexed tail reaches this
884/// many rows; `0` folds every run. Unlike the scalar fold (deferred because Lance
885/// rewrites the whole index file), FTS/vector fold via a cheap delta append - the
886/// reason to batch them is the per-sync S3 round-trip + commit storm, not a
887/// rewrite. Between folds the tail stays fully searchable: the retrievers drop
888/// `fast_search` whenever an unindexed tail exists, so Lance index-probes the
889/// folded rows and flat-scans the (threshold-bounded) tail (fts.md "Index
890/// Maintenance"). The cap bounds that tail-scan cost. `pond optimize`/`pond copy`
891/// fold every run (threshold `0`).
892pub const DEFAULT_SYNC_INDEX_FOLD_ROWS: usize = 5_000;
893
894/// Resolved per-call inputs to the storage-maintenance pass. Built from
895/// `[maintenance]` (and any per-invocation CLI override) at the entry point;
896/// threaded down to `optimize_table_compact` so the substrate never re-reads
897/// `Config` itself.
898#[derive(Debug, Clone, Copy)]
899pub struct MaintenancePolicy {
900    /// See [`DEFAULT_COMPACTION_FRAGMENT_CAP`]; `0` disables the veto.
901    pub compaction_fragment_cap: usize,
902    /// Manifest-retention window handed to `cleanup_old_versions`.
903    pub cleanup_older_than: chrono::Duration,
904    /// Run `cleanup_old_versions` for a table only when its manifest version is
905    /// a multiple of this (`1` = every optimize). The frequent `pond sync` path
906    /// raises it so most syncs skip the version-log walk; see
907    /// [`DEFAULT_SYNC_CLEANUP_INTERVAL`].
908    pub cleanup_interval: u64,
909    /// Defer a scalar (BTree/bitmap) index fold until its unindexed tail reaches
910    /// this many rows; `0` folds every run. The frequent `pond sync` path raises
911    /// it so most syncs skip the full scalar-index rewrite Lance 7.0.0 does on
912    /// every fold; see [`DEFAULT_SYNC_SCALAR_FOLD_ROWS`].
913    pub scalar_fold_row_threshold: usize,
914    /// Defer the FTS + vector (IVF) index fold until its unindexed tail reaches
915    /// this many rows; `0` folds every run. The frequent `pond sync` path raises
916    /// it so most syncs skip the per-fold S3 round-trip storm; recall stays
917    /// complete because the retrievers flat-scan the deferred tail. See
918    /// [`DEFAULT_SYNC_INDEX_FOLD_ROWS`].
919    pub index_fold_row_threshold: usize,
920}
921
922impl MaintenancePolicy {
923    /// Veto off: run every task Lance plans (the optimize tests assume this).
924    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    /// Amortize version cleanup over `interval` commits - the frequent
935    /// `pond sync` path uses this so most syncs skip the version-log walk.
936    #[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    /// Amortize the scalar-index fold over its unindexed tail - the frequent
943    /// `pond sync` path uses this so most syncs skip the full scalar-index
944    /// rewrite Lance 7.0.0 does on every fold.
945    #[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    /// Amortize the FTS + vector fold over its unindexed tail - the frequent
952    /// `pond sync` path uses this so most syncs skip the per-fold S3 round-trip
953    /// storm; recall stays complete via the retrievers' tail flat-scan.
954    #[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    /// The two per-family fold thresholds bundled for the indices phase, so the
961    /// two same-typed `usize`s can't be swapped at a call site.
962    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/// Per-index-family fold-deferral thresholds (rows); `0` folds that family every
971/// run. Bundled so the indices phase takes one param, not two swappable `usize`s.
972#[derive(Debug, Clone, Copy)]
973struct FoldThresholds {
974    scalar: usize,
975    index: usize,
976}
977
978struct FragmentStat {
979    /// `None` when the manifest lacks any file's size.
980    bytes: Option<u64>,
981    rows: u64,
982    deleted_rows: u64,
983}
984
985/// Data-file bytes of one fragment; `None` (poisoning) when any size is
986/// missing from the manifest.
987fn 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
1005/// Candidacy/merge target: HALF the rows the [`TARGET_FRAGMENT_BYTES`] output
1006/// budget holds at the table's average row size. Deriving the target at the
1007/// FULL byte budget made `target == the largest fragment re-encoding could
1008/// reliably produce`: no output could ever satisfy `physical_rows >= target`,
1009/// so the table was re-compacted every sync for a net-zero fragment change
1010/// (measured ~100-120s/sync on the remote store, 30->30 fragments).
1011/// Halving leaves 2x headroom so a byte-capped fragment lands comfortably above
1012/// the target and FREEZES, making compaction productive (merge small -> freeze
1013/// -> stop) instead of perpetual churn. A row floor would make the target
1014/// unreachable again for sufficiently wide rows.
1015fn 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
1032/// Name the first reason an optional compaction task cannot make progress.
1033/// Deletion materialization always passes because removing tombstones is useful.
1034fn 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/// Declarative description of one index pond keeps on a table. Created when
1086/// its trigger fires; folded forward by `pond optimize`.
1087#[derive(Debug, Clone)]
1088pub struct IndexIntent {
1089    /// Stable on-disk name. Must match across runs so existence checks
1090    /// resolve.
1091    pub name: &'static str,
1092    /// Column the index covers.
1093    pub column: &'static str,
1094    /// Condition evaluated against the live dataset before each cycle.
1095    pub trigger: IndexTrigger,
1096    /// How the params are built at create time. Some intents have static
1097    /// params (FTS, scalars); IVF_SQ needs the row count to size partitions.
1098    pub params: IndexParamsKind,
1099}
1100
1101/// When an [`IndexIntent`] should exist on disk.
1102#[derive(Debug, Clone)]
1103pub enum IndexTrigger {
1104    /// Build whenever the table has any rows. Used for FTS and scalar
1105    /// indices: there is no training cost worth delaying.
1106    OnAnyRows,
1107    /// Build when `count(<column> IS NOT NULL) >= threshold`. Used for the
1108    /// IVF_SQ vector index, which trains poorly on too few vectors.
1109    OnNonNullCount {
1110        column: &'static str,
1111        threshold: usize,
1112    },
1113}
1114
1115/// The lance-native shape of an [`IndexIntent`]'s params, dispatched to the
1116/// right `IndexParams` at create time.
1117#[derive(Debug, Clone)]
1118pub enum IndexParamsKind {
1119    /// `BuiltinIndexType::BTree` -> [`IndexType::BTree`];
1120    /// `BuiltinIndexType::Bitmap` -> [`IndexType::Bitmap`]; etc.
1121    Scalar(BuiltinIndexType),
1122    /// `InvertedIndexParams` with the word-level `simple` tokenizer plus
1123    /// English stemming, stop-words off (spec.md#search-language-neutral-index).
1124    /// Word retrieval beats character ngram ~2x on the real corpus at ~4x less
1125    /// index weight; substring/symbol lookup stays on the SQL `LIKE` /
1126    /// `contains_tokens` path, not here.
1127    InvertedFtsWord,
1128    /// `VectorIndexParams::with_ivf_sq_params` with cosine metric (e5 vectors
1129    /// are L2-normalized). 8-bit scalar quantization stores per-dimension codes
1130    /// in the index itself, so kNN computes distances from the prewarmed
1131    /// partition with no refine pass - PQ+refine instead re-reads ~k*factor
1132    /// exact vectors from the data files as scattered per-row GETs, the
1133    /// dominant per-query S3 request storm on a throttling remote store
1134    /// (spec.md#search). `max_iters` caps kmeans; partitions follow LanceDB's
1135    /// documented `num_rows // 4096` guidance, floored at one.
1136    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/// Anyhow-chain sentinel pond attaches when `retry_lance` exhausts attempts
1208/// against an OCC commit-conflict failure (spec.md#protocol). The wire layer
1209/// downcasts to this type to classify the outcome as `conflict` rather than
1210/// the generic `storage_unavailable`.
1211#[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/// Per-phase result for one table's pass through `Handle::optimize_table`.
1229/// spec.md#substrate 3.7 (`lance-index-maintenance`): the indices phase and the
1230/// compaction phase get independent retry budgets and independent commits,
1231/// so a hot writer that starves the Rewrite cannot abort the index Update.
1232#[derive(Debug)]
1233pub enum PhaseOutcome {
1234    /// Phase attempted and committed work.
1235    Ok,
1236    /// Phase attempted; no work was needed.
1237    Noop,
1238    /// Phase attempted; OCC retry budget exhausted on conflict (the operator
1239    /// can rerun later once the hot writer quiesces).
1240    SkippedConflict,
1241    /// Phase failed with a non-conflict error.
1242    Failed(anyhow::Error),
1243    /// Phase not requested by the caller (e.g. compaction skipped under
1244    /// `Store::build_indices_only`).
1245    NotAttempted,
1246}
1247
1248impl PhaseOutcome {
1249    pub fn is_failed(&self) -> bool {
1250        matches!(self, Self::Failed(_))
1251    }
1252}
1253
1254/// What `Handle::optimize_table` did for one table.
1255#[derive(Debug)]
1256pub struct TableOptimizeOutcome {
1257    pub table: Table,
1258    pub indices: PhaseOutcome,
1259    pub compaction: PhaseOutcome,
1260}
1261
1262/// Boundary event during one `Handle::optimize_table` pass. The CLI binds a
1263/// progress callback to render a live spinner; library callers pass `None`.
1264#[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    /// Intra-index liveness, forwarded from Lance's `IndexBuildProgress`
1277    /// callbacks (FTS tokenize/copy, IVF train/shuffle/merge, BTree/Bitmap
1278    /// build stages). Fires many times per index between `PhaseStart` /
1279    /// `PhaseDone`; the spinner just overwrites its message each tick.
1280    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
1311/// `Arc` rather than `Box` so the same callback can be cloned into the
1312/// `PondIndexProgress` Arc that Lance's `IndexBuildProgress` builder demands -
1313/// otherwise intra-index stage events have no path back to the CLI spinner.
1314pub 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
1322/// Bridges Lance's `IndexBuildProgress` async callbacks (`stage_start`,
1323/// `stage_progress`, `stage_complete`) into pond's `OptimizeEvent::IndexStage`
1324/// stream so the CLI spinner can show "fts tokenize_docs 1.4M / 2M rows"
1325/// instead of going dark for 10-20 minutes during a single `create_index` or
1326/// `optimize_indices` call. Remembers the active stage's `total` / `unit` so
1327/// `stage_progress` (which only carries `completed`) can render a full
1328/// fraction. Emissions are throttled to one every 100ms; FTS's per-batch
1329/// `stage_progress` calls would otherwise contend the spinner mutex.
1330struct PondIndexProgress {
1331    callback: OptimizeProgressFn,
1332    table: Table,
1333    index: String,
1334    state: std::sync::Mutex<PondIndexStageState>,
1335}
1336
1337// `IndexBuildProgress` requires `Debug`; the `callback` field is
1338// `Arc<dyn Fn...>` which has no `Debug` impl, so derive doesn't apply.
1339impl 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
1440/// True when the chain root is one of Lance's commit-conflict variants
1441/// (`CommitConflict`, `RetryableCommitConflict`, `TooMuchWriteContention`).
1442/// Everything else (timeouts, IAM denials, disk errors) is not a conflict.
1443pub 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
1454/// True when `retry_lance` exhausted retries against an OCC conflict and
1455/// attached `ConflictExhausted` to the chain head.
1456fn is_conflict_exhausted(error: &anyhow::Error) -> bool {
1457    error.chain().any(|cause| cause.is::<ConflictExhausted>())
1458}
1459
1460/// True when the chain root is Lance's `Index` error class - a structural
1461/// index fault (e.g. delta segments with mismatched posting tail codecs) that
1462/// retry cannot clear and only a from-scratch rebuild repairs.
1463pub 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/// On-disk byte totals for the three session datasets, plus everything else
1470/// under the data-dir root. Sized by listing through Lance's object-store
1471/// layer (spec.md#lance-chokepoints-storage) so `file://` and `s3://` behave alike.
1472#[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/// `data/` bytes on disk vs bytes the latest manifest references; the gap is
1484/// superseded versions awaiting the cleanup retention window.
1485#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1486pub struct DataLiveness {
1487    pub on_disk: u64,
1488    /// `None` when the manifest lacks any referenced file's size.
1489    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 match. Emitted as `regexp_like(<col>, '<pat>')`. Never pushes
1528    /// down to BTREE indexes (Lance's scalar-index-expr parser ignores it),
1529    /// so the filter is a full-scan-with-predicate - acceptable for
1530    /// human-driven `--project re:...` queries, not for hot paths.
1531    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                // Wrap in parens so the disjunction composes safely as a child
1569                // of an outer `And` (SQL `OR` binds looser than `AND`).
1570                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/// Read-side options for `Handle::scan`: optional prefilter predicate and
1594/// optional projection. Default = no filter, all columns.
1595#[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/// Lance cache caps in bytes. `None` lets the substrate pick the backend-aware
1629/// default (local FS gets a tighter cap; object stores stay near Lance's
1630/// defaults). Wired through `Store::open_with_options` from `[runtime]`.
1631#[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
1646/// Local-FS default: tight enough that a long-lived `pond mcp` lands well
1647/// under the 500 MiB target without measurable latency cost vs Lance's 6 GiB
1648/// default (see `benches/serve_mem_bench.rs --cap-sweep`).
1649const LOCAL_INDEX_CACHE_BYTES: usize = 256 * 1024 * 1024;
1650const LOCAL_METADATA_CACHE_BYTES: usize = 128 * 1024 * 1024;
1651/// Object-store defaults: latency to refill is per-page, so keep more in cache
1652/// than local - but bounded above the warm working set, not Lance's 6 GiB.
1653/// Post word-tokenizer FTS that set is ~450 MB (simple invert + IVF_SQ aux), so
1654/// 1 GiB holds both indices warm with headroom while capping the RSS ceiling.
1655const 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    /// One `lance::Session` shared across all three datasets. Carries the
1674    /// metadata + index caches and the `ObjectStoreRegistry` (which holds
1675    /// the underlying object_store / S3 client). Sharing the session means
1676    /// one cache pool covers all three tables and one S3 client serves all
1677    /// three datasets - load-bearing on object-store backends where a
1678    /// per-dataset client would mean 3x the connection pools and 3x the
1679    /// credential refreshes (lance/src/dataset/builder.rs:509-517).
1680    #[allow(dead_code)]
1681    session: Arc<Session>,
1682    /// The `lance-namespace` catalog seam. v1 uses the Directory impl;
1683    /// future hosted pond swaps to "rest" without touching read/write paths
1684    /// (spec.md#lance-chokepoints-catalog).
1685    nm: Arc<dyn LanceNamespace>,
1686    /// Namespace identifier this handle binds to. v1 is always `root()`; the
1687    /// typed seam matches `resolve_namespace`'s return so multi-namespace
1688    /// routing can land without churning call sites (spec.md#wire-namespace-resolution).
1689    nm_ident: NamespaceIdent,
1690    /// Object-store options threaded through every `DatasetBuilder` and
1691    /// `Dataset::write` call so refresh / index-creation paths inherit the
1692    /// same credentials and region as the initial open. Empty on local-FS
1693    /// installs.
1694    storage_options: HashMap<String, String>,
1695    /// Data-dir URL the handle was opened against. `pond status` reads this
1696    /// to display where the bytes live and to decide whether to walk a local
1697    /// directory or issue a remote `LIST` for sizing.
1698    location: Url,
1699    /// Freshness window applied to the lazily-opened `sessions` and `parts`
1700    /// datasets when they first open, matching the eager `messages` open's
1701    /// scheme-keyed `refresh_after`.
1702    lazy_refresh_after: Duration,
1703    /// Object-store wrapper (fsync durability + index disk cache + io-trace)
1704    /// applied on every dataset open, including the lazy sessions/parts opens
1705    /// and any re-open.
1706    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.lance` opens lazily, like `parts`: the search request path
1744    /// reads only `messages`. Writers (ingest), `pond status`, restore, and the
1745    /// daemon's background index-cache GC open it on first use.
1746    sessions: OnceCell<Mutex<CachedDataset>>,
1747    messages: Mutex<CachedDataset>,
1748    /// `parts.lance` opens lazily on the first read or write that needs it:
1749    /// any get read (every mode reads parts to build summaries), grouped
1750    /// search hydrating user-hit summaries, or ingest with Part events. A
1751    /// process that does none of those skips the file, saving its metadata
1752    /// pages and file handle at cold-open. The OnceCell makes init
1753    /// single-flight; the inner `Mutex<CachedDataset>` then behaves identically
1754    /// to the other two.
1755    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/// Outcome of one [`Handle::append_stream`] write. Lance's `execute_stream`
1785/// returns only the new `Dataset` (no write summary), so these totals are
1786/// captured from the cumulative `WriteStats` ticks plus pond's own OCC attempt
1787/// counter.
1788#[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/// Monotonic high-water fold over the cumulative `WriteStats` ticks
1797/// `append_stream` receives. Lance restarts a stream's cumulative counters from
1798/// zero on each OCC retry, so `fetch_max` keeps the fold monotonic - a retry
1799/// contributes nothing until it passes the prior mark, making `AppendStats`
1800/// exact under retries.
1801#[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
1826/// Append-mode write params. Byte-sized fragments, not Lance's 90 GB default:
1827/// kilobyte rows would otherwise pack multi-GiB fragments that compaction
1828/// rewrites wholesale (see `TARGET_FRAGMENT_BYTES`). Reuses the create params so
1829/// appended fragments match the table's storage version / row-id mode.
1830fn 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    /// Open without storage options or explicit cache caps. Backend-aware
1839    /// defaults from `[runtime]` apply.
1840    pub async fn open(location: &Url) -> Result<Self> {
1841        Self::open_with_options(location, HashMap::new(), RuntimeCaps::default()).await
1842    }
1843
1844    /// Live size in bytes of the shared Lance session caches (index + metadata).
1845    /// Walks the caches, so it is not cheap - bench/diagnostic use only.
1846    pub fn lance_cache_bytes(&self) -> u64 {
1847        self.session.size_bytes()
1848    }
1849
1850    /// Open with object-store options handed through to Lance verbatim, plus
1851    /// the resolved `[runtime]` cache caps. Object-store keys are the
1852    /// `object_store` crate's standard config names; pond does not parse them.
1853    /// Opening datasets never performs index work; index lifecycle lives under
1854    /// `Handle::optimize_table`. `sessions.lance` and `parts.lance` open lazily
1855    /// on first use.
1856    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    /// Like [`Self::open_with_options`], plus an `_indices/*` disk cache rooted
1865    /// at `index_cache_dir` (caller supplies it, mirroring `ensure_rowmap`) so a
1866    /// fresh process skips the cold index load. Ignored for local-FS stores.
1867    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        // One Session shared across all three datasets so metadata/index
1884        // caches and the object_store registry (and thus any S3 client) are
1885        // pooled rather than duplicated three times. Caps are sized by the
1886        // `[runtime]` block; explicit values from `caps` win, otherwise the
1887        // local/remote backend default kicks in.
1888        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        // Build the lance-namespace catalog seam once (spec.md#lance-chokepoints-catalog).
1895        // The `root` property is whatever URL the Directory impl understands;
1896        // `uri_to_url` (lance-io/object_store.rs) accepts both bare paths and
1897        // URLs, so passing the scheme-qualified URL for local FS works the
1898        // same as the bare-path form. Trailing slash stripped for clean logs.
1899        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        // Object-store credentials/region/endpoint flow into the namespace
1904        // via the `storage.<key>` property convention (lance-namespace-impls
1905        // dir.rs from_properties: lines 423-436).
1906        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        // spec.md#lance-handle-freshness: refresh window is scheme-keyed. Local-FS
1915        // manifest reads are microsecond-cheap, so `0` (always-refresh) is
1916        // essentially free and removes the stale-read window entirely. Object
1917        // stores have real per-call cost; `5s` caps manifest fetch overhead at
1918        // acceptable lag for human-driven queries.
1919        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    /// Read-only view of the `storage_options` the handle was opened with.
1960    /// `pond status` needs them to instantiate a raw `object_store` client
1961    /// that can `LIST` the remote bucket for sizing.
1962    pub fn storage_options(&self) -> &HashMap<String, String> {
1963        &self.storage_options
1964    }
1965
1966    /// Object-store URI for a `pond_sql` export artifact:
1967    /// `<location>/exports/<name>`. A sibling of the `*.lance` table dirs;
1968    /// the Directory namespace tracks tables in its `__manifest` table rather
1969    /// than by listing prefixes, so this prefix is never seen as a table
1970    /// (lance-namespace-impls dir/manifest.rs). Never `register_table`'d.
1971    fn export_uri(&self, name: &str) -> String {
1972        format!(
1973            "{}/exports/{name}",
1974            self.location.as_str().trim_end_matches('/')
1975        )
1976    }
1977
1978    /// `ObjectStoreParams` carrying the handle's `storage_options` so raw
1979    /// object-store opens (export I/O, `table_sizes` listing) inherit the same
1980    /// credentials/region as the dataset opens. Empty options -> no accessor.
1981    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    /// Write a `pond_sql` export artifact, reusing the handle's
1993    /// storage_options so S3 installs inherit the same credentials.
1994    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    /// Read a `pond_sql` export artifact back (for the
2009    /// `pond-sql-export://` MCP resource).
2010    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    /// Local filesystem path of an export artifact, when the data dir is
2025    /// `file://`. The stdio MCP client shares this filesystem, so it can read
2026    /// the file directly (e.g. duckdb/polars) instead of pulling base64 via
2027    /// `resources/read`. `None` on object-store installs.
2028    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    /// Insert-only merge: append new rows, never overwrite a matched PK.
2045    /// Returns rows inserted. The fold lives separately under
2046    /// `Handle::optimize_table` (spec.md#lance-index-maintenance).
2047    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    /// Insert-only merge that surfaces Lance's full `MergeStats`. Callers that
2059    /// need bytes written, file count, or OCC retry count (e.g. `pond copy`'s
2060    /// progress display) use this; the thin wrapper above keeps the
2061    /// affected-rows return for everyone else.
2062    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    /// Update-only merge: `WhenMatched::UpdateAll` on matched PKs; unmatched
2080    /// rows dropped. The fold lives separately under `Handle::optimize_table`.
2081    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    /// The OCC write-commit seam (spec.md#lance-chokepoints-write): every write -
2100    /// `merge` and the append paths - runs through here. It takes the cached
2101    /// handle's lock, hands `execute` the latest dataset, commits the dataset
2102    /// `execute` returns, and keeps the cache coherent - all under retry.
2103    /// `execute` builds the table-specific builder, runs it, and returns the new
2104    /// dataset plus its own stats payload; it reruns per OCC attempt, so it owns
2105    /// what it needs. Write-type specifics (params, stats, tracing) stay with the
2106    /// caller.
2107    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    /// [`Self::write_committed`] with a retry gate (see
2116    /// [`Self::retry_lance_filtered`]). `merge_insert` is idempotent on retry
2117    /// (`WhenMatched::DoNothing` re-reads and no-ops), so it retries everything;
2118    /// the bare `Append` path passes [`is_commit_conflict`] so a post-commit
2119    /// transient fault surfaces rather than re-appending into a duplicate.
2120    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    /// Shared merge path for [`Self::merge_insert`] and [`Self::merge_update`].
2145    /// Returns Lance's `MergeStats` verbatim so the progress layer can read
2146    /// `bytes_written` / `num_files_written` / `num_attempts` without a second
2147    /// round-trip; the thin wrappers above project to `u64` for callers that
2148    /// only need the affected-rows count.
2149    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                    // pond presents each PK at most once per batch; FirstSeen keeps
2174                    // the first occurrence rather than failing (Lance's default).
2175                    builder.source_dedupe_behavior(SourceDedupeBehavior::FirstSeen);
2176                    // Cleanup is operator-driven via `pond optimize`; the per-commit
2177                    // auto hook would add a LIST per write on remote backends without
2178                    // changing the steady-state retention.
2179                    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    /// Append a streamed source into `table` under a single commit - the
2205    /// bandwidth-bound counterpart to [`Self::merge`]. spec.md#session-durable-copy:
2206    /// rows that cannot collide on the destination (absent sessions) take this
2207    /// path. `Append` never joins or probes the target, so its cost is the
2208    /// bytes written, not the per-batch commit + key-scan that `merge_insert`
2209    /// pays - the fix for store-to-store copy being commit-latency-bound on
2210    /// remote object stores.
2211    ///
2212    /// `make_source` is a *factory*, not a prebuilt stream: a Lance scan stream
2213    /// is one-shot, so an OCC retry rebuilds it. A single per-call `WriteAccum`
2214    /// (shared across attempts, NOT fresh per attempt) makes the row/byte/file
2215    /// fold exact under retries.
2216    ///
2217    /// Unlike [`Self::append_batches`] this keeps the retry-everything
2218    /// `write_committed`: a transient fault during the large streamed upload
2219    /// almost always precedes the manifest commit (the rebuilt source re-uploads
2220    /// and the orphaned fragments are GC'd, no duplicate), so failing a full
2221    /// bulk copy on every transient to close the narrow lost-ack-after-commit
2222    /// window is the wrong trade. That rare window is surfaced by the copy
2223    /// verify's duplicate check instead (spec.md#session-movement-complete).
2224    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    /// [`Self::append_stream`] for batches pond already holds in memory (the sync
2274    /// write path) instead of a source-store scan. Row count is taken from the
2275    /// batches - exact under OCC retry without depending on the progress tick.
2276    ///
2277    /// Retries only on a commit *conflict*, not on transient faults: `Append`
2278    /// has no row-level idempotency, so re-running it after a manifest commit
2279    /// that landed but whose ack was lost would duplicate the rows. A conflict
2280    /// proves the commit did not land (re-append is safe); anything else
2281    /// surfaces and the caller's re-plan-from-current-state re-run heals it
2282    /// without doubling rows (spec.md#lance-deterministic-pk).
2283    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    /// Run the table-local maintenance cycle for the supplied index intents.
2332    /// Every index family folds incrementally via `optimize_indices`; none is
2333    /// rebuilt from scratch (spec.md#lance-index-maintenance).
2334    ///
2335    /// spec.md#substrate 3.7 (`lance-index-maintenance`): indices and compaction
2336    /// commit independently and use independent retry budgets, so a hot writer
2337    /// that starves compaction (Rewrite) does not abort the index build
2338    /// (Update) the operator actually asked for.
2339    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    /// Run only the indices phase for one table. Used by the optimize embed
2360    /// stage's tail
2361    /// to fold newly written vectors into the indices without paying the
2362    /// compaction retry budget while embed itself may still be writing.
2363    pub async fn optimize_table_indices_only(
2364        &self,
2365        table: Table,
2366        intents: &[IndexIntent],
2367        progress: Option<&OptimizeProgressFn>,
2368    ) -> PhaseOutcome {
2369        // Thresholds 0: this tail-fold path always folds every index; only
2370        // `pond sync` batches them (`with_scalar_fold_row_threshold` /
2371        // `with_index_fold_row_threshold`).
2372        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    /// Lance `cleanup_old_versions` for one table: reclaim files no manifest
2470    /// within the retention window references. No compaction and no new commit -
2471    /// it only deletes superseded files, so no OCC retry is needed.
2472    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    /// Build a prefiltered `Scanner` for `table`. Composable read entry
2501    /// point for callers that need to layer extra builder calls
2502    /// (`full_text_search`, `nearest`) on top of pond's predicate seam.
2503    /// Routine scans should prefer `Handle::scan`.
2504    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    /// Single read entry point: prefilter via `predicate`, optionally
2513    /// project, return the prepared `Scanner` (spec.md#lance-chokepoints-read).
2514    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    /// Collect the primary-key (`id`) set for `table`. Storage verification
2549    /// compares these sets across two stores: matching row counts can still
2550    /// hide divergent membership, so proving a destination is a complete
2551    /// superset of a source needs the ids, not the cardinalities
2552    /// (spec.md#substrate, `lance-deterministic-pk`).
2553    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    /// Names of every index on `messages` - the vector-index tests read this.
2564    #[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    /// Whether `messages` carries an index named `name`. Manifest-only and
2572    /// cache-backed (`load_indices` hits the dataset index cache), so it is
2573    /// cheap enough to gate `Scanner::fast_search` per query: fast-search
2574    /// returns an empty plan when the index is absent, so the retrievers must
2575    /// only opt in once it exists.
2576    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    /// Whether `messages` index `name` covers every row - it exists and has no
2583    /// unindexed tail - so `fast_search` (index-only) is complete. When a
2584    /// deferred fold has left a tail, the retrievers must omit `fast_search` so
2585    /// Lance index-probes the folded rows and flat-scans the tail, keeping recall
2586    /// complete (spec.md#search, fts.md "Index Maintenance"). Manifest-only and
2587    /// cache-backed, so it is cheap enough to gate `fast_search` per query.
2588    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    /// Reclaim cached `_indices/<uuid>` dirs no longer referenced by any table's
2606    /// manifest. No-op for local stores or a never-populated cache. Best-effort:
2607    /// a new index version naturally re-fetches, so an over-eager prune only
2608    /// costs one re-download.
2609    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    /// Count rows in `table` not yet covered by `index_name`. Manifest-only;
2631    /// a missing index reports the whole table. Powers `pond status`.
2632    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    /// Which table owns the named index, if any. Used by
2649    /// `pond optimize --drop-index <name>` to route the drop to the right
2650    /// dataset without sequentially probing-and-swallowing errors (the prior
2651    /// loop hid permission/network failures behind "no such index"). Runs the
2652    /// three `load_indices` calls in parallel; an error here is a real I/O
2653    /// failure and propagates with context.
2654    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    /// Drop the named index. Used by the `pond optimize --force-embed` model-swap path
2684    /// to retire an IVF_SQ whose centroids belong to the old distance
2685    /// space, before the next write re-bootstraps it over the new model's
2686    /// vectors. Errors when the index does not exist; callers may swallow
2687    /// that.
2688    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    /// Resolve each table's stored location through the namespace catalog
2700    /// (spec.md#lance-chokepoints-catalog) - no hardcoded `.lance` suffix.
2701    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    /// Whether the store holds synced data yet. `open` eagerly creates only the
2717    /// `messages` dataset; `sessions` and `parts` open lazily on first use
2718    /// (see `open_with_options`), so `parts`' presence is the "has been synced"
2719    /// signal - letting read-only surfaces (`pond status`) render an empty state
2720    /// instead of erroring on the first `parts` describe.
2721    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    /// On-disk byte totals for the three datasets plus the data-dir remainder.
2736    /// Every byte is sized by listing through Lance's object store
2737    /// (spec.md#lance-chokepoints-storage), identical for `file://` and `s3://`.
2738    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                &registry,
2745                &params,
2746                &self.table_location(sessions::SESSIONS).await?,
2747            )
2748            .await?;
2749        let messages = self
2750            .listed_size(
2751                &registry,
2752                &params,
2753                &self.table_location(sessions::MESSAGES).await?,
2754            )
2755            .await?;
2756        let parts = self
2757            .listed_size(
2758                &registry,
2759                &params,
2760                &self.table_location(sessions::PARTS).await?,
2761            )
2762            .await?;
2763        // `other` is whatever sits under the data-dir root but not in the three
2764        // tables (config.toml, stray index temp files): root total minus them.
2765        let root_total = self
2766            .listed_size(&registry, &params, self.location.as_str())
2767            .await?;
2768        let other = root_total.saturating_sub(sessions + messages + parts);
2769        let sessions_data = self
2770            .data_liveness(&registry, &params, Table::Sessions, sessions::SESSIONS)
2771            .await?;
2772        let messages_data = self
2773            .data_liveness(&registry, &params, Table::Messages, sessions::MESSAGES)
2774            .await?;
2775        let parts_data = self
2776            .data_liveness(&registry, &params, 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    /// Sum `ObjectMeta.size` for every object recursively under `uri`.
2810    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    /// Open `sessions.lance` on first use (spec.md#datasets). The search request
2836    /// path reads only `messages`; the daemon's background index-cache GC
2837    /// (`prune_index_cache`) opens this on a `serve`/`mcp` process. Single-flight
2838    /// via `OnceCell`, like `parts`.
2839    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    /// Open `parts.lance` on first use (spec.md#datasets). Single-flight via
2849    /// `OnceCell`; once initialized, behaves identically to the other two.
2850    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    /// Shared lazy-open path for the `sessions`/`parts` `OnceCell`s. `schema` is a
2856    /// thunk so the (local-CPU) schema build happens only on the cold init, not
2857    /// on every cache hit.
2858    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        // Default: retry every transient fault (spec.md#lance-retry-jitter).
2888        self.retry_lance_filtered(label, |_| true, operation).await
2889    }
2890
2891    /// Like [`Self::retry_lance`] but `should_retry` gates which errors are
2892    /// retried. [`Self::append_batches`] passes [`is_commit_conflict`]: a commit
2893    /// conflict means this writer's commit did NOT land, so re-running the
2894    /// operation is safe; any other error (notably a transient fault that may
2895    /// have arrived *after* the manifest commit landed - the lost-ack case) is
2896    /// surfaced instead of retried, because `Append` has no row-level
2897    /// idempotency and a blind re-append would duplicate. The caller's
2898    /// operation re-plans from current state on its own re-run, which is the
2899    /// idempotent recovery (spec.md#lance-deterministic-pk).
2900    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                    // `{:#}` walks anyhow's cause chain inline; `%error` (Display)
2919                    // drops everything below the top-level message.
2920                    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                    // spec.md#protocol: surface OCC failures as a typed `conflict`
2939                    // rather than the generic `storage_unavailable` bucket. The
2940                    // chain root is a `lance::Error` (commit-conflict family) when
2941                    // pond's retry layer exhausted because the manifest could not
2942                    // be advanced; everything else (timeouts, IAM, disk) stays
2943                    // `storage_unavailable`.
2944                    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        // Symmetric +/- `jitter` factor de-correlates concurrent retriers on
2957        // a contended manifest (spec.md#lance-retry-jitter); clamped to `max_backoff`.
2958        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}
2962/// Compaction phase: plan + amplification veto + execute + `cleanup_old_versions`,
2963/// one retry block, separate from the indices phase so a lost Rewrite race
2964/// does not abort index work.
2965///
2966/// Vetoes Lance-planned tasks instead of pre-gating on pond fragment math:
2967/// Lance bins split at index-coverage boundaries, so pond predictions diverge
2968/// from what Lance actually rewrites (the old run-sum gate latched open and
2969/// rewrote a 665 MiB tail fragment every 5-min sync). Only whole planned
2970/// tasks are filtered, so OCC and conflict semantics are untouched.
2971///
2972/// spec.md#lance-index-maintenance mandates FRI on by default, but at
2973/// v7.0.0-beta.16 `defer_index_remap=true` together with `stable-row-ids`
2974/// panics in `optimize.rs::commit_compaction` with "defer_index_remap
2975/// requires row_addrs but none were provided": `rewrite_files` skips
2976/// row_addrs when stable row ids are on, then the FRI builder demands
2977/// them. With stable_row_ids the remap step is already a no-op
2978/// (`optimize.rs:1490`: `needs_remapping = !uses_stable_row_ids() &&
2979/// !defer_index_remap`), so running without FRI is correct - we only
2980/// lose the documented concurrency-with-index-build benefit. Flip to
2981/// `true` once upstream fixes the conflict.
2982async 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        // Binary-copy eligible fragments (concatenate encoded pages, no
2998        // decode/re-encode) and fall back to Reencode automatically for blob
2999        // (parts), deletion-bearing, or schema-varied fragments. ~27% faster on
3000        // the messages/sessions reencode path, safe everywhere else.
3001        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    // Safe GC only. delete_unverified=false keeps Lance's 7-day in-progress
3070    // guard, so this never races a concurrent writer (spec.md#concurrency); GC
3071    // runs outside OCC, so the guard is what makes it safe on any backend.
3072    //
3073    // Gated: the walk over the version log is round-trip-bound on object stores
3074    // (~9s measured on the real corpus) and reclaims ~one version per run, so
3075    // the frequent `pond sync` path amortizes it over `cleanup_interval`
3076    // commits rather than paying it every sync (`pond optimize`/`pond copy`
3077    // keep interval 1). Skipping only delays reclaiming old versions - the next
3078    // due cleanup sweeps the accumulated backlog - so it is always safe.
3079    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        // Lance v7 `cleanup_old_versions` removes orphan files inside
3090        // `_indices/<uuid>/` but does NOT remove the parent dir, so failed/no-op
3091        // index merges accumulate empty UUID dirs forever (one inode each).
3092        // Harmless beyond inode pressure; tracked upstream. No pond-side FS sweep
3093        // here (spec.md#concurrency: Lance-native maintenance only); the sole
3094        // sanctioned direct-FS touches are the durability fsyncs inside Lance's
3095        // wrapper seam (spec.md#local-store-durability) and heal's quarantine
3096        // renames on an already-failed open (spec.md#local-store-self-heal).
3097        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
3114/// Gate for the version-cleanup walk: at interval `<= 1` it runs every optimize;
3115/// otherwise only when the manifest `version` is a multiple of it. A run whose
3116/// version steps past a multiple defers to the next one, so the gap between
3117/// cleanups is bounded and - since version 0 is a multiple of every interval -
3118/// cleanup always eventually fires; it is never skipped indefinitely.
3119fn cleanup_due(version: u64, interval: u64) -> bool {
3120    interval <= 1 || version.is_multiple_of(interval)
3121}
3122
3123/// Indices phase: create absent indexes, then fold trailing fragments into
3124/// every existing index via batched `optimize_indices` (append, or merge once a
3125/// family's delta segments reach `DELTA_MERGE_THRESHOLD`). Returns `true` if
3126/// anything committed.
3127async 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        // A trailing tail (rows written since the last fold) stays fully
3184        // searchable while deferred: the retrievers drop `fast_search` whenever
3185        // an index has an unindexed tail, so Lance index-probes the folded rows
3186        // and flat-scans the tail (spec.md#search, fts.md "Index Maintenance").
3187        // `DELTA_MERGE_THRESHOLD` keeps the per-fold segment count bounded.
3188        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        // Batch folds per family so a tiny sync doesn't pay a per-fold cost that
3197        // dwarfs the delta. Scalar (BTree/bitmap): Lance 7.0.0 rewrites the whole
3198        // index file per fold (O(index size)); defer until the tail is worth one
3199        // rewrite - get/count/sql read the deferred tail via scan. FTS + vector
3200        // (IVF): the fold is a cheap delta append, but each is an S3 round-trip +
3201        // commit storm; defer until the tail is worth one fold - search
3202        // flat-scans the deferred tail so recall stays complete either way.
3203        // `pond optimize`/`pond copy` pass 0 (fold every run).
3204        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        // FTS-only guard: folding a tail with zero non-null values writes an
3219        // empty delta segment, which Lance 7.0.0 reads back with the wrong
3220        // posting-tail codec (`Default` = VarintDelta vs the metadata-absent
3221        // default Fixed32), deterministically failing every later merge. The
3222        // probe is bounded: it reads only the tail fragments the fold itself
3223        // would read, stops at the first non-null value, and runs only after
3224        // the fold threshold already passed.
3225        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        // Every family folds incrementally via `optimize_indices` (the
3237        // append/merge batch below) - no full rebuild. BTree rewrites its index
3238        // file by merging the existing sorted pages with only the new fragments'
3239        // data; Bitmap/FTS/IVF_SQ accumulate delta segments. None re-scans
3240        // already-indexed source (spec.md#lance-index-maintenance).
3241        append_indices.push(intent.name.to_owned());
3242    }
3243
3244    if !append_indices.is_empty() {
3245        // Per-index segment count from the manifest loaded above (delta segments
3246        // share the intent name). Indices that have piled up
3247        // `DELTA_MERGE_THRESHOLD` segments fold with `merge` (collapse to one);
3248        // the rest take the cheap append. Splitting keeps each query reading few
3249        // segments without paying a consolidation on every tiny fold.
3250        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        // FTS delta segments are never merged: Lance 7.0.0's inverted merge
3261        // has two reproducible defects on real segments - "different posting
3262        // tail codecs" (a partitionless segment reports the derive-default
3263        // codec) and an index-out-of-bounds panic in InnerBuilder::merge_from
3264        // (token ids past the resized posting table; crashed the 5-min cron
3265        // sync in a loop). At the same threshold a merge would fire, the FTS
3266        // index instead rebuilds from scratch - the one consolidation path
3267        // Lance executes correctly. Scalar and vector merges are unaffected.
3268        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
3343/// Fragment-scoped scanner over the non-null rows of `column` - the shared
3344/// base of the existence probe and the indexable count below.
3345fn 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
3356/// True when any row in `fragments` holds a non-null value for `column`.
3357/// Existence probe for the FTS fold guard: scans only the given fragments,
3358/// stops at the first hit (`limit 1`), so its read is a subset of what the
3359/// fold it gates would read.
3360async 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
3375/// Count of rows in `fragments` holding a non-null value for `column`. Serves
3376/// the indexable status view; fragment-scoped, so bounded by the tail.
3377async 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        // Content indexes (FTS, IVF) take in only non-null rows - most message
3448        // rows carry a null `search_text`/`vector` (tool/system roles), so the
3449        // raw fragment row count vastly overstates the actionable backlog and
3450        // an all-null tail (which the FTS fold guard skips) would read as
3451        // stuck. The indexable view counts what a fold could actually index;
3452        // opt-in because the count scans the tail, which the per-sync summary
3453        // must not pay.
3454        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
3475/// Open the table at `table_name` via the namespace; create + initialize on
3476/// `TableNotFound`. Schema-checks the on-disk dataset against pond's
3477/// expectation so a stale data dir surfaces early.
3478///
3479/// Probes via `nm.describe_table` directly rather than `DatasetBuilder::from_namespace`:
3480/// the builder re-wraps an already-`Namespace`-wrapped error
3481/// (lance/src/dataset/builder.rs:142), so going through it would force a
3482/// chain-walk to classify `TableNotFound`. The direct probe stays at one
3483/// wrap level and downcasts cleanly. Managed-versioning hookup (REST
3484/// namespace external-manifest commits) is not wired here; v1 ships
3485/// Directory v2 only.
3486/// Diagnostic S3 IO tracing. Inert unless [`io_trace::enable`] is called
3487/// before the store opens; then a shared `IOTracker` is injected as the
3488/// object-store wrapper on every dataset read open, counting exactly how many
3489/// GETs (and bytes, and - under the `io-trace` feature - which paths) each
3490/// query issues against a remote store. Used by `serve_mem_bench --io-trace`
3491/// to attribute the per-query S3 request load. Not a production code path.
3492pub 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    /// Arm tracing. Must run before the store opens so the wrapper is applied
3499    /// when the datasets' object store is built.
3500    pub fn enable() {
3501        let _ = TRACKER.set(IOTracker::default());
3502    }
3503
3504    /// The shared tracker as an object-store wrapper, when armed.
3505    pub(super) fn wrapper() -> Option<Arc<IOTracker>> {
3506        TRACKER.get().map(|tracker| Arc::new(tracker.clone()))
3507    }
3508
3509    /// Read and reset the IO accumulated since the last call.
3510    pub fn take() -> Option<IoStats> {
3511        TRACKER.get().map(IOTracker::incremental_stats)
3512    }
3513}
3514
3515/// On-disk cache for `_indices/*` so a fresh process serves the IVF + FTS index
3516/// from local disk instead of re-loading it from the object store on every
3517/// cold-start (spec.md#search). Scoped to `_indices/*` because those files are
3518/// immutable and UUID-addressed, so a hit is always correct and a new index is
3519/// an automatic miss; data (served by the rowmap) and manifests (need freshness)
3520/// pass through. A `WrappingObjectStore`, so it stays inside the object-store
3521/// layer rather than reaching around it.
3522pub 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    /// Drop conditional headers (etag/if-modified): they reference the remote
3543    /// object and would spuriously fail against the local cache copy.
3544    fn local_opts(options: &GetOptions) -> GetOptions {
3545        GetOptions {
3546            range: options.range.clone(),
3547            head: options.head,
3548            ..Default::default()
3549        }
3550    }
3551
3552    /// `WrappingObjectStore` factory: holds the per-store cache root and hands a
3553    /// `CachingStore` to every dataset open on this store.
3554    #[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        /// `LocalFileSystem` requires the prefix to exist, so create it first.
3562        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        /// Fetch the whole object once, write it (`LocalFileSystem::put` stages +
3605        /// renames atomically), then serve the requested range from the copy. The
3606        /// per-path single-flight coalesces a process's concurrent first reads of
3607        /// one file into a single fetch; cross-process writes race safely since
3608        /// the bytes are identical and the rename is atomic.
3609        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            // Drop the entry so the map stays bounded as index versions churn.
3618            // Unconditionally safe (singleflight idiom): any waiter already holds
3619            // its own `lock` clone, and a later miss re-creates the entry but
3620            // finds the file cached.
3621            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            // Cache write or re-read failed (e.g. disk full): serve from origin.
3647            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                // Through get_opts so the first touch caches the whole object.
3690                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            // Delete both from the origin. The index file is served from the
3776            // local cache; the data file (never cached) is now gone.
3777            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            // A range read of the cached index slices the local copy.
3786            let slice = store.get_range(&index_path, 1..4).await.unwrap();
3787            assert_eq!(slice.as_ref(), b"NDE");
3788        }
3789    }
3790}
3791
3792/// fsync-on-write wrapper for local stores (spec.md#local-store-durability):
3793/// `LocalFileSystem` publishes a name (hard_link/rename) without syncing the
3794/// bytes, so a hard host stop can persist the name over page-cache-only bytes -
3795/// a zero-byte manifest that permanently poisons the table. This wrapper fsyncs
3796/// the written file and its parent directory after the inner write returns, so
3797/// every artifact is durable before Lance proceeds to the next step. Unix only:
3798/// Windows commits route through a different handler with different dir-fsync
3799/// semantics and rely on self-heal instead.
3800#[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    /// fsync the file the write just published plus its parent directory: the
3826    /// file `sync_all` makes the bytes durable, the dir `sync_all` makes the
3827    /// name (the freshly linked/renamed entry) durable. Fsync failures are hard
3828    /// errors - a silently no-op durability layer is worse than none - but a
3829    /// vanished parent dir (concurrent cleanup) is tolerated.
3830    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            // Unix permits fsync on an O_RDONLY directory fd.
3838            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    /// `WrappingObjectStore` factory: stateless, hands an `FsyncStore` to every
3850    /// dataset open on a local store.
3851    #[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        // Override so a rename keeps the inner store's native (atomic) semantics
3938        // and we fsync the destination; the trait default would degrade it to
3939        // copy+delete through our own `copy_opts`.
3940        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    /// Wraps the inner multipart upload so the final object (data and index
3953    /// files) is fsynced once `complete()` publishes it; parts and abort pass
3954    /// straight through.
3955    #[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        // A prefix-less local store, so an object_store `Path` is the absolute
3985        // FS path (leading `/` stripped) - exactly what `to_local_path` inverts.
3986        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            // object_store `Path` is the absolute FS path minus the leading `/`.
3993            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            // Round-trips through the wrapper...
4006            let got = store.get(&path).await.unwrap().bytes().await.unwrap();
4007            assert_eq!(got.as_ref(), b"DURABLE");
4008            // ...and the fsync targeted the real file `to_local_path` maps to.
4009            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
4036/// Stable filesystem-safe key for a store URL: same URL -> same key, so sibling
4037/// pond processes share one on-disk cache and distinct stores never collide.
4038/// Shared by the rowmap (`sessions.rs`), the index disk cache, and the CLI's
4039/// per-store sync lock / last-sync state files.
4040pub fn store_key(location: &Url) -> String {
4041    blake3::hash(location.as_str().as_bytes()).to_hex()[..16].to_owned()
4042}
4043
4044/// Reclaim cached `_indices/<uuid>` dirs whose UUID is not in `keep`. Recurses
4045/// to each `_indices` dir (the bucket prefix varies) and prunes its dead UUID
4046/// children. Best-effort; unlink-safe (POSIX keeps an in-flight reader's inode).
4047fn 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
4073/// The object-store wrapper applied to every dataset open: the fsync-on-write
4074/// durability wrapper (local stores, unix), the `_indices/*` disk cache (remote
4075/// stores only, when a cache dir is supplied), and the diagnostic io-trace
4076/// wrapper. `None` when none is active. The durability and index-cache wrappers
4077/// are backend-exclusive (local vs remote), so they never coexist.
4078fn 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    // Innermost, wrapping the real store directly: fsync must act on the final
4084    // on-disk file the moment the inner write publishes its name, before any
4085    // diagnostic wrapper's post-processing.
4086    #[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                    // A crashed local commit can leave a zero-byte/truncated head
4140                    // manifest that poisons the table permanently (spec.md#local-store-self-heal);
4141                    // self-heal by rolling back to the newest fully readable version.
4142                    // Remote stores never produce this (atomic PUT), so heal is local-only.
4143                    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                // fall through to create
4166            }
4167            _ => {
4168                return Err(anyhow::Error::from(error))
4169                    .with_context(|| format!("failed to describe table {table_name}"));
4170            }
4171        },
4172    }
4173
4174    // Create path: pond seeds an empty dataset with the canonical schema so
4175    // every subsequent open lands on a real Lance dataset, not a phantom.
4176    let mut write_params = sessions::write_params_for_create();
4177    write_params.session = Some(session.clone());
4178    write_params.mode = WriteMode::Create;
4179    // The wrapper must ride the create write too, or a local store's very first
4180    // manifest commit escapes fsync (local opens carry a wrapper but no
4181    // storage_options, so the old `!is_empty()` gate skipped it entirely).
4182    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
4199/// Apply the same session/wrapper/storage-option store params to a
4200/// `DatasetBuilder` that every pond open uses, so the heal probe and retry open
4201/// read the store identically to the real open.
4202fn 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
4224/// Lance's `_versions/` subdirectory name (lance-table commit.rs:70).
4225const VERSIONS_DIR_NAME: &str = "_versions";
4226/// Cap on scan-verify probes during a heal walk; a real crash leaves 1-2 bad
4227/// manifests, so a store needing more is pathological - fall through to the
4228/// enriched error rather than probe unboundedly.
4229const HEAL_MAX_PROBES: usize = 32;
4230
4231/// Parse a `_versions/` manifest filename to its version, following Lance's
4232/// `ManifestNamingScheme` (lance-table commit.rs:114-153): V2 is a 20-digit
4233/// `u64::MAX - version`, V1 is the plain version. Returns `None` for detached
4234/// (`d`-prefixed) manifests and for anything not ending in `.manifest` - which
4235/// skips prior `.corrupt` quarantines and Lance `.tmp_*` staging leftovers.
4236fn 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
4248/// Pin-open a specific version and drain a real scan over every column. The
4249/// pinned open resolves the manifest path deterministically and never lists
4250/// the directory or reads the poisoned head (lance builder.rs:239); draining
4251/// the scan forces data-page reads, which catches a zeroed data file that
4252/// manifest metadata alone hides. Full projection is load-bearing: a
4253/// column-update commit (embed's write shape) puts later-added columns in
4254/// their own per-fragment data files, which a narrower scan would never read
4255/// (spec.md#local-store-self-heal).
4256async 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
4279/// Self-heal a crash-damaged local table: walk `_versions/` head-down to the
4280/// newest fully readable version, quarantine the unreadable manifests above it
4281/// (atomic rename to `*.manifest.corrupt`, never delete), then retry the normal
4282/// open once. Lossless for pond: source histories are truth and the next
4283/// `pond sync` re-ingests the aborted commit (spec.md#local-store-self-heal).
4284/// When nothing is quarantinable, returns the original error enriched (Layer 3).
4285async 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    // Walk head-down to the newest fully readable version, collecting the
4340    // unreadable manifests above it. Rename nothing until a rollback target is
4341    // confirmed - a half-quarantine on a store we cannot repair is worse.
4342    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        // Lance's footer needs >=16 bytes (lance-io utils.rs:128); below that is
4348        // definitively unreadable, so quarantine it without a probe.
4349        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    // No readable version at all (or hit the probe cap): touch nothing.
4367    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    // The head itself is readable: the open failure is not manifest-shaped.
4378    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    // Quarantine each unreadable manifest above the rollback target: atomic
4390    // in-place rename, never delete. A lost-race rename (NotFound) means a
4391    // concurrent heal already moved it - treat as done and proceed.
4392    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    // Retry the normal open once; the head is now the rollback target.
4413    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    // Loud one-line notice: warns render on CLI stderr by default (main.rs
4426    // init_tracing defaults to WARN level).
4427    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
4436/// Layer 3: wrap the raw open error with what heal inspected and the concrete
4437/// recovery, so the caller sees a named fix instead of `Invalid range 0..0`.
4438/// Never quarantines (the raw error stays in the cause chain).
4439fn 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
4449// lance-namespace sometimes nests one `lance::Error::Namespace` inside another
4450// before the underlying `NamespaceError`; walk the whole `.source()` chain
4451// rather than only matching the outer variant.
4452fn 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}
4477/// How the on-disk schema relates to this build's expected schema.
4478enum SchemaFit {
4479    Match,
4480    /// Expected columns absent on disk, every one nullable: the store
4481    /// predates an additive schema change and upgrades in place.
4482    MissingNullable(Vec<lance::deps::arrow_schema::Field>),
4483    /// On-disk columns this build does not know: written by a newer pond.
4484    /// Reads proceed (scans project known columns; extra ones are inert);
4485    /// writes against the newer schema fail at the Lance layer, and the fix
4486    /// is upgrading pond, not editing the store.
4487    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
4527/// Open-time schema reconciliation: a store missing this build's known
4528/// nullable columns is backfilled IN PLACE via `Dataset::add_columns` - the
4529/// values derive from data already stored, so no re-ingest is ever required
4530/// (spec.md#session-durable-copy: a rotated source cannot supply rows again).
4531/// Concurrent openers race benignly: `add_columns` commits through OCC, a
4532/// losing writer sees a conflict, re-checks out latest, and finds the columns
4533/// present.
4534async 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        // Catch a vector-dim change (configured `[embeddings].dim` differs
4559        // from the on-disk vector column width) early with a friendly
4560        // message. Lance would otherwise reject the next write with an
4561        // opaque schema-mismatch error inside the `merge_update` path.
4562        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
4590/// One in-place additive migration pass: derive the missing columns from
4591/// stored data and commit them via `add_columns` (new column files only, no
4592/// row rewrites). The recipe - which columns to read and how to derive the
4593/// values - is consumer knowledge and lives in `sessions::column_backfill`.
4594async 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    // A full-column read over a remote store runs minutes; a silent stall
4603    // reads as a hang, so this one-time event gets a stderr notice (same
4604    // pattern as the embedding-model download notice in embed.rs).
4605    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    // Boxed: `add_columns`' concrete future is enormous (it embeds DataFusion
4611    // planner types), and inlining it here overflows rustc's auto-trait
4612    // solver (E0275) once this future nests inside the open chain.
4613    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                // Another writer migrated (or wrote) concurrently; re-check
4639                // out latest and let the caller re-classify.
4640                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}
4654/// Object-store defaults injected for any non-local pond location. Each key
4655/// is only set when neither the user-provided key nor its env-var-form alias
4656/// is already present, so explicit overrides in `[storage]` always win.
4657/// `aws_unsigned_payload` is gated on a custom endpoint (the marker for
4658/// S3-compatible stores like Hetzner, MinIO, R2), where the SHA256 payload
4659/// signature is wasted work the server does not validate.
4660fn 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    // `request_timeout` bounds a single object-store request (one range GET/PUT),
4673    // not a whole scan - a streaming read issues many small requests, each well
4674    // under this. We keep it deliberately tight as a HARD BARRIER: a single
4675    // request exceeding 60s means a design/infra problem to fix (chunk the read,
4676    // use change-data-feed, fix the endpoint), never something to paper over with
4677    // a longer timeout. An explicit `[storage]` override still wins.
4678    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        // Local file:// stores get the fsync durability wrapper (no index cache
4740        // dir, io-trace unarmed); remote stores get nothing here.
4741        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        // file (Lance's `uri_to_url` appends the trailing slash; `child_uri`
4768        // trims it downstream)
4769        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        // s3 passthrough
4774        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        // s3+https: TLS stays on, virtual-hosted defaults on for domain
4778        // hosts, region defaults deterministically. The endpoint is
4779        // assembled at resolve time with the bucket folded into the host
4780        // (object_store's virtual-hosted convention).
4781        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        // s3+http on an IP host: allow_http flips, path-style auto-selected
4798        // (a bucket subdomain on an IP can't resolve), port survives.
4799        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        // An explicit endpoint in `extra` is the escape hatch and wins.
4812        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        // gs passthrough
4831        let gcs = StorageUrl::parse("gs://bucket/p").unwrap();
4832        assert_eq!(gcs.lance_url().as_str(), "gs://bucket/p");
4833        // az: account folds into options
4834        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        // tests-only schemes pass through untouched
4841        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        // RFC 3986 userinfo is a leak class, never accepted.
4848        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        // Missing bucket.
4856        assert!(StorageUrl::parse("s3+https://host").is_err());
4857        assert!(StorageUrl::parse("az://acct").is_err());
4858        // Unknown scheme names the grammar.
4859        let err = StorageUrl::parse("ftp://host/x")
4860            .expect_err("ftp")
4861            .to_string();
4862        assert!(err.contains("s3+https"), "got: {err}");
4863        // Unrecognized query params die loudly.
4864        let err = StorageUrl::parse("s3://b/p?regoin=x")
4865            .expect_err("typo")
4866            .to_string();
4867        assert!(err.contains("regoin"), "got: {err}");
4868        // Query params on local / in-memory schemes die just as loudly -
4869        // no silent carry into the URL Lance opens.
4870        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        // `?` in a bare path is a filename character, not a query.
4879        assert!(StorageUrl::parse("/tmp/a?b").is_ok());
4880    }
4881
4882    #[test]
4883    fn storage_url_canonicalizes_ports_and_keeps_percent_encoding() {
4884        // Default port strips so scope matching can't split on `:443`.
4885        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        // Non-default port survives into the assembled endpoint.
4889        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        // Percent-encoded prefix passes through to the Lance URL verbatim.
4896        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        // Stripped before Lance sees the URL.
4916        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        // Assembly precedence: scheme < set < query.
4920        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        // virtual_hosted=true (query) -> the bucket rides in the endpoint host.
4926        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        // Longest match wins.
4947        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        // Segment boundary: `/pond` does not match `/pond-2`.
4962        assert_eq!(
4963            bind("s3+https://host/pond-2"),
4964            CredsBinding::Set {
4965                name: "all".to_owned(),
4966                via: BindVia::CatchAll
4967            },
4968        );
4969        // No cross-scheme normalization: the scoped sets don't match s3://.
4970        assert_eq!(
4971            bind("s3://pond/sub"),
4972            CredsBinding::Set {
4973                name: "all".to_owned(),
4974                via: BindVia::CatchAll
4975            },
4976        );
4977        // Default-port spelling matches the portless scope.
4978        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        // `?creds=` pointer beats every scope...
4986        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        // ...and a pointer to a missing set is an error, not a fallback.
4994        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        // No sets at all -> ambient chain; local URLs skip resolution.
5002        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        // A purely local invocation must not nag about remote-only sets.
5034        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                // Two trailing newlines: exactly one is stripped.
5052                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        // A failing command surfaces its text and exit status.
5068        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        // The command cache: one subprocess per command text per process.
5083        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        // Auth-class error with a bound set names the set...
5110        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        // ...and without one, points at the (empty) ambient chain.
5115        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        // Anything else is I/O, set or no set.
5128        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        // Lance wraps an empty-creds chain as a `Generic` error, never the
5137        // typed `Unauthenticated`; the rendered `CredentialsNotLoaded` is the
5138        // signal. Bound -> Auth (the set is wrong), unbound -> NoCreds.
5139        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        // The shape Lance actually produces: bug-report boilerplate, the real
5156        // cause, an internal source location, then the same text re-printed.
5157        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        // Display carries only the fix-naming lead, no chain.
5169        assert!(
5170            !failure.to_string().contains("file a bug report"),
5171            "lead must not trail the chain: {failure}"
5172        );
5173        // OccUnsupported's detail is already curated into Display.
5174        let occ = CheckFailure::OccUnsupported {
5175            detail: "put-if-none-match ignored".to_owned(),
5176        };
5177        assert!(occ.concise_cause().is_none());
5178        // Oversized single-line causes middle-truncate, keeping the tail
5179        // (wrapped transport errors put the root cause at the end).
5180        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        // One large candidate plus tiny appends -> vetoed.
5226        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        // Peer-sized candidates can merge and reach the target.
5229        let peers = [stat(100_000_000), stat(100_000_000)];
5230        assert!(task_is_kept(&peers, derived_target_rows(&peers)));
5231        // Remainder reaches largest / COMPACTION_ABSORB_FACTOR -> kept.
5232        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        // A 5 -> 4 rewrite still loops when all four outputs stay below target.
5320        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        // V1: plain version.
5349        assert_eq!(parse_manifest_version("5.manifest"), Some(5));
5350        assert_eq!(parse_manifest_version("0.manifest"), Some(0));
5351        // V2: 20-digit `u64::MAX - version`. u64::MAX renders as version 0.
5352        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        // Detached (`d`-prefixed) manifests are skipped.
5361        assert_eq!(parse_manifest_version("d123.manifest"), None);
5362        // Prior quarantines and Lance staging leftovers are skipped (not `.manifest`).
5363        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        // Unrelated files.
5369        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        // A column-update commit (embed's write shape) puts the new column in
5376        // its own per-fragment data file; a narrow projection would declare the
5377        // version healthy while that file is crash-zeroed.
5378        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        // Fresh session per probe so nothing is served from cache.
5423        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        // interval <= 1 always cleans (pond optimize / pond copy / tests).
5438        assert!(cleanup_due(0, 1));
5439        assert!(cleanup_due(7, 1));
5440        assert!(cleanup_due(5, 0));
5441        // interval N: only on multiples (the amortized pond sync path).
5442        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        // ~1.3 KiB rows -> ~100k-row target (half the byte budget, for freeze
5453        // headroom over the 256 MiB output cap).
5454        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        // No usable sizes -> Lance default.
5462        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        // Tiny rows clamp at the ceiling.
5472        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        // A row larger than the half-cap still needs a usable target.
5495        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            &not_namespace,
5525            ErrorCode::TableNotFound
5526        ));
5527    }
5528
5529    /// Round-trip: opening a fresh data dir through `lance-namespace`
5530    /// produces all three tables, and `Handle::scan` returns an empty batch
5531    /// for each (no spurious schema mismatch, no namespace error).
5532    #[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        // Each table has its own PK column; project the canonical one so the
5539        // scan is exercised end-to-end (catalog -> dataset -> scanner -> batch).
5540        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}