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 always run, bounding
837/// manifest growth even when the amplification veto would skip them. As
838/// policy cap, 0 disables the veto (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
846const MIN_TARGET_ROWS_PER_FRAGMENT: u64 = 50_000;
847/// Ceiling = Lance's own default.
848const MAX_TARGET_ROWS_PER_FRAGMENT: u64 = 1024 * 1024;
849
850/// Keep a task only when the merged-in remainder is >= largest/this:
851/// size-tiered amortization, O(log n) lifetime rewrites per row.
852pub const COMPACTION_ABSORB_FACTOR: u64 = 4;
853
854/// Default manifest-retention window for the safe cleanup pass. Matches
855/// LanceDB's recommended OSS-operator practice (lancedb docs: performance.mdx,
856/// tables/update.mdx). With `delete_unverified=false`, Lance's 7-day
857/// in-progress guard still protects unverified files regardless of this value
858/// (`UNVERIFIED_THRESHOLD_DAYS` in lance/dataset/cleanup.rs).
859pub fn default_cleanup_older_than() -> chrono::Duration {
860    // Toward Lance's 1 h floor: fewer retained manifest versions = cheaper
861    // remote open (spec.md#search). The append fast-path already curbs the
862    // version churn that earlier forced a wider window.
863    chrono::Duration::hours(1)
864}
865
866/// `pond sync` runs every few minutes; reclaiming old manifest versions on
867/// every run pays the full version-log walk over S3 (~9 s measured on the real
868/// corpus) to free roughly one version. Amortize by cleaning only when a
869/// table's manifest version is a multiple of this many commits. Explicit
870/// `pond optimize` and the one-shot `pond copy` keep interval 1 (clean every
871/// run) so maintenance and durability moves are never skipped.
872pub const DEFAULT_SYNC_CLEANUP_INTERVAL: u64 = 16;
873
874/// `pond sync` defers a scalar (BTree/bitmap) index fold until its unindexed
875/// tail reaches this many rows. Lance 7.0.0 ignores `OptimizeOptions::append()`
876/// for scalar indexes and rewrites the whole index file on every fold
877/// (O(index size), not O(delta)), so folding on every tiny sync pays a full
878/// rewrite for a handful of new rows. Batching amortizes that rewrite; the
879/// deferred tail stays correct for get/count/sql (they scan it) with scan cost
880/// bounded by this cap, and vector/FTS still fold every run so search recall is
881/// unaffected. `pond optimize`/`pond copy` fold every run (threshold `0`).
882pub const DEFAULT_SYNC_SCALAR_FOLD_ROWS: usize = 50_000;
883
884/// Defer the FTS + vector (IVF) index fold until the unindexed tail reaches this
885/// many rows; `0` folds every run. Unlike the scalar fold (deferred because Lance
886/// rewrites the whole index file), FTS/vector fold via a cheap delta append - the
887/// reason to batch them is the per-sync S3 round-trip + commit storm, not a
888/// rewrite. Between folds the tail stays fully searchable: the retrievers drop
889/// `fast_search` whenever an unindexed tail exists, so Lance index-probes the
890/// folded rows and flat-scans the (threshold-bounded) tail (fts.md "Index
891/// Maintenance"). The cap bounds that tail-scan cost. `pond optimize`/`pond copy`
892/// fold every run (threshold `0`).
893pub const DEFAULT_SYNC_INDEX_FOLD_ROWS: usize = 5_000;
894
895/// Resolved per-call inputs to the storage-maintenance pass. Built from
896/// `[maintenance]` (and any per-invocation CLI override) at the entry point;
897/// threaded down to `optimize_table_compact` so the substrate never re-reads
898/// `Config` itself.
899#[derive(Debug, Clone, Copy)]
900pub struct MaintenancePolicy {
901    /// See [`DEFAULT_COMPACTION_FRAGMENT_CAP`]; `0` disables the veto.
902    pub compaction_fragment_cap: usize,
903    /// Manifest-retention window handed to `cleanup_old_versions`.
904    pub cleanup_older_than: chrono::Duration,
905    /// Run `cleanup_old_versions` for a table only when its manifest version is
906    /// a multiple of this (`1` = every optimize). The frequent `pond sync` path
907    /// raises it so most syncs skip the version-log walk; see
908    /// [`DEFAULT_SYNC_CLEANUP_INTERVAL`].
909    pub cleanup_interval: u64,
910    /// Defer a scalar (BTree/bitmap) index fold until its unindexed tail reaches
911    /// this many rows; `0` folds every run. The frequent `pond sync` path raises
912    /// it so most syncs skip the full scalar-index rewrite Lance 7.0.0 does on
913    /// every fold; see [`DEFAULT_SYNC_SCALAR_FOLD_ROWS`].
914    pub scalar_fold_row_threshold: usize,
915    /// Defer the FTS + vector (IVF) index fold until its unindexed tail reaches
916    /// this many rows; `0` folds every run. The frequent `pond sync` path raises
917    /// it so most syncs skip the per-fold S3 round-trip storm; recall stays
918    /// complete because the retrievers flat-scan the deferred tail. See
919    /// [`DEFAULT_SYNC_INDEX_FOLD_ROWS`].
920    pub index_fold_row_threshold: usize,
921}
922
923impl MaintenancePolicy {
924    /// Veto off: run every task Lance plans (the optimize tests assume this).
925    pub fn always_compact() -> Self {
926        Self {
927            compaction_fragment_cap: 0,
928            cleanup_older_than: default_cleanup_older_than(),
929            cleanup_interval: 1,
930            scalar_fold_row_threshold: 0,
931            index_fold_row_threshold: 0,
932        }
933    }
934
935    /// Amortize version cleanup over `interval` commits - the frequent
936    /// `pond sync` path uses this so most syncs skip the version-log walk.
937    #[must_use]
938    pub fn with_cleanup_interval(mut self, interval: u64) -> Self {
939        self.cleanup_interval = interval.max(1);
940        self
941    }
942
943    /// Amortize the scalar-index fold over its unindexed tail - the frequent
944    /// `pond sync` path uses this so most syncs skip the full scalar-index
945    /// rewrite Lance 7.0.0 does on every fold.
946    #[must_use]
947    pub fn with_scalar_fold_row_threshold(mut self, threshold: usize) -> Self {
948        self.scalar_fold_row_threshold = threshold;
949        self
950    }
951
952    /// Amortize the FTS + vector fold over its unindexed tail - the frequent
953    /// `pond sync` path uses this so most syncs skip the per-fold S3 round-trip
954    /// storm; recall stays complete via the retrievers' tail flat-scan.
955    #[must_use]
956    pub fn with_index_fold_row_threshold(mut self, threshold: usize) -> Self {
957        self.index_fold_row_threshold = threshold;
958        self
959    }
960
961    /// The two per-family fold thresholds bundled for the indices phase, so the
962    /// two same-typed `usize`s can't be swapped at a call site.
963    fn fold_thresholds(&self) -> FoldThresholds {
964        FoldThresholds {
965            scalar: self.scalar_fold_row_threshold,
966            index: self.index_fold_row_threshold,
967        }
968    }
969}
970
971/// Per-index-family fold-deferral thresholds (rows); `0` folds that family every
972/// run. Bundled so the indices phase takes one param, not two swappable `usize`s.
973#[derive(Debug, Clone, Copy)]
974struct FoldThresholds {
975    scalar: usize,
976    index: usize,
977}
978
979struct FragmentStat {
980    /// `None` when the manifest lacks any file's size.
981    bytes: Option<u64>,
982    rows: u64,
983    deleted_rows: u64,
984}
985
986/// Data-file bytes of one fragment; `None` (poisoning) when any size is
987/// missing from the manifest.
988fn fragment_bytes(fragment: &lance::table::format::Fragment) -> Option<u64> {
989    fragment.files.iter().try_fold(0u64, |total, file| {
990        Some(total + file.file_size_bytes.get()?.get())
991    })
992}
993
994fn fragment_stat(fragment: &lance::table::format::Fragment) -> FragmentStat {
995    FragmentStat {
996        bytes: fragment_bytes(fragment),
997        rows: fragment.physical_rows.unwrap_or(0) as u64,
998        deleted_rows: fragment
999            .deletion_file
1000            .as_ref()
1001            .and_then(|deletions| deletions.num_deleted_rows)
1002            .unwrap_or(0) as u64,
1003    }
1004}
1005
1006/// Candidacy/merge target: HALF the rows a [`TARGET_FRAGMENT_BYTES`] fragment
1007/// holds at the table's average row size. Compaction byte-caps every output
1008/// fragment at [`TARGET_FRAGMENT_BYTES`] (`max_bytes_per_file`), so deriving the
1009/// target at the FULL byte budget made `target == the largest fragment
1010/// compaction can produce`: no output could ever satisfy `physical_rows >=
1011/// target`, so the table was re-compacted every sync for a net-zero fragment
1012/// change (measured ~100-120s/sync on the remote store, 30->30 fragments).
1013/// Halving leaves 2x headroom so a byte-capped fragment lands comfortably above
1014/// the target and FREEZES, making compaction productive (merge small -> freeze
1015/// -> stop) instead of perpetual churn.
1016fn derived_target_rows(stats: &[FragmentStat]) -> usize {
1017    let (mut bytes, mut rows) = (0u64, 0u64);
1018    for stat in stats {
1019        if let Some(fragment_bytes) = stat.bytes
1020            && stat.rows > 0
1021        {
1022            bytes += fragment_bytes;
1023            rows += stat.rows;
1024        }
1025    }
1026    if bytes == 0 || rows == 0 {
1027        return MAX_TARGET_ROWS_PER_FRAGMENT as usize;
1028    }
1029    let avg_row_bytes = (bytes / rows).max(1);
1030    (TARGET_FRAGMENT_BYTES / 2 / avg_row_bytes)
1031        .clamp(MIN_TARGET_ROWS_PER_FRAGMENT, MAX_TARGET_ROWS_PER_FRAGMENT) as usize
1032}
1033
1034/// Amplification veto: skip tasks that mostly rewrite one big fragment to
1035/// absorb fresh appends. Deletion-materialization tasks always pass (vetoing
1036/// them would leave tombstones unreclaimed forever); compared in bytes when
1037/// every file size is known, rows otherwise.
1038fn keep_task(stats: &[FragmentStat], cap: usize, deletion_threshold: f32) -> bool {
1039    if stats.iter().any(|stat| {
1040        stat.rows > 0 && (stat.deleted_rows as f32 / stat.rows as f32) > deletion_threshold
1041    }) {
1042        return true;
1043    }
1044    if stats.len() >= cap {
1045        return true;
1046    }
1047    let weights: Vec<u64> = if stats.iter().all(|stat| stat.bytes.is_some()) {
1048        stats.iter().filter_map(|stat| stat.bytes).collect()
1049    } else {
1050        stats.iter().map(|stat| stat.rows).collect()
1051    };
1052    let total: u64 = weights.iter().sum();
1053    let largest = weights.iter().copied().max().unwrap_or(0);
1054    (total - largest) * COMPACTION_ABSORB_FACTOR >= largest
1055}
1056
1057/// Declarative description of one index pond keeps on a table. Created when
1058/// its trigger fires; folded forward by `pond optimize`.
1059#[derive(Debug, Clone)]
1060pub struct IndexIntent {
1061    /// Stable on-disk name. Must match across runs so existence checks
1062    /// resolve.
1063    pub name: &'static str,
1064    /// Column the index covers.
1065    pub column: &'static str,
1066    /// Condition evaluated against the live dataset before each cycle.
1067    pub trigger: IndexTrigger,
1068    /// How the params are built at create time. Some intents have static
1069    /// params (FTS, scalars); IVF_SQ needs the row count to size partitions.
1070    pub params: IndexParamsKind,
1071}
1072
1073/// When an [`IndexIntent`] should exist on disk.
1074#[derive(Debug, Clone)]
1075pub enum IndexTrigger {
1076    /// Build whenever the table has any rows. Used for FTS and scalar
1077    /// indices: there is no training cost worth delaying.
1078    OnAnyRows,
1079    /// Build when `count(<column> IS NOT NULL) >= threshold`. Used for the
1080    /// IVF_SQ vector index, which trains poorly on too few vectors.
1081    OnNonNullCount {
1082        column: &'static str,
1083        threshold: usize,
1084    },
1085}
1086
1087/// The lance-native shape of an [`IndexIntent`]'s params, dispatched to the
1088/// right `IndexParams` at create time.
1089#[derive(Debug, Clone)]
1090pub enum IndexParamsKind {
1091    /// `BuiltinIndexType::BTree` -> [`IndexType::BTree`];
1092    /// `BuiltinIndexType::Bitmap` -> [`IndexType::Bitmap`]; etc.
1093    Scalar(BuiltinIndexType),
1094    /// `InvertedIndexParams` with the word-level `simple` tokenizer plus
1095    /// English stemming, stop-words off (spec.md#search-language-neutral-index).
1096    /// Word retrieval beats character ngram ~2x on the real corpus at ~4x less
1097    /// index weight; substring/symbol lookup stays on the SQL `LIKE` /
1098    /// `contains_tokens` path, not here.
1099    InvertedFtsWord,
1100    /// `VectorIndexParams::with_ivf_sq_params` with cosine metric (e5 vectors
1101    /// are L2-normalized). 8-bit scalar quantization stores per-dimension codes
1102    /// in the index itself, so kNN computes distances from the prewarmed
1103    /// partition with no refine pass - PQ+refine instead re-reads ~k*factor
1104    /// exact vectors from the data files as scattered per-row GETs, the
1105    /// dominant per-query S3 request storm on a throttling remote store
1106    /// (spec.md#search). `max_iters` caps kmeans; partitions follow LanceDB's
1107    /// documented `num_rows // 4096` guidance, floored at one.
1108    IvfSqCosine { num_bits: u16, max_iters: usize },
1109}
1110
1111impl IndexTrigger {
1112    async fn should_create(&self, dataset: &Dataset) -> Result<bool> {
1113        match self {
1114            Self::OnAnyRows => Ok(dataset.count_rows(None).await? > 0),
1115            Self::OnNonNullCount { column, threshold } => {
1116                let count = dataset
1117                    .count_rows(Some(format!("{column} IS NOT NULL")))
1118                    .await?;
1119                Ok(count >= *threshold)
1120            }
1121        }
1122    }
1123}
1124
1125impl IndexParamsKind {
1126    fn index_type(&self) -> IndexType {
1127        match self {
1128            Self::Scalar(BuiltinIndexType::Bitmap) => IndexType::Bitmap,
1129            Self::Scalar(BuiltinIndexType::ZoneMap) => IndexType::ZoneMap,
1130            Self::Scalar(_) => IndexType::BTree,
1131            Self::InvertedFtsWord => IndexType::Inverted,
1132            Self::IvfSqCosine { .. } => IndexType::Vector,
1133        }
1134    }
1135
1136    async fn build(&self, dataset: &Dataset) -> Result<Box<dyn lance::index::IndexParams>> {
1137        match self {
1138            Self::Scalar(kind) => Ok(Box::new(ScalarIndexParams::for_builtin(kind.clone()))),
1139            Self::InvertedFtsWord => Ok(Box::new(
1140                InvertedIndexParams::default()
1141                    .base_tokenizer("simple".to_owned())
1142                    .stem(true)
1143                    .remove_stop_words(false),
1144            )),
1145            Self::IvfSqCosine {
1146                num_bits,
1147                max_iters,
1148            } => {
1149                let count = dataset
1150                    .count_rows(Some("vector IS NOT NULL".to_owned()))
1151                    .await?;
1152                let partitions = count.checked_div(4096).unwrap_or(0).max(1);
1153                let mut ivf = IvfBuildParams::new(partitions);
1154                ivf.max_iters = *max_iters;
1155                let sq = SQBuildParams {
1156                    num_bits: *num_bits,
1157                    ..Default::default()
1158                };
1159                Ok(Box::new(VectorIndexParams::with_ivf_sq_params(
1160                    MetricType::Cosine,
1161                    ivf,
1162                    sq,
1163                )))
1164            }
1165        }
1166    }
1167}
1168
1169#[derive(Debug, Clone, PartialEq, Eq)]
1170pub struct IndexStatus {
1171    pub table: Table,
1172    pub intent_name: String,
1173    pub fragments_covered: usize,
1174    pub unindexed_fragments: usize,
1175    pub unindexed_rows: usize,
1176    pub exists: bool,
1177}
1178
1179/// Anyhow-chain sentinel pond attaches when `retry_lance` exhausts attempts
1180/// against an OCC commit-conflict failure (spec.md#protocol). The wire layer
1181/// downcasts to this type to classify the outcome as `conflict` rather than
1182/// the generic `storage_unavailable`.
1183#[derive(Debug, Clone, Copy)]
1184pub struct ConflictExhausted {
1185    pub attempts: u8,
1186}
1187
1188impl std::fmt::Display for ConflictExhausted {
1189    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1190        write!(
1191            formatter,
1192            "commit conflict exhausted after {} attempt(s)",
1193            self.attempts
1194        )
1195    }
1196}
1197
1198impl std::error::Error for ConflictExhausted {}
1199
1200/// Per-phase result for one table's pass through `Handle::optimize_table`.
1201/// spec.md#substrate 3.7 (`lance-index-maintenance`): the indices phase and the
1202/// compaction phase get independent retry budgets and independent commits,
1203/// so a hot writer that starves the Rewrite cannot abort the index Update.
1204#[derive(Debug)]
1205pub enum PhaseOutcome {
1206    /// Phase attempted and committed work.
1207    Ok,
1208    /// Phase attempted; no work was needed.
1209    Noop,
1210    /// Phase attempted; OCC retry budget exhausted on conflict (the operator
1211    /// can rerun later once the hot writer quiesces).
1212    SkippedConflict,
1213    /// Phase failed with a non-conflict error.
1214    Failed(anyhow::Error),
1215    /// Phase not requested by the caller (e.g. compaction skipped under
1216    /// `Store::build_indices_only`).
1217    NotAttempted,
1218}
1219
1220impl PhaseOutcome {
1221    pub fn is_failed(&self) -> bool {
1222        matches!(self, Self::Failed(_))
1223    }
1224}
1225
1226/// What `Handle::optimize_table` did for one table.
1227#[derive(Debug)]
1228pub struct TableOptimizeOutcome {
1229    pub table: Table,
1230    pub indices: PhaseOutcome,
1231    pub compaction: PhaseOutcome,
1232}
1233
1234/// Boundary event during one `Handle::optimize_table` pass. The CLI binds a
1235/// progress callback to render a live spinner; library callers pass `None`.
1236#[derive(Debug, Clone)]
1237pub enum OptimizeEvent {
1238    PhaseStart {
1239        table: Table,
1240        phase: OptimizePhase,
1241        detail: Option<String>,
1242    },
1243    PhaseDone {
1244        table: Table,
1245        phase: OptimizePhase,
1246        elapsed_ms: u64,
1247    },
1248    /// Intra-index liveness, forwarded from Lance's `IndexBuildProgress`
1249    /// callbacks (FTS tokenize/copy, IVF train/shuffle/merge, BTree/Bitmap
1250    /// build stages). Fires many times per index between `PhaseStart` /
1251    /// `PhaseDone`; the spinner just overwrites its message each tick.
1252    IndexStage {
1253        table: Table,
1254        index: String,
1255        stage: String,
1256        completed: u64,
1257        total: Option<u64>,
1258        unit: String,
1259    },
1260}
1261
1262#[derive(Debug, Clone, Copy)]
1263pub enum OptimizePhase {
1264    Compact,
1265    Cleanup,
1266    IndexCreate,
1267    IndexRebuild,
1268    IndexAppend,
1269}
1270
1271impl OptimizePhase {
1272    pub fn label(self) -> &'static str {
1273        match self {
1274            Self::Compact => "compact",
1275            Self::Cleanup => "cleanup",
1276            Self::IndexCreate => "index-create",
1277            Self::IndexRebuild => "index-rebuild",
1278            Self::IndexAppend => "index-append",
1279        }
1280    }
1281}
1282
1283/// `Arc` rather than `Box` so the same callback can be cloned into the
1284/// `PondIndexProgress` Arc that Lance's `IndexBuildProgress` builder demands -
1285/// otherwise intra-index stage events have no path back to the CLI spinner.
1286pub type OptimizeProgressFn = Arc<dyn Fn(OptimizeEvent) + Send + Sync>;
1287
1288fn emit(progress: Option<&OptimizeProgressFn>, event: OptimizeEvent) {
1289    if let Some(callback) = progress {
1290        callback(event);
1291    }
1292}
1293
1294/// Bridges Lance's `IndexBuildProgress` async callbacks (`stage_start`,
1295/// `stage_progress`, `stage_complete`) into pond's `OptimizeEvent::IndexStage`
1296/// stream so the CLI spinner can show "fts tokenize_docs 1.4M / 2M rows"
1297/// instead of going dark for 10-20 minutes during a single `create_index` or
1298/// `optimize_indices` call. Remembers the active stage's `total` / `unit` so
1299/// `stage_progress` (which only carries `completed`) can render a full
1300/// fraction. Emissions are throttled to one every 100ms; FTS's per-batch
1301/// `stage_progress` calls would otherwise contend the spinner mutex.
1302struct PondIndexProgress {
1303    callback: OptimizeProgressFn,
1304    table: Table,
1305    index: String,
1306    state: std::sync::Mutex<PondIndexStageState>,
1307}
1308
1309// `IndexBuildProgress` requires `Debug`; the `callback` field is
1310// `Arc<dyn Fn...>` which has no `Debug` impl, so derive doesn't apply.
1311impl std::fmt::Debug for PondIndexProgress {
1312    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1313        f.debug_struct("PondIndexProgress")
1314            .field("table", &self.table)
1315            .field("index", &self.index)
1316            .finish_non_exhaustive()
1317    }
1318}
1319
1320#[derive(Debug, Default)]
1321struct PondIndexStageState {
1322    total: Option<u64>,
1323    unit: String,
1324    last_emit: Option<Instant>,
1325}
1326
1327impl PondIndexProgress {
1328    fn new(callback: OptimizeProgressFn, table: Table, index: String) -> Arc<Self> {
1329        Arc::new(Self {
1330            callback,
1331            table,
1332            index,
1333            state: std::sync::Mutex::new(PondIndexStageState::default()),
1334        })
1335    }
1336}
1337
1338#[async_trait::async_trait]
1339impl lance_index::progress::IndexBuildProgress for PondIndexProgress {
1340    async fn stage_start(&self, stage: &str, total: Option<u64>, unit: &str) -> lance::Result<()> {
1341        if let Ok(mut state) = self.state.lock() {
1342            state.total = total;
1343            state.unit = unit.to_owned();
1344            state.last_emit = Some(Instant::now());
1345        }
1346        (self.callback)(OptimizeEvent::IndexStage {
1347            table: self.table,
1348            index: self.index.clone(),
1349            stage: stage.to_owned(),
1350            completed: 0,
1351            total,
1352            unit: unit.to_owned(),
1353        });
1354        Ok(())
1355    }
1356
1357    async fn stage_progress(&self, stage: &str, completed: u64) -> lance::Result<()> {
1358        let (total, unit) = {
1359            let Ok(mut state) = self.state.lock() else {
1360                return Ok(());
1361            };
1362            let now = Instant::now();
1363            if let Some(prev) = state.last_emit
1364                && now.duration_since(prev) < Duration::from_millis(100)
1365            {
1366                return Ok(());
1367            }
1368            state.last_emit = Some(now);
1369            (state.total, state.unit.clone())
1370        };
1371        (self.callback)(OptimizeEvent::IndexStage {
1372            table: self.table,
1373            index: self.index.clone(),
1374            stage: stage.to_owned(),
1375            completed,
1376            total,
1377            unit,
1378        });
1379        Ok(())
1380    }
1381
1382    async fn stage_complete(&self, stage: &str) -> lance::Result<()> {
1383        let (total, unit) = {
1384            let Ok(state) = self.state.lock() else {
1385                return Ok(());
1386            };
1387            (state.total, state.unit.clone())
1388        };
1389        (self.callback)(OptimizeEvent::IndexStage {
1390            table: self.table,
1391            index: self.index.clone(),
1392            stage: stage.to_owned(),
1393            completed: total.unwrap_or(0),
1394            total,
1395            unit,
1396        });
1397        Ok(())
1398    }
1399}
1400
1401fn lance_progress(
1402    progress: Option<&OptimizeProgressFn>,
1403    table: Table,
1404    index: &str,
1405) -> Arc<dyn lance_index::progress::IndexBuildProgress> {
1406    match progress {
1407        Some(callback) => PondIndexProgress::new(callback.clone(), table, index.to_owned()),
1408        None => Arc::new(lance_index::progress::NoopIndexBuildProgress),
1409    }
1410}
1411
1412/// True when the chain root is one of Lance's commit-conflict variants
1413/// (`CommitConflict`, `RetryableCommitConflict`, `TooMuchWriteContention`).
1414/// Everything else (timeouts, IAM denials, disk errors) is not a conflict.
1415pub fn is_commit_conflict(error: &anyhow::Error) -> bool {
1416    error.downcast_ref::<lance::Error>().is_some_and(|err| {
1417        matches!(
1418            err,
1419            lance::Error::CommitConflict { .. }
1420                | lance::Error::RetryableCommitConflict { .. }
1421                | lance::Error::TooMuchWriteContention { .. }
1422        )
1423    })
1424}
1425
1426/// True when `retry_lance` exhausted retries against an OCC conflict and
1427/// attached `ConflictExhausted` to the chain head.
1428fn is_conflict_exhausted(error: &anyhow::Error) -> bool {
1429    error.chain().any(|cause| cause.is::<ConflictExhausted>())
1430}
1431
1432/// True when the chain root is Lance's `Index` error class - a structural
1433/// index fault (e.g. delta segments with mismatched posting tail codecs) that
1434/// retry cannot clear and only a from-scratch rebuild repairs.
1435pub fn is_index_error(error: &anyhow::Error) -> bool {
1436    error
1437        .downcast_ref::<lance::Error>()
1438        .is_some_and(|err| matches!(err, lance::Error::Index { .. }))
1439}
1440
1441/// On-disk byte totals for the three session datasets, plus everything else
1442/// under the data-dir root. Sized by listing through Lance's object-store
1443/// layer (spec.md#lance-chokepoints-storage) so `file://` and `s3://` behave alike.
1444#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1445pub struct TableSizes {
1446    pub sessions: u64,
1447    pub messages: u64,
1448    pub parts: u64,
1449    pub other: u64,
1450    pub sessions_data: DataLiveness,
1451    pub messages_data: DataLiveness,
1452    pub parts_data: DataLiveness,
1453}
1454
1455/// `data/` bytes on disk vs bytes the latest manifest references; the gap is
1456/// superseded versions awaiting the cleanup retention window.
1457#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1458pub struct DataLiveness {
1459    pub on_disk: u64,
1460    /// `None` when the manifest lacks any referenced file's size.
1461    pub live: Option<u64>,
1462}
1463
1464impl DataLiveness {
1465    pub fn dead(&self) -> Option<u64> {
1466        self.live.map(|live| self.on_disk.saturating_sub(live))
1467    }
1468}
1469
1470#[derive(Debug, Clone, PartialEq, Eq)]
1471pub enum ScalarValue {
1472    String(String),
1473    Int32(i32),
1474    Raw(String),
1475}
1476impl From<&str> for ScalarValue {
1477    fn from(value: &str) -> Self {
1478        Self::String(value.to_owned())
1479    }
1480}
1481impl From<String> for ScalarValue {
1482    fn from(value: String) -> Self {
1483        Self::String(value)
1484    }
1485}
1486impl From<i32> for ScalarValue {
1487    fn from(value: i32) -> Self {
1488        Self::Int32(value)
1489    }
1490}
1491#[derive(Debug, Clone, PartialEq, Eq)]
1492pub enum Predicate {
1493    Eq(&'static str, ScalarValue),
1494    Ne(&'static str, ScalarValue),
1495    IsNull(&'static str),
1496    IsNotNull(&'static str),
1497    In(&'static str, Vec<ScalarValue>),
1498    LikeContains(&'static str, String),
1499    /// Regex match. Emitted as `regexp_like(<col>, '<pat>')`. Never pushes
1500    /// down to BTREE indexes (Lance's scalar-index-expr parser ignores it),
1501    /// so the filter is a full-scan-with-predicate - acceptable for
1502    /// human-driven `--project re:...` queries, not for hot paths.
1503    Regex(&'static str, String),
1504    Gte(&'static str, ScalarValue),
1505    Lte(&'static str, ScalarValue),
1506    And(Vec<Predicate>),
1507    Or(Vec<Predicate>),
1508    Not(Box<Predicate>),
1509}
1510impl Predicate {
1511    pub fn to_lance(&self) -> String {
1512        match self {
1513            Self::Eq(column, value) => format!("{column} = {}", value.to_lance()),
1514            Self::Ne(column, value) => format!("{column} <> {}", value.to_lance()),
1515            Self::IsNull(column) => format!("{column} IS NULL"),
1516            Self::IsNotNull(column) => format!("{column} IS NOT NULL"),
1517            Self::In(column, values) => {
1518                let values = values
1519                    .iter()
1520                    .map(ScalarValue::to_lance)
1521                    .collect::<Vec<_>>()
1522                    .join(", ");
1523                format!("{column} IN ({values})")
1524            }
1525            Self::LikeContains(column, value) => {
1526                format!("{column} LIKE {} ESCAPE '\\'", like_contains(value))
1527            }
1528            Self::Regex(column, pattern) => {
1529                format!("regexp_like({column}, {})", quoted_string(pattern))
1530            }
1531            Self::Gte(column, value) => format!("{column} >= {}", value.to_lance()),
1532            Self::Lte(column, value) => format!("{column} <= {}", value.to_lance()),
1533            Self::And(predicates) => predicates
1534                .iter()
1535                .map(Self::to_lance)
1536                .filter(|predicate| !predicate.is_empty())
1537                .collect::<Vec<_>>()
1538                .join(" AND "),
1539            Self::Or(predicates) => {
1540                // Wrap in parens so the disjunction composes safely as a child
1541                // of an outer `And` (SQL `OR` binds looser than `AND`).
1542                let body = predicates
1543                    .iter()
1544                    .map(Self::to_lance)
1545                    .filter(|predicate| !predicate.is_empty())
1546                    .collect::<Vec<_>>()
1547                    .join(" OR ");
1548                if body.is_empty() {
1549                    String::new()
1550                } else {
1551                    format!("({body})")
1552                }
1553            }
1554            Self::Not(inner) => {
1555                let body = inner.to_lance();
1556                if body.is_empty() {
1557                    String::new()
1558                } else {
1559                    format!("NOT ({body})")
1560                }
1561            }
1562        }
1563    }
1564}
1565/// Read-side options for `Handle::scan`: optional prefilter predicate and
1566/// optional projection. Default = no filter, all columns.
1567#[derive(Default)]
1568pub struct ScanOpts<'a> {
1569    pub predicate: Option<&'a Predicate>,
1570    pub projection: Option<&'a [&'a str]>,
1571}
1572
1573impl<'a> ScanOpts<'a> {
1574    pub fn project_only(projection: &'a [&'a str]) -> Self {
1575        Self {
1576            predicate: None,
1577            projection: Some(projection),
1578        }
1579    }
1580    pub fn with_predicate_and_projection(
1581        predicate: &'a Predicate,
1582        projection: &'a [&'a str],
1583    ) -> Self {
1584        Self {
1585            predicate: Some(predicate),
1586            projection: Some(projection),
1587        }
1588    }
1589}
1590
1591impl ScalarValue {
1592    fn to_lance(&self) -> String {
1593        match self {
1594            Self::String(value) => quoted_string(value),
1595            Self::Int32(value) => value.to_string(),
1596            Self::Raw(value) => value.clone(),
1597        }
1598    }
1599}
1600/// Lance cache caps in bytes. `None` lets the substrate pick the backend-aware
1601/// default (local FS gets a tighter cap; object stores stay near Lance's
1602/// defaults). Wired through `Store::open_with_options` from `[runtime]`.
1603#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
1604pub struct RuntimeCaps {
1605    pub index_cache_bytes: Option<usize>,
1606    pub metadata_cache_bytes: Option<usize>,
1607}
1608
1609impl RuntimeCaps {
1610    pub fn from_config(config: &crate::config::RuntimeConfig) -> Self {
1611        Self {
1612            index_cache_bytes: config.index_cache_bytes,
1613            metadata_cache_bytes: config.metadata_cache_bytes,
1614        }
1615    }
1616}
1617
1618/// Local-FS default: tight enough that a long-lived `pond mcp` lands well
1619/// under the 500 MiB target without measurable latency cost vs Lance's 6 GiB
1620/// default (see `benches/serve_mem_bench.rs --cap-sweep`).
1621const LOCAL_INDEX_CACHE_BYTES: usize = 256 * 1024 * 1024;
1622const LOCAL_METADATA_CACHE_BYTES: usize = 128 * 1024 * 1024;
1623/// Object-store defaults: latency to refill is per-page, so keep more in cache
1624/// than local - but bounded above the warm working set, not Lance's 6 GiB.
1625/// Post word-tokenizer FTS that set is ~450 MB (simple invert + IVF_SQ aux), so
1626/// 1 GiB holds both indices warm with headroom while capping the RSS ceiling.
1627const REMOTE_INDEX_CACHE_BYTES: usize = 1024 * 1024 * 1024;
1628const REMOTE_METADATA_CACHE_BYTES: usize = 512 * 1024 * 1024;
1629
1630fn resolve_cache_caps(location: &Url, caps: RuntimeCaps) -> (usize, usize) {
1631    let (index_default, metadata_default) = if config::is_local(location) {
1632        (LOCAL_INDEX_CACHE_BYTES, LOCAL_METADATA_CACHE_BYTES)
1633    } else {
1634        (REMOTE_INDEX_CACHE_BYTES, REMOTE_METADATA_CACHE_BYTES)
1635    };
1636    (
1637        caps.index_cache_bytes.unwrap_or(index_default),
1638        caps.metadata_cache_bytes.unwrap_or(metadata_default),
1639    )
1640}
1641
1642pub struct Handle {
1643    datasets: DatasetSet,
1644    retry: RetryPolicy,
1645    /// One `lance::Session` shared across all three datasets. Carries the
1646    /// metadata + index caches and the `ObjectStoreRegistry` (which holds
1647    /// the underlying object_store / S3 client). Sharing the session means
1648    /// one cache pool covers all three tables and one S3 client serves all
1649    /// three datasets - load-bearing on object-store backends where a
1650    /// per-dataset client would mean 3x the connection pools and 3x the
1651    /// credential refreshes (lance/src/dataset/builder.rs:509-517).
1652    #[allow(dead_code)]
1653    session: Arc<Session>,
1654    /// The `lance-namespace` catalog seam. v1 uses the Directory impl;
1655    /// future hosted pond swaps to "rest" without touching read/write paths
1656    /// (spec.md#lance-chokepoints-catalog).
1657    nm: Arc<dyn LanceNamespace>,
1658    /// Namespace identifier this handle binds to. v1 is always `root()`; the
1659    /// typed seam matches `resolve_namespace`'s return so multi-namespace
1660    /// routing can land without churning call sites (spec.md#wire-namespace-resolution).
1661    nm_ident: NamespaceIdent,
1662    /// Object-store options threaded through every `DatasetBuilder` and
1663    /// `Dataset::write` call so refresh / index-creation paths inherit the
1664    /// same credentials and region as the initial open. Empty on local-FS
1665    /// installs.
1666    storage_options: HashMap<String, String>,
1667    /// Data-dir URL the handle was opened against. `pond status` reads this
1668    /// to display where the bytes live and to decide whether to walk a local
1669    /// directory or issue a remote `LIST` for sizing.
1670    location: Url,
1671    /// Freshness window applied to the lazily-opened `sessions` and `parts`
1672    /// datasets when they first open, matching the eager `messages` open's
1673    /// scheme-keyed `refresh_after`.
1674    lazy_refresh_after: Duration,
1675    /// Object-store wrapper (index disk cache + io-trace) applied on every
1676    /// dataset open, including the lazy sessions/parts opens and any re-open.
1677    index_wrapper: Option<Arc<dyn WrappingObjectStore>>,
1678}
1679
1680impl std::fmt::Debug for Handle {
1681    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1682        formatter
1683            .debug_struct("Handle")
1684            .field("datasets", &self.datasets)
1685            .field("retry", &self.retry)
1686            .field("nm_ident", &self.nm_ident)
1687            .field("storage_options", &self.storage_options)
1688            .field("location", &self.location)
1689            .finish()
1690    }
1691}
1692
1693#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1694pub enum Table {
1695    Sessions,
1696    Messages,
1697    Parts,
1698}
1699impl Table {
1700    pub fn as_str(self) -> &'static str {
1701        self.label()
1702    }
1703
1704    fn label(self) -> &'static str {
1705        match self {
1706            Self::Sessions => "sessions",
1707            Self::Messages => "messages",
1708            Self::Parts => "parts",
1709        }
1710    }
1711}
1712#[derive(Debug)]
1713struct DatasetSet {
1714    /// `sessions.lance` opens lazily, like `parts`: the search request path
1715    /// reads only `messages`. Writers (ingest), `pond status`, restore, and the
1716    /// daemon's background index-cache GC open it on first use.
1717    sessions: OnceCell<Mutex<CachedDataset>>,
1718    messages: Mutex<CachedDataset>,
1719    /// `parts.lance` opens lazily on the first read or write that needs it:
1720    /// any `pond_get` (every mode reads parts to build summaries), grouped
1721    /// search hydrating user-hit summaries, or ingest with Part events. A
1722    /// process that does none of those skips the file, saving its metadata
1723    /// pages and file handle at cold-open. The OnceCell makes init
1724    /// single-flight; the inner `Mutex<CachedDataset>` then behaves identically
1725    /// to the other two.
1726    parts: OnceCell<Mutex<CachedDataset>>,
1727}
1728#[derive(Debug)]
1729struct CachedDataset {
1730    dataset: Dataset,
1731    last_refresh: Instant,
1732    refresh_after: Duration,
1733}
1734impl CachedDataset {
1735    fn new(dataset: Dataset, refresh_after: Duration) -> Self {
1736        Self {
1737            dataset,
1738            last_refresh: Instant::now(),
1739            refresh_after,
1740        }
1741    }
1742    async fn latest(&mut self) -> Result<Dataset> {
1743        if self.last_refresh.elapsed() >= self.refresh_after {
1744            self.dataset.checkout_latest().await?;
1745            self.last_refresh = Instant::now();
1746        }
1747        Ok(self.dataset.clone())
1748    }
1749    fn replace(&mut self, dataset: Dataset) {
1750        self.dataset = dataset;
1751        self.last_refresh = Instant::now();
1752    }
1753}
1754
1755/// Outcome of one [`Handle::append_stream`] write. Lance's `execute_stream`
1756/// returns only the new `Dataset` (no write summary), so these totals are
1757/// captured from the cumulative `WriteStats` ticks plus pond's own OCC attempt
1758/// counter.
1759#[derive(Debug, Clone, Copy, Default)]
1760pub struct AppendStats {
1761    pub rows: u64,
1762    pub bytes_written: u64,
1763    pub files_written: u64,
1764    pub attempts: u32,
1765}
1766
1767/// Monotonic high-water fold over the cumulative `WriteStats` ticks
1768/// `append_stream` receives. Lance restarts a stream's cumulative counters from
1769/// zero on each OCC retry, so `fetch_max` keeps the fold monotonic - a retry
1770/// contributes nothing until it passes the prior mark, making `AppendStats`
1771/// exact under retries.
1772#[derive(Default)]
1773struct WriteAccum {
1774    rows: std::sync::atomic::AtomicU64,
1775    bytes: std::sync::atomic::AtomicU64,
1776    files: std::sync::atomic::AtomicU64,
1777}
1778
1779impl WriteAccum {
1780    fn observe(&self, stats: &WriteStats) {
1781        use std::sync::atomic::Ordering::Relaxed;
1782        self.rows.fetch_max(stats.rows_written, Relaxed);
1783        self.bytes.fetch_max(stats.bytes_written, Relaxed);
1784        self.files.fetch_max(stats.files_written as u64, Relaxed);
1785    }
1786    fn rows(&self) -> u64 {
1787        self.rows.load(std::sync::atomic::Ordering::Relaxed)
1788    }
1789    fn bytes(&self) -> u64 {
1790        self.bytes.load(std::sync::atomic::Ordering::Relaxed)
1791    }
1792    fn files(&self) -> u64 {
1793        self.files.load(std::sync::atomic::Ordering::Relaxed)
1794    }
1795}
1796
1797/// Append-mode write params. Byte-sized fragments, not Lance's 90 GB default:
1798/// kilobyte rows would otherwise pack multi-GiB fragments that compaction
1799/// rewrites wholesale (see `TARGET_FRAGMENT_BYTES`). Reuses the create params so
1800/// appended fragments match the table's storage version / row-id mode.
1801fn append_write_params() -> WriteParams {
1802    let mut params = sessions::write_params_for_create();
1803    params.mode = WriteMode::Append;
1804    params.max_bytes_per_file = TARGET_FRAGMENT_BYTES as usize;
1805    params
1806}
1807
1808impl Handle {
1809    /// Open without storage options or explicit cache caps. Backend-aware
1810    /// defaults from `[runtime]` apply.
1811    pub async fn open(location: &Url) -> Result<Self> {
1812        Self::open_with_options(location, HashMap::new(), RuntimeCaps::default()).await
1813    }
1814
1815    /// Live size in bytes of the shared Lance session caches (index + metadata).
1816    /// Walks the caches, so it is not cheap - bench/diagnostic use only.
1817    pub fn lance_cache_bytes(&self) -> u64 {
1818        self.session.size_bytes()
1819    }
1820
1821    /// Open with object-store options handed through to Lance verbatim, plus
1822    /// the resolved `[runtime]` cache caps. Object-store keys are the
1823    /// `object_store` crate's standard config names; pond does not parse them.
1824    /// Opening datasets never performs index work; index lifecycle lives under
1825    /// `Handle::optimize_table`. `sessions.lance` and `parts.lance` open lazily
1826    /// on first use.
1827    pub async fn open_with_options(
1828        location: &Url,
1829        storage_options: HashMap<String, String>,
1830        caps: RuntimeCaps,
1831    ) -> Result<Self> {
1832        Self::open_with_options_cached(location, storage_options, caps, None).await
1833    }
1834
1835    /// Like [`Self::open_with_options`], plus an `_indices/*` disk cache rooted
1836    /// at `index_cache_dir` (caller supplies it, mirroring `ensure_rowmap`) so a
1837    /// fresh process skips the cold index load. Ignored for local-FS stores.
1838    pub async fn open_with_options_cached(
1839        location: &Url,
1840        mut storage_options: HashMap<String, String>,
1841        caps: RuntimeCaps,
1842        index_cache_dir: Option<PathBuf>,
1843    ) -> Result<Self> {
1844        if let Some(path) = config::local_path(location) {
1845            tokio::fs::create_dir_all(&path).await.with_context(|| {
1846                format!(
1847                    "failed to create data dir {}; fix the storage destination ([storage].path in config) or re-run `pond init`",
1848                    path.display()
1849                )
1850            })?;
1851        } else {
1852            apply_remote_storage_defaults(&mut storage_options);
1853        }
1854        // One Session shared across all three datasets so metadata/index
1855        // caches and the object_store registry (and thus any S3 client) are
1856        // pooled rather than duplicated three times. Caps are sized by the
1857        // `[runtime]` block; explicit values from `caps` win, otherwise the
1858        // local/remote backend default kicks in.
1859        let (index_cache_bytes, metadata_cache_bytes) = resolve_cache_caps(location, caps);
1860        let session = Arc::new(Session::new(
1861            index_cache_bytes,
1862            metadata_cache_bytes,
1863            Arc::new(ObjectStoreRegistry::default()),
1864        ));
1865        // Build the lance-namespace catalog seam once (spec.md#lance-chokepoints-catalog).
1866        // The `root` property is whatever URL the Directory impl understands;
1867        // `uri_to_url` (lance-io/object_store.rs) accepts both bare paths and
1868        // URLs, so passing the scheme-qualified URL for local FS works the
1869        // same as the bare-path form. Trailing slash stripped for clean logs.
1870        let root = location.as_str().trim_end_matches('/').to_string();
1871        let mut connect = ConnectBuilder::new("dir")
1872            .property("root", root)
1873            .session(session.clone());
1874        // Object-store credentials/region/endpoint flow into the namespace
1875        // via the `storage.<key>` property convention (lance-namespace-impls
1876        // dir.rs from_properties: lines 423-436).
1877        for (key, value) in &storage_options {
1878            connect = connect.property(format!("storage.{key}"), value.clone());
1879        }
1880        let nm: Arc<dyn LanceNamespace> = connect
1881            .connect()
1882            .await
1883            .context("failed to connect lance Directory namespace")?;
1884        let nm_ident = NamespaceIdent::root();
1885        // spec.md#lance-handle-freshness: refresh window is scheme-keyed. Local-FS
1886        // manifest reads are microsecond-cheap, so `0` (always-refresh) is
1887        // essentially free and removes the stale-read window entirely. Object
1888        // stores have real per-call cost; `5s` caps manifest fetch overhead at
1889        // acceptable lag for human-driven queries.
1890        let refresh_after = if config::is_local(location) {
1891            Duration::ZERO
1892        } else {
1893            Duration::from_secs(5)
1894        };
1895        let index_wrapper = index_store_wrapper(location, index_cache_dir.as_deref());
1896        let handle = Self {
1897            datasets: DatasetSet {
1898                sessions: OnceCell::new(),
1899                messages: Mutex::new(CachedDataset::new(
1900                    open_or_create_via_ns(
1901                        &nm,
1902                        &nm_ident,
1903                        sessions::MESSAGES,
1904                        sessions::message_schema(),
1905                        &session,
1906                        &storage_options,
1907                        index_wrapper.clone(),
1908                    )
1909                    .await?,
1910                    refresh_after,
1911                )),
1912                parts: OnceCell::new(),
1913            },
1914            retry: RetryPolicy::default(),
1915            session,
1916            nm,
1917            nm_ident,
1918            storage_options,
1919            location: location.clone(),
1920            lazy_refresh_after: refresh_after,
1921            index_wrapper,
1922        };
1923        Ok(handle)
1924    }
1925
1926    pub fn location(&self) -> &Url {
1927        &self.location
1928    }
1929
1930    /// Read-only view of the `storage_options` the handle was opened with.
1931    /// `pond status` needs them to instantiate a raw `object_store` client
1932    /// that can `LIST` the remote bucket for sizing.
1933    pub fn storage_options(&self) -> &HashMap<String, String> {
1934        &self.storage_options
1935    }
1936
1937    /// Object-store URI for a `pond_sql_query` export artifact:
1938    /// `<location>/exports/<name>`. A sibling of the `*.lance` table dirs;
1939    /// the Directory namespace tracks tables in its `__manifest` table rather
1940    /// than by listing prefixes, so this prefix is never seen as a table
1941    /// (lance-namespace-impls dir/manifest.rs). Never `register_table`'d.
1942    fn export_uri(&self, name: &str) -> String {
1943        format!(
1944            "{}/exports/{name}",
1945            self.location.as_str().trim_end_matches('/')
1946        )
1947    }
1948
1949    /// `ObjectStoreParams` carrying the handle's `storage_options` so raw
1950    /// object-store opens (export I/O, `table_sizes` listing) inherit the same
1951    /// credentials/region as the dataset opens. Empty options -> no accessor.
1952    fn object_store_params(&self) -> ObjectStoreParams {
1953        ObjectStoreParams {
1954            storage_options_accessor: (!self.storage_options.is_empty()).then(|| {
1955                Arc::new(StorageOptionsAccessor::with_static_options(
1956                    self.storage_options.clone(),
1957                ))
1958            }),
1959            ..Default::default()
1960        }
1961    }
1962
1963    /// Write a `pond_sql_query` export artifact, reusing the handle's
1964    /// storage_options so S3 installs inherit the same credentials.
1965    pub(crate) async fn export_write(&self, name: &str, bytes: &[u8]) -> Result<()> {
1966        let uri = self.export_uri(name);
1967        let registry = Arc::new(ObjectStoreRegistry::default());
1968        let (store, path) =
1969            ObjectStore::from_uri_and_params(registry, &uri, &self.object_store_params())
1970                .await
1971                .with_context(|| format!("failed to open object store for {uri}"))?;
1972        store
1973            .put(&path, bytes)
1974            .await
1975            .with_context(|| format!("failed to write export {uri}"))?;
1976        Ok(())
1977    }
1978
1979    /// Read a `pond_sql_query` export artifact back (for the
1980    /// `pond-sql-export://` MCP resource).
1981    pub(crate) async fn export_read(&self, name: &str) -> Result<Vec<u8>> {
1982        let uri = self.export_uri(name);
1983        let registry = Arc::new(ObjectStoreRegistry::default());
1984        let (store, path) =
1985            ObjectStore::from_uri_and_params(registry, &uri, &self.object_store_params())
1986                .await
1987                .with_context(|| format!("failed to open object store for {uri}"))?;
1988        let bytes = store
1989            .read_one_all(&path)
1990            .await
1991            .with_context(|| format!("failed to read export {uri}"))?;
1992        Ok(bytes.to_vec())
1993    }
1994
1995    /// Local filesystem path of an export artifact, when the data dir is
1996    /// `file://`. The stdio MCP client shares this filesystem, so it can read
1997    /// the file directly (e.g. duckdb/polars) instead of pulling base64 via
1998    /// `resources/read`. `None` on object-store installs.
1999    pub(crate) fn export_local_path(&self, name: &str) -> Option<std::path::PathBuf> {
2000        if self.location.scheme() != "file" {
2001            return None;
2002        }
2003        let dir = self.location.to_file_path().ok()?;
2004        Some(dir.join("exports").join(name))
2005    }
2006
2007    pub async fn row_counts(&self) -> Result<(usize, usize, usize)> {
2008        Ok((
2009            self.count_rows(Table::Sessions).await?,
2010            self.count_rows(Table::Messages).await?,
2011            self.count_rows(Table::Parts).await?,
2012        ))
2013    }
2014
2015    /// Insert-only merge: append new rows, never overwrite a matched PK.
2016    /// Returns rows inserted. The fold lives separately under
2017    /// `Handle::optimize_table` (spec.md#lance-index-maintenance).
2018    pub(crate) async fn merge_insert(
2019        &self,
2020        table: Table,
2021        batch: RecordBatch,
2022        row_count: usize,
2023    ) -> Result<u64> {
2024        self.merge_insert_stats(table, batch, row_count)
2025            .await
2026            .map(|stats| stats.num_inserted_rows + stats.num_updated_rows)
2027    }
2028
2029    /// Insert-only merge that surfaces Lance's full `MergeStats`. Callers that
2030    /// need bytes written, file count, or OCC retry count (e.g. `pond copy`'s
2031    /// progress display) use this; the thin wrapper above keeps the
2032    /// affected-rows return for everyone else.
2033    pub(crate) async fn merge_insert_stats(
2034        &self,
2035        table: Table,
2036        batch: RecordBatch,
2037        row_count: usize,
2038    ) -> Result<MergeStats> {
2039        self.merge(
2040            table,
2041            batch,
2042            row_count,
2043            "merge_insert",
2044            WhenMatched::DoNothing,
2045            WhenNotMatched::InsertAll,
2046        )
2047        .await
2048    }
2049
2050    /// Update-only merge: `WhenMatched::UpdateAll` on matched PKs; unmatched
2051    /// rows dropped. The fold lives separately under `Handle::optimize_table`.
2052    pub(crate) async fn merge_update(
2053        &self,
2054        table: Table,
2055        batch: RecordBatch,
2056        row_count: usize,
2057    ) -> Result<u64> {
2058        self.merge(
2059            table,
2060            batch,
2061            row_count,
2062            "merge_update",
2063            WhenMatched::UpdateAll,
2064            WhenNotMatched::DoNothing,
2065        )
2066        .await
2067        .map(|stats| stats.num_inserted_rows + stats.num_updated_rows)
2068    }
2069
2070    /// The OCC write-commit seam (spec.md#lance-chokepoints-write): every write -
2071    /// `merge` and the append paths - runs through here. It takes the cached
2072    /// handle's lock, hands `execute` the latest dataset, commits the dataset
2073    /// `execute` returns, and keeps the cache coherent - all under retry.
2074    /// `execute` builds the table-specific builder, runs it, and returns the new
2075    /// dataset plus its own stats payload; it reruns per OCC attempt, so it owns
2076    /// what it needs. Write-type specifics (params, stats, tracing) stay with the
2077    /// caller.
2078    async fn write_committed<E, Fut, P>(&self, table: Table, execute: E) -> Result<P>
2079    where
2080        E: Fn(Arc<Dataset>) -> Fut,
2081        Fut: std::future::Future<Output = Result<(Dataset, P)>>,
2082    {
2083        self.write_committed_with(table, |_| true, execute).await
2084    }
2085
2086    /// [`Self::write_committed`] with a retry gate (see
2087    /// [`Self::retry_lance_filtered`]). `merge_insert` is idempotent on retry
2088    /// (`WhenMatched::DoNothing` re-reads and no-ops), so it retries everything;
2089    /// the bare `Append` path passes [`is_commit_conflict`] so a post-commit
2090    /// transient fault surfaces rather than re-appending into a duplicate.
2091    async fn write_committed_with<E, Fut, P, R>(
2092        &self,
2093        table: Table,
2094        should_retry: R,
2095        execute: E,
2096    ) -> Result<P>
2097    where
2098        E: Fn(Arc<Dataset>) -> Fut,
2099        Fut: std::future::Future<Output = Result<(Dataset, P)>>,
2100        R: Fn(&anyhow::Error) -> bool,
2101    {
2102        self.retry_lance_filtered(table.label(), should_retry, || {
2103            let execute = &execute;
2104            async move {
2105                let mut cached = self.cached(table).await?.lock().await;
2106                let existing = cached.latest().await?;
2107                let (dataset, payload) = execute(Arc::new(existing)).await?;
2108                cached.replace(dataset);
2109                Ok(payload)
2110            }
2111        })
2112        .await
2113    }
2114
2115    /// Shared merge path for [`Self::merge_insert`] and [`Self::merge_update`].
2116    /// Returns Lance's `MergeStats` verbatim so the progress layer can read
2117    /// `bytes_written` / `num_files_written` / `num_attempts` without a second
2118    /// round-trip; the thin wrappers above project to `u64` for callers that
2119    /// only need the affected-rows count.
2120    async fn merge(
2121        &self,
2122        table: Table,
2123        batch: RecordBatch,
2124        row_count: usize,
2125        op: &'static str,
2126        when_matched: WhenMatched,
2127        when_not_matched: WhenNotMatched,
2128    ) -> Result<MergeStats> {
2129        if row_count == 0 {
2130            return Ok(MergeStats::default());
2131        }
2132        let started = Instant::now();
2133        let result = self
2134            .write_committed(table, |existing| {
2135                let batch = batch.clone();
2136                let when_matched = when_matched.clone();
2137                let when_not_matched = when_not_matched.clone();
2138                async move {
2139                    let schema = batch.schema();
2140                    let reader = RecordBatchIterator::new([Ok(batch)], schema);
2141                    let mut builder = MergeInsertBuilder::try_new(existing, Vec::new())?;
2142                    builder.when_matched(when_matched);
2143                    builder.when_not_matched(when_not_matched);
2144                    // pond presents each PK at most once per batch; FirstSeen keeps
2145                    // the first occurrence rather than failing (Lance's default).
2146                    builder.source_dedupe_behavior(SourceDedupeBehavior::FirstSeen);
2147                    // Cleanup is operator-driven via `pond optimize`; the per-commit
2148                    // auto hook would add a LIST per write on remote backends without
2149                    // changing the steady-state retention.
2150                    builder.skip_auto_cleanup(true);
2151                    let (dataset, stats) = builder
2152                        .try_build()?
2153                        .execute_reader(Box::new(reader))
2154                        .await?;
2155                    Ok((dataset.as_ref().clone(), stats))
2156                }
2157            })
2158            .await;
2159        let skipped = result
2160            .as_ref()
2161            .map(|s| s.num_skipped_duplicates)
2162            .unwrap_or(0);
2163        tracing::info!(
2164            target: "pond::perf",
2165            op,
2166            table = %table.label(),
2167            rows = row_count,
2168            elapsed_ms = started.elapsed().as_millis() as u64,
2169            skipped,
2170            "merge",
2171        );
2172        result
2173    }
2174
2175    /// Append a streamed source into `table` under a single commit - the
2176    /// bandwidth-bound counterpart to [`Self::merge`]. spec.md#session-durable-copy:
2177    /// rows that cannot collide on the destination (absent sessions) take this
2178    /// path. `Append` never joins or probes the target, so its cost is the
2179    /// bytes written, not the per-batch commit + key-scan that `merge_insert`
2180    /// pays - the fix for store-to-store copy being commit-latency-bound on
2181    /// remote object stores.
2182    ///
2183    /// `make_source` is a *factory*, not a prebuilt stream: a Lance scan stream
2184    /// is one-shot, so an OCC retry rebuilds it. A single per-call `WriteAccum`
2185    /// (shared across attempts, NOT fresh per attempt) makes the row/byte/file
2186    /// fold exact under retries.
2187    ///
2188    /// Unlike [`Self::append_batches`] this keeps the retry-everything
2189    /// `write_committed`: a transient fault during the large streamed upload
2190    /// almost always precedes the manifest commit (the rebuilt source re-uploads
2191    /// and the orphaned fragments are GC'd, no duplicate), so failing a full
2192    /// bulk copy on every transient to close the narrow lost-ack-after-commit
2193    /// window is the wrong trade. That rare window is surfaced by the copy
2194    /// verify's duplicate check instead (spec.md#session-movement-complete).
2195    pub(crate) async fn append_stream<F, Fut>(
2196        &self,
2197        table: Table,
2198        make_source: F,
2199    ) -> Result<AppendStats>
2200    where
2201        F: Fn() -> Fut,
2202        Fut: std::future::Future<Output = Result<SendableRecordBatchStream>>,
2203    {
2204        let cum = Arc::new(WriteAccum::default());
2205        let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
2206        let started = Instant::now();
2207        self.write_committed(table, |existing| {
2208            let make_source = &make_source;
2209            let cum = cum.clone();
2210            let attempts = attempts.clone();
2211            async move {
2212                attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2213                let stream = make_source().await?;
2214                let dataset = InsertBuilder::new(existing)
2215                    .with_params(&append_write_params())
2216                    .progress(move |stats| cum.observe(&stats))
2217                    .execute_stream(stream)
2218                    .await?;
2219                Ok((dataset, ()))
2220            }
2221        })
2222        .await?;
2223
2224        let attempts = attempts.load(std::sync::atomic::Ordering::Relaxed);
2225        let stats = AppendStats {
2226            rows: cum.rows(),
2227            bytes_written: cum.bytes(),
2228            files_written: cum.files(),
2229            attempts,
2230        };
2231        tracing::info!(
2232            target: "pond::perf",
2233            op = "append",
2234            table = %table.label(),
2235            rows = stats.rows,
2236            files = stats.files_written,
2237            attempts,
2238            elapsed_ms = started.elapsed().as_millis() as u64,
2239            "append",
2240        );
2241        Ok(stats)
2242    }
2243
2244    /// [`Self::append_stream`] for batches pond already holds in memory (the sync
2245    /// write path) instead of a source-store scan. Row count is taken from the
2246    /// batches - exact under OCC retry without depending on the progress tick.
2247    ///
2248    /// Retries only on a commit *conflict*, not on transient faults: `Append`
2249    /// has no row-level idempotency, so re-running it after a manifest commit
2250    /// that landed but whose ack was lost would duplicate the rows. A conflict
2251    /// proves the commit did not land (re-append is safe); anything else
2252    /// surfaces and the caller's re-plan-from-current-state re-run heals it
2253    /// without doubling rows (spec.md#lance-deterministic-pk).
2254    pub(crate) async fn append_batches(
2255        &self,
2256        table: Table,
2257        batches: Vec<RecordBatch>,
2258    ) -> Result<AppendStats> {
2259        let total_rows: u64 = batches.iter().map(|batch| batch.num_rows() as u64).sum();
2260        if total_rows == 0 {
2261            return Ok(AppendStats::default());
2262        }
2263        let cum = Arc::new(WriteAccum::default());
2264        let attempts = Arc::new(std::sync::atomic::AtomicU32::new(0));
2265        let started = Instant::now();
2266        self.write_committed_with(table, is_commit_conflict, |existing| {
2267            let cum = cum.clone();
2268            let attempts = attempts.clone();
2269            let batches = batches.clone();
2270            async move {
2271                attempts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2272                let dataset = InsertBuilder::new(existing)
2273                    .with_params(&append_write_params())
2274                    .progress(move |stats| cum.observe(&stats))
2275                    .execute(batches)
2276                    .await?;
2277                Ok((dataset, ()))
2278            }
2279        })
2280        .await?;
2281
2282        let attempts = attempts.load(std::sync::atomic::Ordering::Relaxed);
2283        let stats = AppendStats {
2284            rows: total_rows,
2285            bytes_written: cum.bytes(),
2286            files_written: cum.files(),
2287            attempts,
2288        };
2289        tracing::info!(
2290            target: "pond::perf",
2291            op = "append_batches",
2292            table = %table.label(),
2293            rows = stats.rows,
2294            files = stats.files_written,
2295            attempts,
2296            elapsed_ms = started.elapsed().as_millis() as u64,
2297            "append",
2298        );
2299        Ok(stats)
2300    }
2301
2302    /// Run the table-local maintenance cycle for the supplied index intents.
2303    /// Every index family folds incrementally via `optimize_indices`; none is
2304    /// rebuilt from scratch (spec.md#lance-index-maintenance).
2305    ///
2306    /// spec.md#substrate 3.7 (`lance-index-maintenance`): indices and compaction
2307    /// commit independently and use independent retry budgets, so a hot writer
2308    /// that starves compaction (Rewrite) does not abort the index build
2309    /// (Update) the operator actually asked for.
2310    pub async fn optimize_table(
2311        &self,
2312        table: Table,
2313        intents: &[IndexIntent],
2314        progress: Option<&OptimizeProgressFn>,
2315        policy: &MaintenancePolicy,
2316    ) -> TableOptimizeOutcome {
2317        let compaction = self
2318            .run_optimize_compact_phase(table, progress, policy)
2319            .await;
2320        let indices = self
2321            .run_optimize_indices_phase(table, intents, progress, policy.fold_thresholds())
2322            .await;
2323        TableOptimizeOutcome {
2324            table,
2325            indices,
2326            compaction,
2327        }
2328    }
2329
2330    /// Run only the indices phase for one table. Used by the optimize embed
2331    /// stage's tail
2332    /// to fold newly written vectors into the indices without paying the
2333    /// compaction retry budget while embed itself may still be writing.
2334    pub async fn optimize_table_indices_only(
2335        &self,
2336        table: Table,
2337        intents: &[IndexIntent],
2338        progress: Option<&OptimizeProgressFn>,
2339    ) -> PhaseOutcome {
2340        // Thresholds 0: this tail-fold path always folds every index; only
2341        // `pond sync` batches them (`with_scalar_fold_row_threshold` /
2342        // `with_index_fold_row_threshold`).
2343        self.run_optimize_indices_phase(
2344            table,
2345            intents,
2346            progress,
2347            FoldThresholds {
2348                scalar: 0,
2349                index: 0,
2350            },
2351        )
2352        .await
2353    }
2354
2355    async fn run_optimize_indices_phase(
2356        &self,
2357        table: Table,
2358        intents: &[IndexIntent],
2359        progress: Option<&OptimizeProgressFn>,
2360        folds: FoldThresholds,
2361    ) -> PhaseOutcome {
2362        if intents.is_empty() {
2363            return PhaseOutcome::Noop;
2364        }
2365        let result = self
2366            .retry_lance(table.label(), || async {
2367                let mut guard = self.cached(table).await?.lock().await;
2368                let mut dataset = guard.latest().await?;
2369                let did_work =
2370                    optimize_table_indices(&mut dataset, intents, table, progress, folds).await?;
2371                guard.replace(dataset);
2372                Ok::<_, anyhow::Error>(did_work)
2373            })
2374            .await;
2375        match result {
2376            Ok(true) => PhaseOutcome::Ok,
2377            Ok(false) => PhaseOutcome::Noop,
2378            Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
2379            Err(error) => PhaseOutcome::Failed(error),
2380        }
2381    }
2382
2383    async fn run_optimize_compact_phase(
2384        &self,
2385        table: Table,
2386        progress: Option<&OptimizeProgressFn>,
2387        policy: &MaintenancePolicy,
2388    ) -> PhaseOutcome {
2389        let result = self
2390            .retry_lance(table.label(), || async {
2391                let mut guard = self.cached(table).await?.lock().await;
2392                let mut dataset = guard.latest().await?;
2393                optimize_table_compact(&mut dataset, table, progress, policy).await?;
2394                guard.replace(dataset);
2395                Ok::<_, anyhow::Error>(())
2396            })
2397            .await;
2398        match result {
2399            Ok(()) => PhaseOutcome::Ok,
2400            Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
2401            Err(error) => PhaseOutcome::Failed(error),
2402        }
2403    }
2404
2405    pub async fn rebuild_index(
2406        &self,
2407        table: Table,
2408        intent: &IndexIntent,
2409        progress: Option<&OptimizeProgressFn>,
2410    ) -> Result<()> {
2411        emit(
2412            progress,
2413            OptimizeEvent::PhaseStart {
2414                table,
2415                phase: OptimizePhase::IndexRebuild,
2416                detail: Some(intent.name.to_owned()),
2417            },
2418        );
2419        let started = Instant::now();
2420        let result = self
2421            .retry_lance(table.label(), || async {
2422                let mut guard = self.cached(table).await?.lock().await;
2423                let mut dataset = guard.latest().await?;
2424                rebuild_index(&mut dataset, intent, progress, table).await?;
2425                guard.replace(dataset);
2426                Ok(())
2427            })
2428            .await;
2429        emit(
2430            progress,
2431            OptimizeEvent::PhaseDone {
2432                table,
2433                phase: OptimizePhase::IndexRebuild,
2434                elapsed_ms: started.elapsed().as_millis() as u64,
2435            },
2436        );
2437        result
2438    }
2439
2440    /// Lance `cleanup_old_versions` for one table: reclaim files no manifest
2441    /// within the retention window references. No compaction and no new commit -
2442    /// it only deletes superseded files, so no OCC retry is needed.
2443    pub async fn cleanup_table_versions(
2444        &self,
2445        table: Table,
2446        older_than: chrono::Duration,
2447    ) -> Result<()> {
2448        let mut guard = self.cached(table).await?.lock().await;
2449        let dataset = guard.latest().await?;
2450        dataset
2451            .cleanup_old_versions(older_than, Some(false), Some(false))
2452            .await
2453            .with_context(|| format!("cleanup_old_versions failed for {}", table.label()))?;
2454        Ok(())
2455    }
2456
2457    pub async fn index_status(
2458        &self,
2459        table: Table,
2460        intents: &[IndexIntent],
2461        indexable_only: bool,
2462    ) -> Result<Vec<IndexStatus>> {
2463        let dataset = self.dataset(table).await?;
2464        index_status(table, &dataset, intents, indexable_only).await
2465    }
2466
2467    pub(crate) async fn dataset(&self, table: Table) -> Result<Dataset> {
2468        let mut cached = self.cached(table).await?.lock().await;
2469        cached.latest().await
2470    }
2471    /// Build a prefiltered `Scanner` for `table`. Composable read entry
2472    /// point for callers that need to layer extra builder calls
2473    /// (`full_text_search`, `nearest`) on top of pond's predicate seam.
2474    /// Routine scans should prefer `Handle::scan`.
2475    pub(crate) async fn scanner(
2476        &self,
2477        table: Table,
2478        predicate: Option<&Predicate>,
2479    ) -> Result<lance::dataset::scanner::Scanner> {
2480        let dataset = self.dataset(table).await?;
2481        scanner_with_prefilter(&dataset, predicate)
2482    }
2483    /// Single read entry point: prefilter via `predicate`, optionally
2484    /// project, return the prepared `Scanner` (spec.md#lance-chokepoints-read).
2485    pub async fn scan(
2486        &self,
2487        table: Table,
2488        opts: ScanOpts<'_>,
2489    ) -> Result<lance::dataset::scanner::Scanner> {
2490        let mut scanner = self.scanner(table, opts.predicate).await?;
2491        if let Some(projection) = opts.projection {
2492            scanner.project(projection)?;
2493        }
2494        Ok(scanner)
2495    }
2496    pub(crate) async fn scan_batch(
2497        &self,
2498        table: Table,
2499        predicate: Option<&Predicate>,
2500        projection: &[&str],
2501    ) -> Result<RecordBatch> {
2502        let opts = ScanOpts {
2503            predicate,
2504            projection: (!projection.is_empty()).then_some(projection),
2505        };
2506        self.scan(table, opts)
2507            .await?
2508            .try_into_batch()
2509            .await
2510            .context("scan failed")
2511    }
2512    pub async fn count_rows(&self, table: Table) -> Result<usize> {
2513        self.dataset(table)
2514            .await?
2515            .count_rows(None)
2516            .await
2517            .map_err(Into::into)
2518    }
2519    /// Collect the primary-key (`id`) set for `table`. Storage verification
2520    /// compares these sets across two stores: matching row counts can still
2521    /// hide divergent membership, so proving a destination is a complete
2522    /// superset of a source needs the ids, not the cardinalities
2523    /// (spec.md#substrate, `lance-deterministic-pk`).
2524    pub async fn collect_ids(&self, table: Table) -> Result<std::collections::HashSet<String>> {
2525        let batch = self.scan_batch(table, None, &["id"]).await?;
2526        let ids = batch
2527            .column_by_name("id")
2528            .context("scan projection dropped the id column")?
2529            .as_any()
2530            .downcast_ref::<StringArray>()
2531            .context("id column is not Utf8")?;
2532        Ok(ids.iter().flatten().map(str::to_owned).collect())
2533    }
2534    /// Names of every index on `messages` - the vector-index tests read this.
2535    #[cfg(test)]
2536    pub(crate) async fn messages_index_names(&self) -> Result<Vec<String>> {
2537        let dataset = self.dataset(Table::Messages).await?;
2538        let indices = dataset.load_indices().await?;
2539        Ok(indices.iter().map(|index| index.name.clone()).collect())
2540    }
2541
2542    /// Whether `messages` carries an index named `name`. Manifest-only and
2543    /// cache-backed (`load_indices` hits the dataset index cache), so it is
2544    /// cheap enough to gate `Scanner::fast_search` per query: fast-search
2545    /// returns an empty plan when the index is absent, so the retrievers must
2546    /// only opt in once it exists.
2547    pub(crate) async fn messages_has_index(&self, name: &str) -> Result<bool> {
2548        let dataset = self.dataset(Table::Messages).await?;
2549        let indices = dataset.load_indices().await?;
2550        Ok(indices.iter().any(|index| index.name == name))
2551    }
2552
2553    /// Whether `messages` index `name` covers every row - it exists and has no
2554    /// unindexed tail - so `fast_search` (index-only) is complete. When a
2555    /// deferred fold has left a tail, the retrievers must omit `fast_search` so
2556    /// Lance index-probes the folded rows and flat-scans the tail, keeping recall
2557    /// complete (spec.md#search, fts.md "Index Maintenance"). Manifest-only and
2558    /// cache-backed, so it is cheap enough to gate `fast_search` per query.
2559    pub(crate) async fn messages_fast_search_ready(&self, name: &str) -> Result<bool> {
2560        let dataset = self.dataset(Table::Messages).await?;
2561        if !dataset
2562            .load_indices()
2563            .await?
2564            .iter()
2565            .any(|index| index.name == name)
2566        {
2567            return Ok(false);
2568        }
2569        let unindexed = dataset
2570            .unindexed_fragments(name)
2571            .await
2572            .with_context(|| format!("unindexed_fragments failed for {name}"))?;
2573        Ok(unindexed.is_empty())
2574    }
2575
2576    /// Reclaim cached `_indices/<uuid>` dirs no longer referenced by any table's
2577    /// manifest. No-op for local stores or a never-populated cache. Best-effort:
2578    /// a new index version naturally re-fetches, so an over-eager prune only
2579    /// costs one re-download.
2580    pub(crate) async fn prune_index_cache(&self, cache_dir: &std::path::Path) {
2581        if config::is_local(&self.location) {
2582            return;
2583        }
2584        let root = cache_dir.join(store_key(&self.location)).join("indices");
2585        if !root.exists() {
2586            return;
2587        }
2588        let mut keep = std::collections::HashSet::new();
2589        for table in [Table::Sessions, Table::Messages, Table::Parts] {
2590            let Ok(dataset) = self.dataset(table).await else {
2591                return;
2592            };
2593            let Ok(indices) = dataset.load_indices().await else {
2594                return;
2595            };
2596            keep.extend(indices.iter().map(|index| index.uuid.to_string()));
2597        }
2598        prune_stale_uuid_dirs(&root, &keep);
2599    }
2600
2601    /// Count rows in `table` not yet covered by `index_name`. Manifest-only;
2602    /// a missing index reports the whole table. Powers `pond status`.
2603    pub(crate) async fn unindexed_row_count(
2604        &self,
2605        table: Table,
2606        index_name: &str,
2607    ) -> Result<usize> {
2608        let dataset = self.dataset(table).await?;
2609        let fragments = dataset
2610            .unindexed_fragments(index_name)
2611            .await
2612            .with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
2613        Ok(fragments
2614            .iter()
2615            .map(|fragment| fragment.num_rows().unwrap_or(0))
2616            .sum())
2617    }
2618
2619    /// Which table owns the named index, if any. Used by
2620    /// `pond optimize --drop-index <name>` to route the drop to the right
2621    /// dataset without sequentially probing-and-swallowing errors (the prior
2622    /// loop hid permission/network failures behind "no such index"). Runs the
2623    /// three `load_indices` calls in parallel; an error here is a real I/O
2624    /// failure and propagates with context.
2625    pub(crate) async fn find_index_owner(&self, name: &str) -> Result<Option<Table>> {
2626        let list = |table: Table| async move {
2627            let dataset = self.dataset(table).await?;
2628            let names: Vec<String> = dataset
2629                .load_indices()
2630                .await
2631                .with_context(|| format!("load_indices failed for {}", table.label()))?
2632                .iter()
2633                .map(|index| index.name.clone())
2634                .collect();
2635            Ok::<_, anyhow::Error>(names)
2636        };
2637        let (sessions, messages, parts) = tokio::try_join!(
2638            list(Table::Sessions),
2639            list(Table::Messages),
2640            list(Table::Parts),
2641        )?;
2642        for (table, names) in [
2643            (Table::Sessions, sessions),
2644            (Table::Messages, messages),
2645            (Table::Parts, parts),
2646        ] {
2647            if names.iter().any(|n| n == name) {
2648                return Ok(Some(table));
2649            }
2650        }
2651        Ok(None)
2652    }
2653
2654    /// Drop the named index. Used by the `pond optimize --force-embed` model-swap path
2655    /// to retire an IVF_SQ whose centroids belong to the old distance
2656    /// space, before the next write re-bootstraps it over the new model's
2657    /// vectors. Errors when the index does not exist; callers may swallow
2658    /// that.
2659    pub(crate) async fn drop_index(&self, table: Table, name: &str) -> Result<()> {
2660        let mut guard = self.cached(table).await?.lock().await;
2661        let mut dataset = guard.latest().await?;
2662        dataset
2663            .drop_index(name)
2664            .await
2665            .with_context(|| format!("drop_index({name}) failed for {}", table.label()))?;
2666        guard.replace(dataset);
2667        Ok(())
2668    }
2669
2670    /// Resolve each table's stored location through the namespace catalog
2671    /// (spec.md#lance-chokepoints-catalog) - no hardcoded `.lance` suffix.
2672    async fn table_location(&self, table_name: &str) -> Result<String> {
2673        let request = DescribeTableRequest {
2674            id: Some(self.nm_ident.as_table_id(table_name)),
2675            ..Default::default()
2676        };
2677        let response = self
2678            .nm
2679            .describe_table(request)
2680            .await
2681            .with_context(|| format!("failed to describe table {table_name}"))?;
2682        response
2683            .location
2684            .with_context(|| format!("namespace returned no location for table {table_name}"))
2685    }
2686
2687    /// Whether the store holds synced data yet. `open` eagerly creates only the
2688    /// `messages` dataset; `sessions` and `parts` open lazily on first use
2689    /// (see `open_with_options`), so `parts`' presence is the "has been synced"
2690    /// signal - letting read-only surfaces (`pond status`) render an empty state
2691    /// instead of erroring on the first `parts` describe.
2692    pub async fn initialized(&self) -> Result<bool> {
2693        let request = DescribeTableRequest {
2694            id: Some(self.nm_ident.as_table_id(sessions::PARTS)),
2695            ..Default::default()
2696        };
2697        match self.nm.describe_table(request).await {
2698            Ok(_) => Ok(true),
2699            Err(error) if is_namespace_error_code(&error, ErrorCode::TableNotFound) => Ok(false),
2700            Err(error) => {
2701                Err(anyhow::Error::from(error)).context("failed to probe table existence")
2702            }
2703        }
2704    }
2705
2706    /// On-disk byte totals for the three datasets plus the data-dir remainder.
2707    /// Every byte is sized by listing through Lance's object store
2708    /// (spec.md#lance-chokepoints-storage), identical for `file://` and `s3://`.
2709    pub async fn table_sizes(&self) -> Result<TableSizes> {
2710        let registry = Arc::new(ObjectStoreRegistry::default());
2711        let params = self.object_store_params();
2712
2713        let sessions = self
2714            .listed_size(
2715                &registry,
2716                &params,
2717                &self.table_location(sessions::SESSIONS).await?,
2718            )
2719            .await?;
2720        let messages = self
2721            .listed_size(
2722                &registry,
2723                &params,
2724                &self.table_location(sessions::MESSAGES).await?,
2725            )
2726            .await?;
2727        let parts = self
2728            .listed_size(
2729                &registry,
2730                &params,
2731                &self.table_location(sessions::PARTS).await?,
2732            )
2733            .await?;
2734        // `other` is whatever sits under the data-dir root but not in the three
2735        // tables (config.toml, stray index temp files): root total minus them.
2736        let root_total = self
2737            .listed_size(&registry, &params, self.location.as_str())
2738            .await?;
2739        let other = root_total.saturating_sub(sessions + messages + parts);
2740        let sessions_data = self
2741            .data_liveness(&registry, &params, Table::Sessions, sessions::SESSIONS)
2742            .await?;
2743        let messages_data = self
2744            .data_liveness(&registry, &params, Table::Messages, sessions::MESSAGES)
2745            .await?;
2746        let parts_data = self
2747            .data_liveness(&registry, &params, Table::Parts, sessions::PARTS)
2748            .await?;
2749        Ok(TableSizes {
2750            sessions,
2751            messages,
2752            parts,
2753            other,
2754            sessions_data,
2755            messages_data,
2756            parts_data,
2757        })
2758    }
2759
2760    async fn data_liveness(
2761        &self,
2762        registry: &Arc<ObjectStoreRegistry>,
2763        params: &ObjectStoreParams,
2764        table: Table,
2765        table_name: &str,
2766    ) -> Result<DataLiveness> {
2767        let location = self.table_location(table_name).await?;
2768        let data_dir = format!("{}/data", location.trim_end_matches('/'));
2769        let on_disk = self.listed_size(registry, params, &data_dir).await?;
2770        let dataset = self.dataset(table).await?;
2771        let live = dataset
2772            .get_fragments()
2773            .iter()
2774            .try_fold(0u64, |total, fragment| {
2775                Some(total + fragment_bytes(fragment.metadata())?)
2776            });
2777        Ok(DataLiveness { on_disk, live })
2778    }
2779
2780    /// Sum `ObjectMeta.size` for every object recursively under `uri`.
2781    async fn listed_size(
2782        &self,
2783        registry: &Arc<ObjectStoreRegistry>,
2784        params: &ObjectStoreParams,
2785        uri: &str,
2786    ) -> Result<u64> {
2787        let (store, base) = ObjectStore::from_uri_and_params(registry.clone(), uri, params)
2788            .await
2789            .with_context(|| format!("failed to open object store for {uri}"))?;
2790        let mut listing = store.list(Some(base));
2791        let mut total = 0u64;
2792        while let Some(meta) = listing.next().await {
2793            let meta = meta.with_context(|| format!("listing {uri} failed"))?;
2794            total += meta.size;
2795        }
2796        Ok(total)
2797    }
2798    async fn cached(&self, table: Table) -> Result<&Mutex<CachedDataset>> {
2799        match table {
2800            Table::Sessions => self.sessions_cached().await,
2801            Table::Messages => Ok(&self.datasets.messages),
2802            Table::Parts => self.parts_cached().await,
2803        }
2804    }
2805
2806    /// Open `sessions.lance` on first use (spec.md#datasets). The search request
2807    /// path reads only `messages`; the daemon's background index-cache GC
2808    /// (`prune_index_cache`) opens this on a `serve`/`mcp` process. Single-flight
2809    /// via `OnceCell`, like `parts`.
2810    async fn sessions_cached(&self) -> Result<&Mutex<CachedDataset>> {
2811        self.lazy_cached(
2812            &self.datasets.sessions,
2813            sessions::SESSIONS,
2814            sessions::session_schema,
2815        )
2816        .await
2817    }
2818
2819    /// Open `parts.lance` on first use (spec.md#datasets). Single-flight via
2820    /// `OnceCell`; once initialized, behaves identically to the other two.
2821    async fn parts_cached(&self) -> Result<&Mutex<CachedDataset>> {
2822        self.lazy_cached(&self.datasets.parts, sessions::PARTS, sessions::part_schema)
2823            .await
2824    }
2825
2826    /// Shared lazy-open path for the `sessions`/`parts` `OnceCell`s. `schema` is a
2827    /// thunk so the (local-CPU) schema build happens only on the cold init, not
2828    /// on every cache hit.
2829    async fn lazy_cached<'a>(
2830        &self,
2831        cell: &'a OnceCell<Mutex<CachedDataset>>,
2832        table_name: &str,
2833        schema: fn() -> lance::deps::arrow_schema::SchemaRef,
2834    ) -> Result<&'a Mutex<CachedDataset>> {
2835        cell.get_or_try_init(|| async {
2836            let dataset = open_or_create_via_ns(
2837                &self.nm,
2838                &self.nm_ident,
2839                table_name,
2840                schema(),
2841                &self.session,
2842                &self.storage_options,
2843                self.index_wrapper.clone(),
2844            )
2845            .await?;
2846            Ok::<_, anyhow::Error>(Mutex::new(CachedDataset::new(
2847                dataset,
2848                self.lazy_refresh_after,
2849            )))
2850        })
2851        .await
2852    }
2853    async fn retry_lance<T, Fut, Op>(&self, label: &str, operation: Op) -> Result<T>
2854    where
2855        Fut: std::future::Future<Output = Result<T>>,
2856        Op: FnMut() -> Fut,
2857    {
2858        // Default: retry every transient fault (spec.md#lance-retry-jitter).
2859        self.retry_lance_filtered(label, |_| true, operation).await
2860    }
2861
2862    /// Like [`Self::retry_lance`] but `should_retry` gates which errors are
2863    /// retried. [`Self::append_batches`] passes [`is_commit_conflict`]: a commit
2864    /// conflict means this writer's commit did NOT land, so re-running the
2865    /// operation is safe; any other error (notably a transient fault that may
2866    /// have arrived *after* the manifest commit landed - the lost-ack case) is
2867    /// surfaced instead of retried, because `Append` has no row-level
2868    /// idempotency and a blind re-append would duplicate. The caller's
2869    /// operation re-plans from current state on its own re-run, which is the
2870    /// idempotent recovery (spec.md#lance-deterministic-pk).
2871    async fn retry_lance_filtered<T, Fut, Op, R>(
2872        &self,
2873        label: &str,
2874        should_retry: R,
2875        mut operation: Op,
2876    ) -> Result<T>
2877    where
2878        Fut: std::future::Future<Output = Result<T>>,
2879        Op: FnMut() -> Fut,
2880        R: Fn(&anyhow::Error) -> bool,
2881    {
2882        let mut attempt = 0u8;
2883        loop {
2884            attempt = attempt.saturating_add(1);
2885            match operation().await {
2886                Ok(value) => return Ok(value),
2887                Err(error) if attempt < self.retry.attempts && should_retry(&error) => {
2888                    let backoff = self.backoff(attempt);
2889                    // `{:#}` walks anyhow's cause chain inline; `%error` (Display)
2890                    // drops everything below the top-level message.
2891                    let error_chain = format!("{error:#}");
2892                    tracing::warn!(
2893                        label,
2894                        attempt,
2895                        ?backoff,
2896                        error = %error_chain,
2897                        "retrying Lance operation"
2898                    );
2899                    tokio::time::sleep(backoff).await;
2900                }
2901                Err(error) => {
2902                    let error_chain = format!("{error:#}");
2903                    tracing::warn!(
2904                        label,
2905                        attempt,
2906                        error = %error_chain,
2907                        "Lance operation exhausted retries"
2908                    );
2909                    // spec.md#protocol: surface OCC failures as a typed `conflict`
2910                    // rather than the generic `storage_unavailable` bucket. The
2911                    // chain root is a `lance::Error` (commit-conflict family) when
2912                    // pond's retry layer exhausted because the manifest could not
2913                    // be advanced; everything else (timeouts, IAM, disk) stays
2914                    // `storage_unavailable`.
2915                    if is_commit_conflict(&error) {
2916                        return Err(error.context(ConflictExhausted { attempts: attempt }));
2917                    }
2918                    return Err(error);
2919                }
2920            }
2921        }
2922    }
2923    fn backoff(&self, attempt: u8) -> Duration {
2924        let shift = u32::from(attempt.saturating_sub(1));
2925        let multiplier = 1u32.checked_shl(shift).unwrap_or(u32::MAX);
2926        let base = self.retry.initial_backoff.saturating_mul(multiplier);
2927        // Symmetric +/- `jitter` factor de-correlates concurrent retriers on
2928        // a contended manifest (spec.md#lance-retry-jitter); clamped to `max_backoff`.
2929        let factor = (1.0 + self.retry.jitter * (fastrand::f64() * 2.0 - 1.0)).max(0.0);
2930        base.mul_f64(factor).min(self.retry.max_backoff)
2931    }
2932}
2933/// Compaction phase: plan + amplification veto + execute + `cleanup_old_versions`,
2934/// one retry block, separate from the indices phase so a lost Rewrite race
2935/// does not abort index work.
2936///
2937/// Vetoes Lance-planned tasks instead of pre-gating on pond fragment math:
2938/// Lance bins split at index-coverage boundaries, so pond predictions diverge
2939/// from what Lance actually rewrites (the old run-sum gate latched open and
2940/// rewrote a 665 MiB tail fragment every 5-min sync). Only whole planned
2941/// tasks are filtered, so OCC and conflict semantics are untouched.
2942///
2943/// spec.md#lance-index-maintenance mandates FRI on by default, but at
2944/// v7.0.0-beta.16 `defer_index_remap=true` together with `stable-row-ids`
2945/// panics in `optimize.rs::commit_compaction` with "defer_index_remap
2946/// requires row_addrs but none were provided": `rewrite_files` skips
2947/// row_addrs when stable row ids are on, then the FRI builder demands
2948/// them. With stable_row_ids the remap step is already a no-op
2949/// (`optimize.rs:1490`: `needs_remapping = !uses_stable_row_ids() &&
2950/// !defer_index_remap`), so running without FRI is correct - we only
2951/// lose the documented concurrency-with-index-build benefit. Flip to
2952/// `true` once upstream fixes the conflict.
2953async fn optimize_table_compact(
2954    dataset: &mut Dataset,
2955    table: Table,
2956    progress: Option<&OptimizeProgressFn>,
2957    policy: &MaintenancePolicy,
2958) -> Result<()> {
2959    let stats: Vec<FragmentStat> = dataset
2960        .get_fragments()
2961        .iter()
2962        .map(|fragment| fragment_stat(fragment.metadata()))
2963        .collect();
2964    let compaction = CompactionOptions {
2965        target_rows_per_fragment: derived_target_rows(&stats),
2966        max_bytes_per_file: Some(TARGET_FRAGMENT_BYTES as usize),
2967        defer_index_remap: false,
2968        // Binary-copy eligible fragments (concatenate encoded pages, no
2969        // decode/re-encode) and fall back to Reencode automatically for blob
2970        // (parts), deletion-bearing, or schema-varied fragments. ~27% faster on
2971        // the messages/sessions reencode path, safe everywhere else.
2972        compaction_mode: Some(CompactionMode::TryBinaryCopy),
2973        ..CompactionOptions::default()
2974    };
2975
2976    let mut plan = plan_compaction(dataset, &compaction).await?;
2977    if policy.compaction_fragment_cap > 0 {
2978        plan.tasks.retain(|task| {
2979            let task_stats: Vec<FragmentStat> = task.fragments.iter().map(fragment_stat).collect();
2980            let keep = keep_task(
2981                &task_stats,
2982                policy.compaction_fragment_cap,
2983                compaction.materialize_deletions_threshold,
2984            );
2985            if !keep {
2986                tracing::debug!(
2987                    target: "pond::perf",
2988                    table = table.as_str(),
2989                    fragments = task_stats.len(),
2990                    "compaction task vetoed: merge dominated by one large fragment",
2991                );
2992            }
2993            keep
2994        });
2995    }
2996    if plan.tasks.is_empty() {
2997        tracing::debug!(
2998            target: "pond::perf",
2999            table = table.as_str(),
3000            "compaction skipped: no task to run",
3001        );
3002    } else {
3003        emit(
3004            progress,
3005            OptimizeEvent::PhaseStart {
3006                table,
3007                phase: OptimizePhase::Compact,
3008                detail: None,
3009            },
3010        );
3011        let started = Instant::now();
3012        let mut completed = Vec::with_capacity(plan.tasks.len());
3013        for task in plan.compaction_tasks() {
3014            completed.push(task.execute(dataset).await?);
3015        }
3016        commit_compaction(
3017            dataset,
3018            completed,
3019            Arc::new(DatasetIndexRemapperOptions::default()),
3020            &compaction,
3021        )
3022        .await?;
3023        emit(
3024            progress,
3025            OptimizeEvent::PhaseDone {
3026                table,
3027                phase: OptimizePhase::Compact,
3028                elapsed_ms: started.elapsed().as_millis() as u64,
3029            },
3030        );
3031    }
3032
3033    // Safe GC only. delete_unverified=false keeps Lance's 7-day in-progress
3034    // guard, so this never races a concurrent writer (spec.md#concurrency); GC
3035    // runs outside OCC, so the guard is what makes it safe on any backend.
3036    //
3037    // Gated: the walk over the version log is round-trip-bound on object stores
3038    // (~9s measured on the real corpus) and reclaims ~one version per run, so
3039    // the frequent `pond sync` path amortizes it over `cleanup_interval`
3040    // commits rather than paying it every sync (`pond optimize`/`pond copy`
3041    // keep interval 1). Skipping only delays reclaiming old versions - the next
3042    // due cleanup sweeps the accumulated backlog - so it is always safe.
3043    if cleanup_due(dataset.version_id(), policy.cleanup_interval) {
3044        emit(
3045            progress,
3046            OptimizeEvent::PhaseStart {
3047                table,
3048                phase: OptimizePhase::Cleanup,
3049                detail: None,
3050            },
3051        );
3052        let started = Instant::now();
3053        // Lance v7 `cleanup_old_versions` removes orphan files inside
3054        // `_indices/<uuid>/` but does NOT remove the parent dir, so failed/no-op
3055        // index merges accumulate empty UUID dirs forever (one inode each).
3056        // Harmless beyond inode pressure; tracked upstream. No pond-side FS sweep
3057        // here (spec.md#concurrency: Lance-native maintenance only).
3058        dataset
3059            .cleanup_old_versions(policy.cleanup_older_than, Some(false), Some(false))
3060            .await
3061            .context("cleanup_old_versions failed during index optimize")?;
3062        emit(
3063            progress,
3064            OptimizeEvent::PhaseDone {
3065                table,
3066                phase: OptimizePhase::Cleanup,
3067                elapsed_ms: started.elapsed().as_millis() as u64,
3068            },
3069        );
3070    }
3071
3072    Ok(())
3073}
3074
3075/// Gate for the version-cleanup walk: at interval `<= 1` it runs every optimize;
3076/// otherwise only when the manifest `version` is a multiple of it. A run whose
3077/// version steps past a multiple defers to the next one, so the gap between
3078/// cleanups is bounded and - since version 0 is a multiple of every interval -
3079/// cleanup always eventually fires; it is never skipped indefinitely.
3080fn cleanup_due(version: u64, interval: u64) -> bool {
3081    interval <= 1 || version.is_multiple_of(interval)
3082}
3083
3084/// Indices phase: create absent indexes, then fold trailing fragments into
3085/// every existing index via batched `optimize_indices` (append, or merge once a
3086/// family's delta segments reach `DELTA_MERGE_THRESHOLD`). Returns `true` if
3087/// anything committed.
3088async fn optimize_table_indices(
3089    dataset: &mut Dataset,
3090    intents: &[IndexIntent],
3091    table: Table,
3092    progress: Option<&OptimizeProgressFn>,
3093    folds: FoldThresholds,
3094) -> Result<bool> {
3095    let existing = dataset.load_indices().await?;
3096    let existing_names: std::collections::HashSet<String> =
3097        existing.iter().map(|index| index.name.clone()).collect();
3098
3099    let mut append_indices: Vec<String> = Vec::new();
3100    let mut did_work = false;
3101
3102    for intent in intents {
3103        let exists = existing_names.contains(intent.name);
3104
3105        if !exists {
3106            if !intent.trigger.should_create(dataset).await? {
3107                continue;
3108            }
3109            let params = intent.params.build(dataset).await?;
3110            let index_type = intent.params.index_type();
3111            tracing::info!(
3112                index = intent.name,
3113                column = intent.column,
3114                "creating Lance index (trigger fired)",
3115            );
3116            emit(
3117                progress,
3118                OptimizeEvent::PhaseStart {
3119                    table,
3120                    phase: OptimizePhase::IndexCreate,
3121                    detail: Some(intent.name.to_owned()),
3122                },
3123            );
3124            let started = Instant::now();
3125            dataset
3126                .create_index_builder(&[intent.column], index_type, params.as_ref())
3127                .name(intent.name.to_owned())
3128                .replace(false)
3129                .progress(lance_progress(progress, table, intent.name))
3130                .await
3131                .with_context(|| format!("failed to create index {}", intent.name))?;
3132            emit(
3133                progress,
3134                OptimizeEvent::PhaseDone {
3135                    table,
3136                    phase: OptimizePhase::IndexCreate,
3137                    elapsed_ms: started.elapsed().as_millis() as u64,
3138                },
3139            );
3140            did_work = true;
3141            continue;
3142        }
3143
3144        // A trailing tail (rows written since the last fold) stays fully
3145        // searchable while deferred: the retrievers drop `fast_search` whenever
3146        // an index has an unindexed tail, so Lance index-probes the folded rows
3147        // and flat-scans the tail (spec.md#search, fts.md "Index Maintenance").
3148        // `DELTA_MERGE_THRESHOLD` keeps the per-fold segment count bounded.
3149        let unindexed = dataset.unindexed_fragments(intent.name).await?;
3150        if unindexed.is_empty() {
3151            continue;
3152        }
3153        let tail_rows: usize = unindexed
3154            .iter()
3155            .map(|fragment| fragment.num_rows().unwrap_or(0))
3156            .sum();
3157        // Batch folds per family so a tiny sync doesn't pay a per-fold cost that
3158        // dwarfs the delta. Scalar (BTree/bitmap): Lance 7.0.0 rewrites the whole
3159        // index file per fold (O(index size)); defer until the tail is worth one
3160        // rewrite - get/count/sql read the deferred tail via scan. FTS + vector
3161        // (IVF): the fold is a cheap delta append, but each is an S3 round-trip +
3162        // commit storm; defer until the tail is worth one fold - search
3163        // flat-scans the deferred tail so recall stays complete either way.
3164        // `pond optimize`/`pond copy` pass 0 (fold every run).
3165        let fold_threshold = match intent.params {
3166            IndexParamsKind::Scalar(_) => folds.scalar,
3167            IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. } => folds.index,
3168        };
3169        if fold_threshold > 0 && tail_rows < fold_threshold {
3170            tracing::debug!(
3171                target: "pond::perf",
3172                index = intent.name,
3173                tail_rows,
3174                threshold = fold_threshold,
3175                "deferring index fold (unindexed tail below threshold)",
3176            );
3177            continue;
3178        }
3179        // FTS-only guard: folding a tail with zero non-null values writes an
3180        // empty delta segment, which Lance 7.0.0 reads back with the wrong
3181        // posting-tail codec (`Default` = VarintDelta vs the metadata-absent
3182        // default Fixed32), deterministically failing every later merge. The
3183        // probe is bounded: it reads only the tail fragments the fold itself
3184        // would read, stops at the first non-null value, and runs only after
3185        // the fold threshold already passed.
3186        if matches!(intent.params, IndexParamsKind::InvertedFtsWord)
3187            && !column_has_values(dataset, intent.column, &unindexed).await?
3188        {
3189            tracing::debug!(
3190                target: "pond::perf",
3191                index = intent.name,
3192                tail_rows,
3193                "skipping FTS fold (tail has no indexable values)",
3194            );
3195            continue;
3196        }
3197        // Every family folds incrementally via `optimize_indices` (the
3198        // append/merge batch below) - no full rebuild. BTree rewrites its index
3199        // file by merging the existing sorted pages with only the new fragments'
3200        // data; Bitmap/FTS/IVF_SQ accumulate delta segments. None re-scans
3201        // already-indexed source (spec.md#lance-index-maintenance).
3202        append_indices.push(intent.name.to_owned());
3203    }
3204
3205    if !append_indices.is_empty() {
3206        // Per-index segment count from the manifest loaded above (delta segments
3207        // share the intent name). Indices that have piled up
3208        // `DELTA_MERGE_THRESHOLD` segments fold with `merge` (collapse to one);
3209        // the rest take the cheap append. Splitting keeps each query reading few
3210        // segments without paying a consolidation on every tiny fold.
3211        let segment_count = |name: &str| {
3212            existing
3213                .iter()
3214                .filter(|index| index.name.as_str() == name)
3215                .count()
3216        };
3217        let (consolidate, to_append): (Vec<String>, Vec<String>) = append_indices
3218            .iter()
3219            .cloned()
3220            .partition(|name| segment_count(name) >= DELTA_MERGE_THRESHOLD);
3221        // FTS delta segments are never merged: Lance 7.0.0's inverted merge
3222        // has two reproducible defects on real segments - "different posting
3223        // tail codecs" (a partitionless segment reports the derive-default
3224        // codec) and an index-out-of-bounds panic in InnerBuilder::merge_from
3225        // (token ids past the resized posting table; crashed the 5-min cron
3226        // sync in a loop). At the same threshold a merge would fire, the FTS
3227        // index instead rebuilds from scratch - the one consolidation path
3228        // Lance executes correctly. Scalar and vector merges are unaffected.
3229        let mut fts_rebuilds: Vec<&IndexIntent> = Vec::new();
3230        let mut to_merge: Vec<String> = Vec::new();
3231        for name in consolidate {
3232            let fts_intent = intents.iter().find(|intent| {
3233                intent.name == name && matches!(intent.params, IndexParamsKind::InvertedFtsWord)
3234            });
3235            match fts_intent {
3236                Some(intent) => fts_rebuilds.push(intent),
3237                None => to_merge.push(name),
3238            }
3239        }
3240
3241        emit(
3242            progress,
3243            OptimizeEvent::PhaseStart {
3244                table,
3245                phase: OptimizePhase::IndexAppend,
3246                detail: Some(append_indices.join(", ")),
3247            },
3248        );
3249        let started = Instant::now();
3250        if !to_append.is_empty() {
3251            dataset
3252                .optimize_indices(&OptimizeOptions::append().index_names(to_append))
3253                .await
3254                .context("optimize_indices(append) failed during index optimize")?;
3255        }
3256        if !to_merge.is_empty() {
3257            dataset
3258                .optimize_indices(
3259                    &OptimizeOptions::merge(DELTA_MERGE_THRESHOLD).index_names(to_merge),
3260                )
3261                .await
3262                .context("optimize_indices(merge) failed during index optimize")?;
3263        }
3264        emit(
3265            progress,
3266            OptimizeEvent::PhaseDone {
3267                table,
3268                phase: OptimizePhase::IndexAppend,
3269                elapsed_ms: started.elapsed().as_millis() as u64,
3270            },
3271        );
3272        for intent in &fts_rebuilds {
3273            emit(
3274                progress,
3275                OptimizeEvent::PhaseStart {
3276                    table,
3277                    phase: OptimizePhase::IndexRebuild,
3278                    detail: Some(intent.name.to_owned()),
3279                },
3280            );
3281            let rebuild_started = Instant::now();
3282            rebuild_index(dataset, intent, progress, table).await?;
3283            emit(
3284                progress,
3285                OptimizeEvent::PhaseDone {
3286                    table,
3287                    phase: OptimizePhase::IndexRebuild,
3288                    elapsed_ms: rebuild_started.elapsed().as_millis() as u64,
3289                },
3290            );
3291        }
3292        tracing::debug!(
3293            target: "pond::perf",
3294            indices = ?append_indices,
3295            rebuilt = ?fts_rebuilds,
3296            "folded trailing fragments into indices",
3297        );
3298        did_work = true;
3299    }
3300
3301    Ok(did_work)
3302}
3303
3304/// Fragment-scoped scanner over the non-null rows of `column` - the shared
3305/// base of the existence probe and the indexable count below.
3306fn non_null_scanner(
3307    dataset: &Dataset,
3308    column: &'static str,
3309    fragments: &[lance::table::format::Fragment],
3310) -> Result<lance::dataset::scanner::Scanner> {
3311    let mut scanner = dataset.scan();
3312    scanner.with_fragments(fragments.to_vec());
3313    scanner.filter(&Predicate::IsNotNull(column).to_lance())?;
3314    Ok(scanner)
3315}
3316
3317/// True when any row in `fragments` holds a non-null value for `column`.
3318/// Existence probe for the FTS fold guard: scans only the given fragments,
3319/// stops at the first hit (`limit 1`), so its read is a subset of what the
3320/// fold it gates would read.
3321async fn column_has_values(
3322    dataset: &Dataset,
3323    column: &'static str,
3324    fragments: &[lance::table::format::Fragment],
3325) -> Result<bool> {
3326    let mut scanner = non_null_scanner(dataset, column, fragments)?;
3327    scanner.project(&[column])?;
3328    scanner.limit(Some(1), None)?;
3329    let batch = scanner
3330        .try_into_batch()
3331        .await
3332        .with_context(|| format!("non-null probe on {column} failed"))?;
3333    Ok(batch.num_rows() > 0)
3334}
3335
3336/// Count of rows in `fragments` holding a non-null value for `column`. Serves
3337/// the indexable status view; fragment-scoped, so bounded by the tail.
3338async fn column_value_count(
3339    dataset: &Dataset,
3340    column: &'static str,
3341    fragments: &[lance::table::format::Fragment],
3342) -> Result<usize> {
3343    let count = non_null_scanner(dataset, column, fragments)?
3344        .count_rows()
3345        .await
3346        .with_context(|| format!("non-null count on {column} failed"))?;
3347    Ok(count as usize)
3348}
3349
3350async fn rebuild_index(
3351    dataset: &mut Dataset,
3352    intent: &IndexIntent,
3353    progress: Option<&OptimizeProgressFn>,
3354    table: Table,
3355) -> Result<()> {
3356    if !intent.trigger.should_create(dataset).await? {
3357        return Ok(());
3358    }
3359    let params = intent.params.build(dataset).await?;
3360    dataset
3361        .create_index_builder(
3362            &[intent.column],
3363            intent.params.index_type(),
3364            params.as_ref(),
3365        )
3366        .name(intent.name.to_owned())
3367        .replace(true)
3368        .progress(lance_progress(progress, table, intent.name))
3369        .await
3370        .with_context(|| format!("failed to rebuild index {}", intent.name))?;
3371    Ok(())
3372}
3373
3374async fn index_status(
3375    table: Table,
3376    dataset: &Dataset,
3377    intents: &[IndexIntent],
3378    indexable_only: bool,
3379) -> Result<Vec<IndexStatus>> {
3380    let existing = dataset.load_indices().await?;
3381    let existing_names: std::collections::HashSet<String> =
3382        existing.iter().map(|index| index.name.clone()).collect();
3383    let total_fragments = dataset.get_fragments().len();
3384    let total_rows = dataset.count_rows(None).await?;
3385    let mut statuses = Vec::with_capacity(intents.len());
3386    for intent in intents {
3387        let exists = existing_names.contains(intent.name);
3388        if !exists {
3389            statuses.push(IndexStatus {
3390                table,
3391                intent_name: intent.name.to_owned(),
3392                fragments_covered: 0,
3393                unindexed_fragments: total_fragments,
3394                unindexed_rows: total_rows,
3395                exists,
3396            });
3397            continue;
3398        }
3399        let unindexed = dataset
3400            .unindexed_fragments(intent.name)
3401            .await
3402            .with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
3403        let unindexed_fragments = unindexed.len();
3404        let mut unindexed_rows: usize = unindexed
3405            .iter()
3406            .map(|fragment| fragment.num_rows().unwrap_or(0))
3407            .sum();
3408        // Content indexes (FTS, IVF) take in only non-null rows - most message
3409        // rows carry a null `search_text`/`vector` (tool/system roles), so the
3410        // raw fragment row count vastly overstates the actionable backlog and
3411        // an all-null tail (which the FTS fold guard skips) would read as
3412        // stuck. The indexable view counts what a fold could actually index;
3413        // opt-in because the count scans the tail, which the per-sync summary
3414        // must not pay.
3415        if indexable_only
3416            && unindexed_rows > 0
3417            && matches!(
3418                intent.params,
3419                IndexParamsKind::InvertedFtsWord | IndexParamsKind::IvfSqCosine { .. }
3420            )
3421        {
3422            unindexed_rows = column_value_count(dataset, intent.column, &unindexed).await?;
3423        }
3424        statuses.push(IndexStatus {
3425            table,
3426            intent_name: intent.name.to_owned(),
3427            fragments_covered: total_fragments.saturating_sub(unindexed_fragments),
3428            unindexed_fragments,
3429            unindexed_rows,
3430            exists,
3431        });
3432    }
3433    Ok(statuses)
3434}
3435
3436/// Open the table at `table_name` via the namespace; create + initialize on
3437/// `TableNotFound`. Schema-checks the on-disk dataset against pond's
3438/// expectation so a stale data dir surfaces early.
3439///
3440/// Probes via `nm.describe_table` directly rather than `DatasetBuilder::from_namespace`:
3441/// the builder re-wraps an already-`Namespace`-wrapped error
3442/// (lance/src/dataset/builder.rs:142), so going through it would force a
3443/// chain-walk to classify `TableNotFound`. The direct probe stays at one
3444/// wrap level and downcasts cleanly. Managed-versioning hookup (REST
3445/// namespace external-manifest commits) is not wired here; v1 ships
3446/// Directory v2 only.
3447/// Diagnostic S3 IO tracing. Inert unless [`io_trace::enable`] is called
3448/// before the store opens; then a shared `IOTracker` is injected as the
3449/// object-store wrapper on every dataset read open, counting exactly how many
3450/// GETs (and bytes, and - under the `io-trace` feature - which paths) each
3451/// query issues against a remote store. Used by `serve_mem_bench --io-trace`
3452/// to attribute the per-query S3 request load. Not a production code path.
3453pub mod io_trace {
3454    use lance_io::utils::tracking_store::{IOTracker, IoStats};
3455    use std::sync::{Arc, OnceLock};
3456
3457    static TRACKER: OnceLock<IOTracker> = OnceLock::new();
3458
3459    /// Arm tracing. Must run before the store opens so the wrapper is applied
3460    /// when the datasets' object store is built.
3461    pub fn enable() {
3462        let _ = TRACKER.set(IOTracker::default());
3463    }
3464
3465    /// The shared tracker as an object-store wrapper, when armed.
3466    pub(super) fn wrapper() -> Option<Arc<IOTracker>> {
3467        TRACKER.get().map(|tracker| Arc::new(tracker.clone()))
3468    }
3469
3470    /// Read and reset the IO accumulated since the last call.
3471    pub fn take() -> Option<IoStats> {
3472        TRACKER.get().map(IOTracker::incremental_stats)
3473    }
3474}
3475
3476/// On-disk cache for `_indices/*` so a fresh process serves the IVF + FTS index
3477/// from local disk instead of re-loading it from the object store on every
3478/// cold-start (spec.md#search). Scoped to `_indices/*` because those files are
3479/// immutable and UUID-addressed, so a hit is always correct and a new index is
3480/// an automatic miss; data (served by the rowmap) and manifests (need freshness)
3481/// pass through. A `WrappingObjectStore`, so it stays inside the object-store
3482/// layer rather than reaching around it.
3483pub mod index_cache {
3484    use object_store::local::LocalFileSystem;
3485    use object_store::path::Path as ObjPath;
3486    use object_store::{
3487        CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
3488        ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result as OsResult,
3489    };
3490    use std::collections::HashMap;
3491    use std::ops::Range;
3492    use std::path::PathBuf;
3493    use std::sync::{Arc, Mutex};
3494
3495    use bytes::Bytes;
3496    use futures::stream::BoxStream;
3497    use lance_io::object_store::WrappingObjectStore;
3498
3499    fn is_index_path(location: &ObjPath) -> bool {
3500        AsRef::<str>::as_ref(location).contains("_indices/")
3501    }
3502
3503    /// Drop conditional headers (etag/if-modified): they reference the remote
3504    /// object and would spuriously fail against the local cache copy.
3505    fn local_opts(options: &GetOptions) -> GetOptions {
3506        GetOptions {
3507            range: options.range.clone(),
3508            head: options.head,
3509            ..Default::default()
3510        }
3511    }
3512
3513    /// `WrappingObjectStore` factory: holds the per-store cache root and hands a
3514    /// `CachingStore` to every dataset open on this store.
3515    #[derive(Debug)]
3516    pub struct IndexDiskCache {
3517        local: Arc<LocalFileSystem>,
3518        inflight: Arc<Mutex<HashMap<ObjPath, Arc<tokio::sync::Mutex<()>>>>>,
3519    }
3520
3521    impl IndexDiskCache {
3522        /// `LocalFileSystem` requires the prefix to exist, so create it first.
3523        pub fn new(root: PathBuf) -> std::io::Result<Self> {
3524            std::fs::create_dir_all(&root)?;
3525            Ok(Self {
3526                local: Arc::new(LocalFileSystem::new_with_prefix(&root)?),
3527                inflight: Arc::new(Mutex::new(HashMap::new())),
3528            })
3529        }
3530    }
3531
3532    impl WrappingObjectStore for IndexDiskCache {
3533        fn wrap(&self, _store_prefix: &str, inner: Arc<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
3534            Arc::new(CachingStore {
3535                inner,
3536                local: self.local.clone(),
3537                inflight: self.inflight.clone(),
3538            })
3539        }
3540    }
3541
3542    #[derive(Debug)]
3543    struct CachingStore {
3544        inner: Arc<dyn ObjectStore>,
3545        local: Arc<LocalFileSystem>,
3546        inflight: Arc<Mutex<HashMap<ObjPath, Arc<tokio::sync::Mutex<()>>>>>,
3547    }
3548
3549    impl std::fmt::Display for CachingStore {
3550        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3551            write!(f, "CachingStore({})", self.inner)
3552        }
3553    }
3554
3555    impl CachingStore {
3556        fn flight_lock(&self, location: &ObjPath) -> Arc<tokio::sync::Mutex<()>> {
3557            self.inflight
3558                .lock()
3559                .unwrap_or_else(|poison| poison.into_inner())
3560                .entry(location.clone())
3561                .or_default()
3562                .clone()
3563        }
3564
3565        /// Fetch the whole object once, write it (`LocalFileSystem::put` stages +
3566        /// renames atomically), then serve the requested range from the copy. The
3567        /// per-path single-flight coalesces a process's concurrent first reads of
3568        /// one file into a single fetch; cross-process writes race safely since
3569        /// the bytes are identical and the rename is atomic.
3570        async fn populate_and_serve(
3571            &self,
3572            location: &ObjPath,
3573            options: GetOptions,
3574        ) -> OsResult<GetResult> {
3575            let lock = self.flight_lock(location);
3576            let _guard = lock.lock().await;
3577            let result = self.fetch_under_flight(location, options).await;
3578            // Drop the entry so the map stays bounded as index versions churn.
3579            // Unconditionally safe (singleflight idiom): any waiter already holds
3580            // its own `lock` clone, and a later miss re-creates the entry but
3581            // finds the file cached.
3582            self.inflight
3583                .lock()
3584                .unwrap_or_else(|p| p.into_inner())
3585                .remove(location);
3586            result
3587        }
3588
3589        async fn fetch_under_flight(
3590            &self,
3591            location: &ObjPath,
3592            options: GetOptions,
3593        ) -> OsResult<GetResult> {
3594            if let Ok(result) = self.local.get_opts(location, local_opts(&options)).await {
3595                return Ok(result);
3596            }
3597            let bytes = self.inner.get(location).await?.bytes().await?;
3598            if self
3599                .local
3600                .put(location, PutPayload::from_bytes(bytes))
3601                .await
3602                .is_ok()
3603                && let Ok(result) = self.local.get_opts(location, local_opts(&options)).await
3604            {
3605                return Ok(result);
3606            }
3607            // Cache write or re-read failed (e.g. disk full): serve from origin.
3608            self.inner.get_opts(location, options).await
3609        }
3610    }
3611
3612    #[async_trait::async_trait]
3613    impl ObjectStore for CachingStore {
3614        async fn get_opts(&self, location: &ObjPath, options: GetOptions) -> OsResult<GetResult> {
3615            if !is_index_path(location) {
3616                return self.inner.get_opts(location, options).await;
3617            }
3618            match self.local.get_opts(location, local_opts(&options)).await {
3619                Ok(result) => Ok(result),
3620                Err(object_store::Error::NotFound { .. }) => {
3621                    self.populate_and_serve(location, options).await
3622                }
3623                Err(_) => self.inner.get_opts(location, options).await,
3624            }
3625        }
3626
3627        async fn put_opts(
3628            &self,
3629            location: &ObjPath,
3630            payload: PutPayload,
3631            opts: PutOptions,
3632        ) -> OsResult<PutResult> {
3633            self.inner.put_opts(location, payload, opts).await
3634        }
3635
3636        async fn put_multipart_opts(
3637            &self,
3638            location: &ObjPath,
3639            opts: PutMultipartOptions,
3640        ) -> OsResult<Box<dyn MultipartUpload>> {
3641            self.inner.put_multipart_opts(location, opts).await
3642        }
3643
3644        async fn get_ranges(
3645            &self,
3646            location: &ObjPath,
3647            ranges: &[Range<u64>],
3648        ) -> OsResult<Vec<Bytes>> {
3649            if is_index_path(location) {
3650                // Through get_opts so the first touch caches the whole object.
3651                let mut out = Vec::with_capacity(ranges.len());
3652                for range in ranges {
3653                    let opts = GetOptions {
3654                        range: Some(range.clone().into()),
3655                        ..Default::default()
3656                    };
3657                    out.push(self.get_opts(location, opts).await?.bytes().await?);
3658                }
3659                return Ok(out);
3660            }
3661            self.inner.get_ranges(location, ranges).await
3662        }
3663
3664        fn delete_stream(
3665            &self,
3666            locations: BoxStream<'static, OsResult<ObjPath>>,
3667        ) -> BoxStream<'static, OsResult<ObjPath>> {
3668            self.inner.delete_stream(locations)
3669        }
3670
3671        fn list(&self, prefix: Option<&ObjPath>) -> BoxStream<'static, OsResult<ObjectMeta>> {
3672            self.inner.list(prefix)
3673        }
3674
3675        fn list_with_offset(
3676            &self,
3677            prefix: Option<&ObjPath>,
3678            offset: &ObjPath,
3679        ) -> BoxStream<'static, OsResult<ObjectMeta>> {
3680            self.inner.list_with_offset(prefix, offset)
3681        }
3682
3683        async fn list_with_delimiter(&self, prefix: Option<&ObjPath>) -> OsResult<ListResult> {
3684            self.inner.list_with_delimiter(prefix).await
3685        }
3686
3687        async fn copy_opts(&self, from: &ObjPath, to: &ObjPath, opts: CopyOptions) -> OsResult<()> {
3688            self.inner.copy_opts(from, to, opts).await
3689        }
3690    }
3691
3692    #[cfg(test)]
3693    mod tests {
3694        #![allow(clippy::unwrap_used)]
3695        use super::*;
3696        use object_store::memory::InMemory;
3697
3698        async fn read(store: &Arc<dyn ObjectStore>, path: &ObjPath) -> Option<Vec<u8>> {
3699            store
3700                .get(path)
3701                .await
3702                .ok()?
3703                .bytes()
3704                .await
3705                .ok()
3706                .map(|b| b.to_vec())
3707        }
3708
3709        #[tokio::test]
3710        async fn caches_index_files_and_passes_data_through() {
3711            let temp = tempfile::tempdir().unwrap();
3712            let inner: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
3713            let index_path = ObjPath::from("d/messages.lance/_indices/uuid1/index.idx");
3714            let data_path = ObjPath::from("d/messages.lance/data/x.lance");
3715            inner
3716                .put(&index_path, PutPayload::from_static(b"INDEX"))
3717                .await
3718                .unwrap();
3719            inner
3720                .put(&data_path, PutPayload::from_static(b"DATA"))
3721                .await
3722                .unwrap();
3723
3724            let cache = IndexDiskCache::new(temp.path().join("indices")).unwrap();
3725            let store = cache.wrap("test", inner.clone());
3726
3727            assert_eq!(
3728                read(&store, &index_path).await.as_deref(),
3729                Some(&b"INDEX"[..])
3730            );
3731            assert_eq!(
3732                read(&store, &data_path).await.as_deref(),
3733                Some(&b"DATA"[..])
3734            );
3735
3736            // Delete both from the origin. The index file is served from the
3737            // local cache; the data file (never cached) is now gone.
3738            inner.delete(&index_path).await.unwrap();
3739            inner.delete(&data_path).await.unwrap();
3740            assert_eq!(
3741                read(&store, &index_path).await.as_deref(),
3742                Some(&b"INDEX"[..])
3743            );
3744            assert_eq!(read(&store, &data_path).await, None);
3745
3746            // A range read of the cached index slices the local copy.
3747            let slice = store.get_range(&index_path, 1..4).await.unwrap();
3748            assert_eq!(slice.as_ref(), b"NDE");
3749        }
3750    }
3751}
3752
3753/// Stable filesystem-safe key for a store URL: same URL -> same key, so sibling
3754/// pond processes share one on-disk cache and distinct stores never collide.
3755/// Shared by the rowmap (`sessions.rs`), the index disk cache, and the CLI's
3756/// per-store sync lock / last-sync state files.
3757pub fn store_key(location: &Url) -> String {
3758    blake3::hash(location.as_str().as_bytes()).to_hex()[..16].to_owned()
3759}
3760
3761/// Reclaim cached `_indices/<uuid>` dirs whose UUID is not in `keep`. Recurses
3762/// to each `_indices` dir (the bucket prefix varies) and prunes its dead UUID
3763/// children. Best-effort; unlink-safe (POSIX keeps an in-flight reader's inode).
3764fn prune_stale_uuid_dirs(dir: &std::path::Path, keep: &std::collections::HashSet<String>) {
3765    let Ok(entries) = std::fs::read_dir(dir) else {
3766        return;
3767    };
3768    for entry in entries.flatten() {
3769        let path = entry.path();
3770        if !path.is_dir() {
3771            continue;
3772        }
3773        if entry.file_name() == "_indices" {
3774            let Ok(children) = std::fs::read_dir(&path) else {
3775                continue;
3776            };
3777            for child in children.flatten() {
3778                if child.path().is_dir()
3779                    && !keep.contains(child.file_name().to_string_lossy().as_ref())
3780                {
3781                    let _ = std::fs::remove_dir_all(child.path());
3782                }
3783            }
3784        } else {
3785            prune_stale_uuid_dirs(&path, keep);
3786        }
3787    }
3788}
3789
3790/// The object-store wrapper applied to every dataset open: the `_indices/*`
3791/// disk cache (remote stores only, when a cache dir is supplied) chained with
3792/// the diagnostic io-trace wrapper. `None` when neither is active.
3793fn index_store_wrapper(
3794    location: &Url,
3795    index_cache_dir: Option<&std::path::Path>,
3796) -> Option<Arc<dyn WrappingObjectStore>> {
3797    let mut wrappers: Vec<Arc<dyn WrappingObjectStore>> = Vec::new();
3798    if let Some(dir) = index_cache_dir
3799        && !config::is_local(location)
3800    {
3801        let root = dir.join(store_key(location)).join("indices");
3802        match index_cache::IndexDiskCache::new(root) {
3803            Ok(cache) => wrappers.push(Arc::new(cache)),
3804            Err(error) => tracing::warn!(%error, "index disk cache disabled; reads hit the store"),
3805        }
3806    }
3807    if let Some(tracker) = io_trace::wrapper() {
3808        wrappers.push(tracker);
3809    }
3810    match wrappers.len() {
3811        0 => None,
3812        1 => Some(wrappers.remove(0)),
3813        _ => Some(Arc::new(ChainedWrappingObjectStore::new(wrappers))),
3814    }
3815}
3816
3817async fn open_or_create_via_ns(
3818    nm: &Arc<dyn LanceNamespace>,
3819    nm_ident: &NamespaceIdent,
3820    table_name: &str,
3821    schema: lance::deps::arrow_schema::SchemaRef,
3822    session: &Arc<Session>,
3823    storage_options: &HashMap<String, String>,
3824    wrapper: Option<Arc<dyn WrappingObjectStore>>,
3825) -> Result<Dataset> {
3826    let table_id = nm_ident.as_table_id(table_name);
3827
3828    let request = DescribeTableRequest {
3829        id: Some(table_id.clone()),
3830        ..Default::default()
3831    };
3832    match nm.describe_table(request).await {
3833        Ok(response) => {
3834            let location = response.location.with_context(|| {
3835                format!("namespace returned no location for table {table_name}")
3836            })?;
3837            let mut builder = DatasetBuilder::from_uri(&location).with_session(session.clone());
3838            match wrapper {
3839                Some(wrapper) => {
3840                    builder = builder.with_store_params(ObjectStoreParams {
3841                        object_store_wrapper: Some(wrapper),
3842                        storage_options_accessor: (!storage_options.is_empty()).then(|| {
3843                            Arc::new(StorageOptionsAccessor::with_static_options(
3844                                storage_options.clone(),
3845                            ))
3846                        }),
3847                        ..Default::default()
3848                    });
3849                }
3850                None => {
3851                    if !storage_options.is_empty() {
3852                        builder = builder.with_storage_options(storage_options.clone());
3853                    }
3854                }
3855            }
3856            let mut dataset = builder
3857                .load()
3858                .await
3859                .with_context(|| format!("failed to open table {table_name}"))?;
3860            ensure_current_schema(&mut dataset, schema.as_ref(), table_name).await?;
3861            return Ok(dataset);
3862        }
3863        Err(error) => match &error {
3864            error if is_namespace_error_code(error, ErrorCode::TableNotFound) => {
3865                // fall through to create
3866            }
3867            _ => {
3868                return Err(anyhow::Error::from(error))
3869                    .with_context(|| format!("failed to describe table {table_name}"));
3870            }
3871        },
3872    }
3873
3874    // Create path: pond seeds an empty dataset with the canonical schema so
3875    // every subsequent open lands on a real Lance dataset, not a phantom.
3876    let mut write_params = sessions::write_params_for_create();
3877    write_params.session = Some(session.clone());
3878    write_params.mode = WriteMode::Create;
3879    if !storage_options.is_empty() {
3880        write_params.store_params = Some(ObjectStoreParams {
3881            storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(
3882                storage_options.clone(),
3883            ))),
3884            ..Default::default()
3885        });
3886    }
3887    let reader = sessions::empty_reader(schema)?;
3888    Dataset::write_into_namespace(reader, nm.clone(), table_id, Some(write_params))
3889        .await
3890        .with_context(|| format!("failed to create table {table_name}"))
3891}
3892
3893// lance-namespace sometimes nests one `lance::Error::Namespace` inside another
3894// before the underlying `NamespaceError`; walk the whole `.source()` chain
3895// rather than only matching the outer variant.
3896fn is_namespace_error_code(error: &lance::Error, code: ErrorCode) -> bool {
3897    if !matches!(error, lance::Error::Namespace { .. }) {
3898        return false;
3899    }
3900    std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |link| {
3901        link.source()
3902    })
3903    .filter_map(|link| link.downcast_ref::<NamespaceError>())
3904    .any(|inner| inner.code() == code)
3905}
3906
3907fn scanner_with_prefilter(
3908    dataset: &Dataset,
3909    predicate: Option<&Predicate>,
3910) -> Result<lance::dataset::scanner::Scanner> {
3911    let mut scanner = dataset.scan();
3912    scanner.prefilter(true);
3913    if let Some(predicate) = predicate {
3914        let filter = predicate.to_lance();
3915        if !filter.is_empty() {
3916            scanner.filter(&filter)?;
3917        }
3918    }
3919    Ok(scanner)
3920}
3921/// How the on-disk schema relates to this build's expected schema.
3922enum SchemaFit {
3923    Match,
3924    /// Expected columns absent on disk, every one nullable: the store
3925    /// predates an additive schema change and upgrades in place.
3926    MissingNullable(Vec<lance::deps::arrow_schema::Field>),
3927    /// On-disk columns this build does not know: written by a newer pond.
3928    /// Reads proceed (scans project known columns; extra ones are inert);
3929    /// writes against the newer schema fail at the Lance layer, and the fix
3930    /// is upgrading pond, not editing the store.
3931    UnknownExtra(Vec<String>),
3932}
3933
3934fn classify_schema(
3935    actual: &lance::deps::arrow_schema::Schema,
3936    expected: &lance::deps::arrow_schema::Schema,
3937    table_name: &str,
3938) -> Result<SchemaFit> {
3939    use std::collections::BTreeSet;
3940    let actual_names: BTreeSet<&str> = actual.fields().iter().map(|f| f.name().as_str()).collect();
3941    let expected_names: BTreeSet<&str> = expected
3942        .fields()
3943        .iter()
3944        .map(|f| f.name().as_str())
3945        .collect();
3946    let missing: Vec<_> = expected
3947        .fields()
3948        .iter()
3949        .filter(|f| !actual_names.contains(f.name().as_str()))
3950        .map(|f| f.as_ref().clone())
3951        .collect();
3952    let extra: Vec<String> = actual_names
3953        .difference(&expected_names)
3954        .map(|name| (*name).to_owned())
3955        .collect();
3956    match (missing.is_empty(), extra.is_empty()) {
3957        (true, true) => Ok(SchemaFit::Match),
3958        (false, true) if missing.iter().all(|f| f.is_nullable()) => {
3959            Ok(SchemaFit::MissingNullable(missing))
3960        }
3961        (true, false) => Ok(SchemaFit::UnknownExtra(extra)),
3962        _ => anyhow::bail!(
3963            "table {table_name} has columns {actual_names:?} but this pond build expects \
3964             {expected_names:?}, and the difference is not an additive nullable-column \
3965             change this build can migrate - upgrade pond, or restore the store from a \
3966             `pond copy` snapshot taken by the version that wrote it",
3967        ),
3968    }
3969}
3970
3971/// Open-time schema reconciliation: a store missing this build's known
3972/// nullable columns is backfilled IN PLACE via `Dataset::add_columns` - the
3973/// values derive from data already stored, so no re-ingest is ever required
3974/// (spec.md#session-durable-copy: a rotated source cannot supply rows again).
3975/// Concurrent openers race benignly: `add_columns` commits through OCC, a
3976/// losing writer sees a conflict, re-checks out latest, and finds the columns
3977/// present.
3978async fn ensure_current_schema(
3979    dataset: &mut Dataset,
3980    expected: &lance::deps::arrow_schema::Schema,
3981    table_name: &str,
3982) -> Result<()> {
3983    use lance::deps::arrow_schema::DataType;
3984    const MAX_MIGRATION_ATTEMPTS: usize = 3;
3985    for _ in 0..MAX_MIGRATION_ATTEMPTS {
3986        let actual = lance::deps::arrow_schema::Schema::from(dataset.schema());
3987        match classify_schema(&actual, expected, table_name)? {
3988            SchemaFit::MissingNullable(missing) => {
3989                backfill_missing_columns(dataset, table_name, missing).await?;
3990                continue;
3991            }
3992            SchemaFit::Match => {}
3993            SchemaFit::UnknownExtra(extra) => {
3994                tracing::warn!(
3995                    table = table_name,
3996                    ?extra,
3997                    "store carries columns unknown to this pond build (written by a newer \
3998                     version); reads proceed, writes need the newer pond",
3999                );
4000            }
4001        }
4002        // Catch a vector-dim change (configured `[embeddings].dim` differs
4003        // from the on-disk vector column width) early with a friendly
4004        // message. Lance would otherwise reject the next write with an
4005        // opaque schema-mismatch error inside the `merge_update` path.
4006        for actual_field in actual.fields() {
4007            let Some(expected_field) = expected.field_with_name(actual_field.name()).ok() else {
4008                continue;
4009            };
4010            if let (
4011                DataType::FixedSizeList(_, actual_dim),
4012                DataType::FixedSizeList(_, expected_dim),
4013            ) = (actual_field.data_type(), expected_field.data_type())
4014                && actual_dim != expected_dim
4015            {
4016                tracing::warn!(
4017                    table = table_name,
4018                    column = actual_field.name(),
4019                    actual_dim,
4020                    expected_dim,
4021                    "embedding dimension differs from config; open proceeds because model swaps are operator-driven",
4022                );
4023            }
4024        }
4025        return Ok(());
4026    }
4027    anyhow::bail!(
4028        "schema migration for table {table_name} did not converge after \
4029         {MAX_MIGRATION_ATTEMPTS} attempts (a concurrent writer kept changing \
4030         the schema); re-run once the other pond process finishes",
4031    )
4032}
4033
4034/// One in-place additive migration pass: derive the missing columns from
4035/// stored data and commit them via `add_columns` (new column files only, no
4036/// row rewrites). The recipe - which columns to read and how to derive the
4037/// values - is consumer knowledge and lives in `sessions::column_backfill`.
4038async fn backfill_missing_columns(
4039    dataset: &mut Dataset,
4040    table_name: &str,
4041    missing: Vec<lance::deps::arrow_schema::Field>,
4042) -> Result<()> {
4043    use lance::dataset::{BatchUDF, NewColumnTransform};
4044    let names: Vec<&str> = missing.iter().map(|f| f.name().as_str()).collect();
4045    let spec = sessions::column_backfill(table_name, &missing)?;
4046    // A full-column read over a remote store runs minutes; a silent stall
4047    // reads as a hang, so this one-time event gets a stderr notice (same
4048    // pattern as the embedding-model download notice in embed.rs).
4049    let _ = crate::output::line_err(&format!(
4050        "migrating {table_name}: backfilling {names:?} from stored data (one-time, in place)...",
4051    ));
4052    let started = std::time::Instant::now();
4053    let mapper = spec.mapper;
4054    // Boxed: `add_columns`' concrete future is enormous (it embeds DataFusion
4055    // planner types), and inlining it here overflows rustc's auto-trait
4056    // solver (E0275) once this future nests inside the open chain.
4057    let migration: std::pin::Pin<
4058        Box<dyn std::future::Future<Output = lance::Result<()>> + Send + '_>,
4059    > = Box::pin(dataset.add_columns(
4060        NewColumnTransform::BatchUDF(BatchUDF {
4061            mapper: Box::new(move |batch| {
4062                mapper(batch).map_err(|error| lance::Error::io(format!("{error:#}")))
4063            }),
4064            output_schema: spec.output_schema,
4065            result_checkpoint: None,
4066        }),
4067        Some(spec.read_columns),
4068        None,
4069    ));
4070    let result = migration.await;
4071    match result {
4072        Ok(()) => {
4073            let _ = crate::output::line_err(&format!(
4074                "migrated {table_name} in {:.1}s",
4075                started.elapsed().as_secs_f64(),
4076            ));
4077            Ok(())
4078        }
4079        Err(error) => {
4080            let error = anyhow::Error::from(error);
4081            if is_commit_conflict(&error) {
4082                // Another writer migrated (or wrote) concurrently; re-check
4083                // out latest and let the caller re-classify.
4084                dataset.checkout_latest().await?;
4085                Ok(())
4086            } else {
4087                Err(error).with_context(|| {
4088                    format!(
4089                        "schema backfill failed for {table_name}; the one-time migration \
4090                         writes new column files, so it needs write access to the store - \
4091                         re-run any pond command with write-capable credentials to complete it",
4092                    )
4093                })
4094            }
4095        }
4096    }
4097}
4098/// Object-store defaults injected for any non-local pond location. Each key
4099/// is only set when neither the user-provided key nor its env-var-form alias
4100/// is already present, so explicit overrides in `[storage]` always win.
4101/// `aws_unsigned_payload` is gated on a custom endpoint (the marker for
4102/// S3-compatible stores like Hetzner, MinIO, R2), where the SHA256 payload
4103/// signature is wasted work the server does not validate.
4104fn apply_remote_storage_defaults(options: &mut HashMap<String, String>) {
4105    fn set_default(options: &mut HashMap<String, String>, aliases: &[&str], value: &str) {
4106        if aliases
4107            .iter()
4108            .any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)))
4109        {
4110            return;
4111        }
4112        options.insert(aliases[0].to_owned(), value.to_owned());
4113    }
4114    set_default(options, &["pool_idle_timeout"], "300 seconds");
4115    set_default(options, &["connect_timeout"], "10 seconds");
4116    // `request_timeout` bounds a single object-store request (one range GET/PUT),
4117    // not a whole scan - a streaming read issues many small requests, each well
4118    // under this. We keep it deliberately tight as a HARD BARRIER: a single
4119    // request exceeding 60s means a design/infra problem to fix (chunk the read,
4120    // use change-data-feed, fix the endpoint), never something to paper over with
4121    // a longer timeout. An explicit `[storage]` override still wins.
4122    set_default(options, &["request_timeout"], "60 seconds");
4123    let has_custom_endpoint = ["aws_endpoint", "endpoint"]
4124        .iter()
4125        .any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)));
4126    if has_custom_endpoint {
4127        set_default(
4128            options,
4129            &["aws_unsigned_payload", "unsigned_payload"],
4130            "true",
4131        );
4132    }
4133}
4134
4135fn quoted_string(value: &str) -> String {
4136    format!("'{}'", value.replace('\'', "''"))
4137}
4138fn like_contains(value: &str) -> String {
4139    let escaped = value
4140        .replace('\\', "\\\\")
4141        .replace('%', "\\%")
4142        .replace('_', "\\_")
4143        .replace('\'', "''");
4144    format!("'%{escaped}%'")
4145}
4146
4147#[cfg(test)]
4148mod tests {
4149    #![allow(clippy::expect_used, clippy::unwrap_used)]
4150
4151    use super::*;
4152    use tempfile::TempDir;
4153
4154    #[test]
4155    fn is_index_error_matches_lance_index_class_through_context() {
4156        let index_fault = anyhow::Error::from(lance::Error::index(
4157            "cannot merge inverted index segments with different posting tail codecs",
4158        ))
4159        .context("optimize_indices(merge) failed during index optimize");
4160        assert!(is_index_error(&index_fault));
4161
4162        let io_fault = anyhow::anyhow!("connection reset").context("optimize_indices failed");
4163        assert!(!is_index_error(&io_fault));
4164    }
4165
4166    #[test]
4167    fn prune_keeps_live_uuid_dirs_and_drops_dead_ones() {
4168        let temp = TempDir::new().unwrap();
4169        let indices = temp.path().join("bkt/messages.lance/_indices");
4170        for uuid in ["live", "dead"] {
4171            std::fs::create_dir_all(indices.join(uuid)).unwrap();
4172            std::fs::write(indices.join(uuid).join("index.idx"), b"x").unwrap();
4173        }
4174        let keep = std::collections::HashSet::from(["live".to_owned()]);
4175        prune_stale_uuid_dirs(temp.path(), &keep);
4176        assert!(indices.join("live").exists());
4177        assert!(!indices.join("dead").exists());
4178    }
4179
4180    fn set(scope: Option<&str>) -> CredsSet {
4181        CredsSet {
4182            scope: scope.map(str::to_owned),
4183            access_key_id: Some("AKIA".to_owned()),
4184            secret_access_key: Some("shh".to_owned()),
4185            ..CredsSet::default()
4186        }
4187    }
4188
4189    fn opts(resolved: &ResolvedStorage, key: &str) -> Option<String> {
4190        resolved.options.get(key).cloned()
4191    }
4192
4193    #[test]
4194    fn storage_url_translation_table() {
4195        // file (Lance's `uri_to_url` appends the trailing slash; `child_uri`
4196        // trims it downstream)
4197        let local = StorageUrl::parse("/srv/pond").unwrap();
4198        assert_eq!(local.lance_url().as_str(), "file:///srv/pond/");
4199        assert!(local.is_local());
4200        assert!(local.scheme_options.is_empty());
4201        // s3 passthrough
4202        let aws = StorageUrl::parse("s3://bucket/prefix").unwrap();
4203        assert_eq!(aws.lance_url().as_str(), "s3://bucket/prefix");
4204        assert!(aws.scheme_options.is_empty());
4205        // s3+https: TLS stays on, virtual-hosted defaults on for domain
4206        // hosts, region defaults deterministically. The endpoint is
4207        // assembled at resolve time with the bucket folded into the host
4208        // (object_store's virtual-hosted convention).
4209        let fat = StorageUrl::parse("s3+https://nbg1.example.com/my-pond/sub").unwrap();
4210        assert_eq!(fat.lance_url().as_str(), "s3://my-pond/sub");
4211        assert_eq!(
4212            fat.scheme_options,
4213            vec![
4214                ("allow_http", "false".to_owned()),
4215                ("virtual_hosted_style_request", "true".to_owned()),
4216                ("region", "us-east-1".to_owned()),
4217            ],
4218        );
4219        let resolved = fat.resolve(&BTreeMap::new()).unwrap();
4220        assert_eq!(
4221            opts(&resolved, "endpoint").as_deref(),
4222            Some("https://my-pond.nbg1.example.com"),
4223        );
4224        assert_eq!(opts(&resolved, "region").as_deref(), Some("us-east-1"));
4225        // s3+http on an IP host: allow_http flips, path-style auto-selected
4226        // (a bucket subdomain on an IP can't resolve), port survives.
4227        let plain = StorageUrl::parse("s3+http://127.0.0.1:9000/pond").unwrap();
4228        assert_eq!(plain.lance_url().as_str(), "s3://pond/");
4229        assert_eq!(plain.scheme_options[0], ("allow_http", "true".to_owned()));
4230        assert_eq!(
4231            plain.scheme_options[1],
4232            ("virtual_hosted_style_request", "false".to_owned()),
4233        );
4234        let resolved = plain.resolve(&BTreeMap::new()).unwrap();
4235        assert_eq!(
4236            opts(&resolved, "endpoint").as_deref(),
4237            Some("http://127.0.0.1:9000"),
4238        );
4239        // An explicit endpoint in `extra` is the escape hatch and wins.
4240        let mut pinned = BTreeMap::new();
4241        pinned.insert(
4242            "default".to_owned(),
4243            CredsSet {
4244                extra: [(
4245                    "endpoint".to_owned(),
4246                    "https://pinned.example.com".to_owned(),
4247                )]
4248                .into_iter()
4249                .collect(),
4250                ..CredsSet::default()
4251            },
4252        );
4253        let resolved = fat.resolve(&pinned).unwrap();
4254        assert_eq!(
4255            opts(&resolved, "endpoint").as_deref(),
4256            Some("https://pinned.example.com"),
4257        );
4258        // gs passthrough
4259        let gcs = StorageUrl::parse("gs://bucket/p").unwrap();
4260        assert_eq!(gcs.lance_url().as_str(), "gs://bucket/p");
4261        // az: account folds into options
4262        let azure = StorageUrl::parse("az://acct/container/p").unwrap();
4263        assert_eq!(azure.lance_url().as_str(), "az://container/p");
4264        assert_eq!(
4265            azure.scheme_options,
4266            vec![("account_name", "acct".to_owned())]
4267        );
4268        // tests-only schemes pass through untouched
4269        let shared = StorageUrl::parse("shared-memory://pond-test-x/").unwrap();
4270        assert_eq!(shared.lance_url().as_str(), "shared-memory://pond-test-x/");
4271    }
4272
4273    #[test]
4274    fn storage_url_rejects_bad_shapes() {
4275        // RFC 3986 userinfo is a leak class, never accepted.
4276        let err = StorageUrl::parse("s3+https://user:pass@host/bucket")
4277            .expect_err("userinfo must be rejected")
4278            .to_string();
4279        assert!(
4280            err.contains("creds"),
4281            "error must name the alternative: {err}"
4282        );
4283        // Missing bucket.
4284        assert!(StorageUrl::parse("s3+https://host").is_err());
4285        assert!(StorageUrl::parse("az://acct").is_err());
4286        // Unknown scheme names the grammar.
4287        let err = StorageUrl::parse("ftp://host/x")
4288            .expect_err("ftp")
4289            .to_string();
4290        assert!(err.contains("s3+https"), "got: {err}");
4291        // Unrecognized query params die loudly.
4292        let err = StorageUrl::parse("s3://b/p?regoin=x")
4293            .expect_err("typo")
4294            .to_string();
4295        assert!(err.contains("regoin"), "got: {err}");
4296        // Query params on local / in-memory schemes die just as loudly -
4297        // no silent carry into the URL Lance opens.
4298        let err = StorageUrl::parse("memory://x?creds=y")
4299            .expect_err("memory query")
4300            .to_string();
4301        assert!(err.contains("query params"), "got: {err}");
4302        let err = StorageUrl::parse("file:///x?creds=y")
4303            .expect_err("file query")
4304            .to_string();
4305        assert!(err.contains("query params"), "got: {err}");
4306        // `?` in a bare path is a filename character, not a query.
4307        assert!(StorageUrl::parse("/tmp/a?b").is_ok());
4308    }
4309
4310    #[test]
4311    fn storage_url_canonicalizes_ports_and_keeps_percent_encoding() {
4312        // Default port strips so scope matching can't split on `:443`.
4313        let with_port = StorageUrl::parse("s3+https://host:443/bucket/p").unwrap();
4314        let without = StorageUrl::parse("s3+https://host/bucket/p").unwrap();
4315        assert_eq!(with_port.canonical(), without.canonical());
4316        // Non-default port survives into the assembled endpoint.
4317        let odd = StorageUrl::parse("s3+https://host:8443/bucket").unwrap();
4318        let resolved = odd.resolve(&BTreeMap::new()).unwrap();
4319        assert_eq!(
4320            resolved.options.get("endpoint").map(String::as_str),
4321            Some("https://bucket.host:8443"),
4322        );
4323        // Percent-encoded prefix passes through to the Lance URL verbatim.
4324        let encoded = StorageUrl::parse("s3+https://host/bucket/pre%20fix").unwrap();
4325        assert_eq!(encoded.lance_url().as_str(), "s3://bucket/pre%20fix");
4326    }
4327
4328    #[test]
4329    fn query_params_strip_and_apply_over_set_fields() {
4330        let mut creds = BTreeMap::new();
4331        creds.insert(
4332            "default".to_owned(),
4333            CredsSet {
4334                region: Some("from-set".to_owned()),
4335                virtual_hosted_style_request: Some(false),
4336                ..set(None)
4337            },
4338        );
4339        let url = StorageUrl::parse(
4340            "s3+https://host/bucket/p?region=from-query&virtual_hosted_style_request=true",
4341        )
4342        .unwrap();
4343        // Stripped before Lance sees the URL.
4344        assert_eq!(url.lance_url().as_str(), "s3://bucket/p");
4345        assert!(url.canonical().query().is_none());
4346        let resolved = url.resolve(&creds).unwrap();
4347        // Assembly precedence: scheme < set < query.
4348        assert_eq!(opts(&resolved, "region").as_deref(), Some("from-query"));
4349        assert_eq!(
4350            opts(&resolved, "virtual_hosted_style_request").as_deref(),
4351            Some("true"),
4352        );
4353        // virtual_hosted=true (query) -> the bucket rides in the endpoint host.
4354        assert_eq!(
4355            opts(&resolved, "endpoint").as_deref(),
4356            Some("https://bucket.host"),
4357        );
4358    }
4359
4360    #[test]
4361    fn scope_matching_binds_by_longest_prefix_at_segment_boundaries() {
4362        let mut creds = BTreeMap::new();
4363        creds.insert("all".to_owned(), set(None));
4364        creds.insert("bucket".to_owned(), set(Some("s3+https://host/pond/")));
4365        creds.insert("deep".to_owned(), set(Some("s3+https://host/pond/sub")));
4366
4367        let bind = |input: &str| {
4368            StorageUrl::parse(input)
4369                .unwrap()
4370                .resolve(&creds)
4371                .unwrap()
4372                .binding
4373        };
4374        // Longest match wins.
4375        assert_eq!(
4376            bind("s3+https://host/pond/sub/x"),
4377            CredsBinding::Set {
4378                name: "deep".to_owned(),
4379                via: BindVia::Scope
4380            },
4381        );
4382        assert_eq!(
4383            bind("s3+https://host/pond/other"),
4384            CredsBinding::Set {
4385                name: "bucket".to_owned(),
4386                via: BindVia::Scope
4387            },
4388        );
4389        // Segment boundary: `/pond` does not match `/pond-2`.
4390        assert_eq!(
4391            bind("s3+https://host/pond-2"),
4392            CredsBinding::Set {
4393                name: "all".to_owned(),
4394                via: BindVia::CatchAll
4395            },
4396        );
4397        // No cross-scheme normalization: the scoped sets don't match s3://.
4398        assert_eq!(
4399            bind("s3://pond/sub"),
4400            CredsBinding::Set {
4401                name: "all".to_owned(),
4402                via: BindVia::CatchAll
4403            },
4404        );
4405        // Default-port spelling matches the portless scope.
4406        assert_eq!(
4407            bind("s3+https://host:443/pond/x"),
4408            CredsBinding::Set {
4409                name: "bucket".to_owned(),
4410                via: BindVia::Scope
4411            },
4412        );
4413        // `?creds=` pointer beats every scope...
4414        assert_eq!(
4415            bind("s3+https://host/pond/sub/x?creds=all"),
4416            CredsBinding::Set {
4417                name: "all".to_owned(),
4418                via: BindVia::Pointer
4419            },
4420        );
4421        // ...and a pointer to a missing set is an error, not a fallback.
4422        let err = StorageUrl::parse("s3://b/p?creds=nope")
4423            .unwrap()
4424            .resolve(&creds)
4425            .expect_err("missing set")
4426            .to_string();
4427        assert!(err.contains("creds=nope"), "got: {err}");
4428
4429        // No sets at all -> ambient chain; local URLs skip resolution.
4430        let empty = BTreeMap::new();
4431        assert_eq!(
4432            StorageUrl::parse("s3://b/p")
4433                .unwrap()
4434                .resolve(&empty)
4435                .unwrap()
4436                .binding,
4437            CredsBinding::Ambient,
4438        );
4439        assert_eq!(
4440            StorageUrl::parse("/srv/pond")
4441                .unwrap()
4442                .resolve(&creds)
4443                .unwrap()
4444                .binding,
4445            CredsBinding::NotApplicable,
4446        );
4447    }
4448
4449    #[test]
4450    fn unmatched_sets_are_reported_only_on_remote_invocations() {
4451        let mut creds = BTreeMap::new();
4452        creds.insert("used".to_owned(), set(Some("s3://bucket/")));
4453        creds.insert("idle".to_owned(), set(Some("s3://other/")));
4454
4455        let remote = StorageUrl::parse("s3://bucket/p")
4456            .unwrap()
4457            .resolve(&creds)
4458            .unwrap();
4459        assert_eq!(unmatched_creds_sets(&[&remote], &creds), vec!["idle"]);
4460
4461        // A purely local invocation must not nag about remote-only sets.
4462        let local = StorageUrl::parse("/srv/pond")
4463            .unwrap()
4464            .resolve(&creds)
4465            .unwrap();
4466        assert!(unmatched_creds_sets(&[&local], &creds).is_empty());
4467    }
4468
4469    #[test]
4470    fn secrets_materialize_from_file_and_command() {
4471        let dir = TempDir::new().unwrap();
4472        let key_path = dir.path().join("key");
4473        std::fs::write(&key_path, "from-file\n").unwrap();
4474        let mut creds = BTreeMap::new();
4475        creds.insert(
4476            "default".to_owned(),
4477            CredsSet {
4478                access_key_id_file: Some(key_path),
4479                // Two trailing newlines: exactly one is stripped.
4480                secret_access_key_command: Some("printf 'from-command\\n\\n'".to_owned()),
4481                ..CredsSet::default()
4482            },
4483        );
4484        let url = StorageUrl::parse("s3://bucket/p").unwrap();
4485        let resolved = url.resolve(&creds).unwrap();
4486        assert_eq!(
4487            opts(&resolved, "access_key_id").as_deref(),
4488            Some("from-file")
4489        );
4490        assert_eq!(
4491            opts(&resolved, "secret_access_key").as_deref(),
4492            Some("from-command\n"),
4493        );
4494
4495        // A failing command surfaces its text and exit status.
4496        let mut failing = BTreeMap::new();
4497        failing.insert(
4498            "default".to_owned(),
4499            CredsSet {
4500                secret_access_key_command: Some("exit 3".to_owned()),
4501                ..CredsSet::default()
4502            },
4503        );
4504        let err = url
4505            .resolve(&failing)
4506            .expect_err("command must fail")
4507            .to_string();
4508        assert!(err.contains("exit 3"), "got: {err}");
4509
4510        // The command cache: one subprocess per command text per process.
4511        let marker = dir.path().join("runs");
4512        let command = format!("echo run >> {} && echo secret", marker.display());
4513        let mut counted = BTreeMap::new();
4514        counted.insert(
4515            "default".to_owned(),
4516            CredsSet {
4517                secret_access_key_command: Some(command),
4518                ..CredsSet::default()
4519            },
4520        );
4521        url.resolve(&counted).unwrap();
4522        url.resolve(&counted).unwrap();
4523        let runs = std::fs::read_to_string(&marker).unwrap();
4524        assert_eq!(runs.lines().count(), 1, "command must run exactly once");
4525    }
4526
4527    #[test]
4528    fn check_errors_classify_by_kind_and_binding() {
4529        let auth_error = || object_store::Error::Unauthenticated {
4530            path: "k".to_owned(),
4531            source: "denied".into(),
4532        };
4533        let bound = CredsBinding::Set {
4534            name: "work".to_owned(),
4535            via: BindVia::Scope,
4536        };
4537        // Auth-class error with a bound set names the set...
4538        match classify_check_error(auth_error(), &bound, "put") {
4539            CheckFailure::Auth { set, .. } => assert_eq!(set, "work"),
4540            other => panic!("want Auth, got {other:?}"),
4541        }
4542        // ...and without one, points at the (empty) ambient chain.
4543        assert!(matches!(
4544            classify_check_error(auth_error(), &CredsBinding::Ambient, "put"),
4545            CheckFailure::NoCreds { .. },
4546        ));
4547        let denied = object_store::Error::PermissionDenied {
4548            path: "k".to_owned(),
4549            source: "403".into(),
4550        };
4551        assert!(matches!(
4552            classify_check_error(denied, &bound, "put"),
4553            CheckFailure::Auth { .. },
4554        ));
4555        // Anything else is I/O, set or no set.
4556        let missing = object_store::Error::NotFound {
4557            path: "k".to_owned(),
4558            source: "404".into(),
4559        };
4560        assert!(matches!(
4561            classify_check_error(missing, &bound, "get"),
4562            CheckFailure::Io { .. },
4563        ));
4564        // Lance wraps an empty-creds chain as a `Generic` error, never the
4565        // typed `Unauthenticated`; the rendered `CredentialsNotLoaded` is the
4566        // signal. Bound -> Auth (the set is wrong), unbound -> NoCreds.
4567        let no_creds = || object_store::Error::Generic {
4568            store: "S3",
4569            source: "Failed to get AWS credentials: CredentialsNotLoaded".into(),
4570        };
4571        assert!(matches!(
4572            classify_check_error(no_creds(), &bound, "put"),
4573            CheckFailure::Auth { .. },
4574        ));
4575        assert!(matches!(
4576            classify_check_error(no_creds(), &CredsBinding::Ambient, "put"),
4577            CheckFailure::NoCreds { .. },
4578        ));
4579    }
4580
4581    #[test]
4582    fn concise_cause_strips_upstream_noise_to_one_line() {
4583        // The shape Lance actually produces: bug-report boilerplate, the real
4584        // cause, an internal source location, then the same text re-printed.
4585        let inner = "Encountered internal error. Please file a bug report at \
4586                     https://github.com/lance-format/lance/issues. Failed to get AWS \
4587                     credentials: CredentialsNotLoaded, <WORKSPACE>/src/object_store/providers/aws.rs:401:21: \
4588                     Encountered internal error. Please file a bug report at \
4589                     https://github.com/lance-format/lance/issues. Failed to get AWS \
4590                     credentials: CredentialsNotLoaded";
4591        let failure = CheckFailure::NoCreds {
4592            source: anyhow!(inner.to_owned()).context("initial conditional put"),
4593        };
4594        let cause = failure.concise_cause().expect("auth-class carries a cause");
4595        assert_eq!(cause, "Failed to get AWS credentials: CredentialsNotLoaded");
4596        // Display carries only the fix-naming lead, no chain.
4597        assert!(
4598            !failure.to_string().contains("file a bug report"),
4599            "lead must not trail the chain: {failure}"
4600        );
4601        // OccUnsupported's detail is already curated into Display.
4602        let occ = CheckFailure::OccUnsupported {
4603            detail: "put-if-none-match ignored".to_owned(),
4604        };
4605        assert!(occ.concise_cause().is_none());
4606        // Oversized single-line causes middle-truncate, keeping the tail
4607        // (wrapped transport errors put the root cause at the end).
4608        let long = CheckFailure::Io {
4609            source: anyhow!(format!("{} dns error: lookup failed", "x".repeat(500))),
4610        };
4611        let cause = long.concise_cause().expect("io carries a cause");
4612        assert!(cause.contains(" ... "), "long causes truncate: {cause}");
4613        assert!(
4614            cause.ends_with("dns error: lookup failed"),
4615            "the tail survives: {cause}"
4616        );
4617    }
4618
4619    #[tokio::test]
4620    async fn storage_check_passes_on_memory_backend() {
4621        let resolved = StorageUrl::parse("memory://check/probe")
4622            .unwrap()
4623            .resolve(&BTreeMap::new())
4624            .unwrap();
4625        storage_check(&resolved).await.expect("memory probe passes");
4626    }
4627
4628    fn stat(bytes: u64) -> FragmentStat {
4629        FragmentStat {
4630            bytes: Some(bytes),
4631            rows: bytes / 1_000,
4632            deleted_rows: 0,
4633        }
4634    }
4635
4636    #[test]
4637    fn compaction_veto_blocks_absorb_keeps_peers() {
4638        // One 665 MiB tail fragment + tiny appends -> vetoed.
4639        let absorb = [stat(665_000_000), stat(1_000_000), stat(2_000_000)];
4640        assert!(!keep_task(&absorb, 64, 0.1));
4641        // Peer-sized merge halves fragment count -> kept.
4642        let peers = [stat(300_000_000), stat(300_000_000)];
4643        assert!(keep_task(&peers, 64, 0.1));
4644        // Remainder reaches largest / COMPACTION_ABSORB_FACTOR -> kept.
4645        let tiered = [stat(400_000), stat(60_000), stat(40_000)];
4646        assert!(keep_task(&tiered, 64, 0.1));
4647    }
4648
4649    #[test]
4650    fn compaction_veto_passes_deletions_and_cap() {
4651        let mut deleting = stat(665_000_000);
4652        deleting.deleted_rows = deleting.rows / 5;
4653        assert!(keep_task(&[deleting, stat(1_000)], 64, 0.1));
4654
4655        let wide: Vec<FragmentStat> = std::iter::once(stat(665_000_000))
4656            .chain(std::iter::repeat_with(|| stat(1_000)).take(63))
4657            .collect();
4658        assert!(keep_task(&wide, 64, 0.1));
4659    }
4660
4661    #[test]
4662    fn compaction_veto_falls_back_to_rows_on_unknown_sizes() {
4663        let mut unknown = stat(665_000_000);
4664        unknown.bytes = None;
4665        // Rows comparison: 665k vs 3k -> still vetoed.
4666        assert!(!keep_task(
4667            &[unknown, stat(1_000_000), stat(2_000_000)],
4668            64,
4669            0.1
4670        ));
4671    }
4672
4673    #[test]
4674    fn cleanup_due_gates_on_version_interval() {
4675        // interval <= 1 always cleans (pond optimize / pond copy / tests).
4676        assert!(cleanup_due(0, 1));
4677        assert!(cleanup_due(7, 1));
4678        assert!(cleanup_due(5, 0));
4679        // interval N: only on multiples (the amortized pond sync path).
4680        assert!(cleanup_due(0, 16));
4681        assert!(cleanup_due(16, 16));
4682        assert!(cleanup_due(48, 16));
4683        assert!(!cleanup_due(15, 16));
4684        assert!(!cleanup_due(17, 16));
4685        assert!(!cleanup_due(31, 16));
4686    }
4687
4688    #[test]
4689    fn derived_target_rows_tracks_row_size_and_clamps() {
4690        // ~1.3 KiB rows -> ~100k-row target (half the byte budget, for freeze
4691        // headroom over the 256 MiB output cap).
4692        let parts_like = [FragmentStat {
4693            bytes: Some(665_000_000),
4694            rows: 511_000,
4695            deleted_rows: 0,
4696        }];
4697        let target = derived_target_rows(&parts_like);
4698        assert!((80_000..150_000).contains(&target), "{target}");
4699        // No usable sizes -> Lance default.
4700        let unknown = [FragmentStat {
4701            bytes: None,
4702            rows: 511_000,
4703            deleted_rows: 0,
4704        }];
4705        assert_eq!(
4706            derived_target_rows(&unknown),
4707            MAX_TARGET_ROWS_PER_FRAGMENT as usize
4708        );
4709        // Tiny rows clamp at the ceiling, huge rows at the floor.
4710        let tiny = [FragmentStat {
4711            bytes: Some(1_000_000),
4712            rows: 100_000,
4713            deleted_rows: 0,
4714        }];
4715        assert_eq!(
4716            derived_target_rows(&tiny),
4717            MAX_TARGET_ROWS_PER_FRAGMENT as usize
4718        );
4719        let huge = [FragmentStat {
4720            bytes: Some(1_000_000_000),
4721            rows: 100,
4722            deleted_rows: 0,
4723        }];
4724        assert_eq!(
4725            derived_target_rows(&huge),
4726            MIN_TARGET_ROWS_PER_FRAGMENT as usize
4727        );
4728    }
4729
4730    #[test]
4731    fn namespace_error_code_walks_wrapped_chain() {
4732        let direct = lance::Error::namespace_source(Box::new(NamespaceError::TableNotFound {
4733            message: "missing".into(),
4734        }));
4735        assert!(is_namespace_error_code(&direct, ErrorCode::TableNotFound));
4736
4737        let wrapped = lance::Error::namespace_source(Box::new(direct));
4738        assert!(is_namespace_error_code(&wrapped, ErrorCode::TableNotFound));
4739
4740        let other_code =
4741            lance::Error::namespace_source(Box::new(NamespaceError::NamespaceNotFound {
4742                message: "nope".into(),
4743            }));
4744        assert!(!is_namespace_error_code(
4745            &other_code,
4746            ErrorCode::TableNotFound
4747        ));
4748
4749        let not_namespace = lance::Error::internal("unrelated");
4750        assert!(!is_namespace_error_code(
4751            &not_namespace,
4752            ErrorCode::TableNotFound
4753        ));
4754    }
4755
4756    /// Round-trip: opening a fresh data dir through `lance-namespace`
4757    /// produces all three tables, and `Handle::scan` returns an empty batch
4758    /// for each (no spurious schema mismatch, no namespace error).
4759    #[tokio::test]
4760    async fn store_opens_via_namespace_and_scan_works() -> Result<()> {
4761        let temp = TempDir::new()?;
4762        let url = Url::from_directory_path(temp.path())
4763            .map_err(|()| anyhow::anyhow!("temp path is not absolute"))?;
4764        let handle = Handle::open(&url).await?;
4765        // Each table has its own PK column; project the canonical one so the
4766        // scan is exercised end-to-end (catalog -> dataset -> scanner -> batch).
4767        let cases: [(Table, &[&str]); 3] = [
4768            (Table::Sessions, &["id"]),
4769            (Table::Messages, &["id"]),
4770            (Table::Parts, &["id"]),
4771        ];
4772        for (table, projection) in cases {
4773            let scanner = handle
4774                .scan(table, ScanOpts::project_only(projection))
4775                .await?;
4776            let batch = scanner.try_into_batch().await?;
4777            assert_eq!(batch.num_rows(), 0, "fresh table should be empty");
4778        }
4779        Ok(())
4780    }
4781}