Skip to main content

lancedb/
table.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The LanceDB Authors
3
4//! LanceDB Table APIs
5
6use crate::blob::BlobFile;
7use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
8use arrow_schema::{Schema, SchemaRef};
9use async_trait::async_trait;
10use datafusion_execution::TaskContext;
11use datafusion_expr::Expr;
12use datafusion_physical_plan::ExecutionPlan;
13use datafusion_physical_plan::display::DisplayableExecutionPlan;
14use futures::StreamExt;
15use futures::stream::FuturesUnordered;
16pub use lance::dataset::ColumnAlteration;
17pub use lance::dataset::NewColumnTransform;
18pub use lance::dataset::ReadParams;
19pub use lance::dataset::Version;
20use lance::dataset::WriteMode;
21use lance::dataset::builder::DatasetBuilder;
22use lance::dataset::{InsertBuilder, WriteParams};
23use lance::index::DatasetIndexExt;
24use lance::index::scalar::load_segment_params;
25use lance::io::{ObjectStoreParams, WrappingObjectStore};
26use lance_datafusion::utils::StreamingWriteSource;
27use lance_index::IndexCriteria;
28use lance_io::object_store::{LanceNamespaceStorageOptionsProvider, StorageOptionsAccessor};
29pub use query::AnyQuery;
30
31use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
32use lance_index::scalar::InvertedIndexParams;
33use lance_index::scalar::inverted::query::collect_query_tokens;
34use lance_namespace::LanceNamespace;
35use lance_namespace::error::NamespaceError;
36use lance_namespace::models::DescribeTableRequest;
37use lance_table::format::Manifest;
38use lance_table::io::commit::CommitHandler;
39use lance_table::io::commit::ManifestNamingScheme;
40use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
41use serde::{Deserialize, Serialize};
42use std::collections::{HashMap, HashSet};
43use std::format;
44use std::path::Path;
45use std::sync::Arc;
46
47use crate::connection::NamespaceClientPushdownOperation;
48
49use crate::DistanceType;
50use crate::blob::BlobRangeRequest;
51use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions};
52use crate::database::Database;
53use crate::database::listing::LANCE_FILE_EXTENSION;
54use crate::database::read_freshness::TableFreshness;
55use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry};
56use crate::error::{Error, Result};
57use crate::index::IndexStatistics;
58use crate::index::{Index, IndexBuilder};
59use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType};
60use crate::job::Job;
61use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery};
62use crate::table::datafusion::insert::InsertExec;
63use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
64
65use self::dataset::DatasetConsistencyWrapper;
66use self::merge::MergeInsertBuilder;
67
68pub mod add_columns;
69mod add_data;
70pub mod branch_merge;
71pub mod checkpoint;
72mod create_index;
73pub mod datafusion;
74pub(crate) mod dataset;
75pub mod delete;
76pub mod lsm_stats;
77pub mod merge;
78pub mod optimize;
79mod primary_key;
80pub mod query;
81pub mod schema_evolution;
82pub mod update;
83pub mod write_progress;
84use crate::index::waiter::wait_for_index;
85pub use add_columns::AddColumnsBuilder;
86#[cfg(feature = "remote")]
87pub(crate) use add_data::PreprocessingOutput;
88pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior};
89pub use branch_merge::{
90    BranchDiff, ColumnChange, ColumnSummary, IndexSummary, MergeBlocker, MergeBlockerCode,
91    MergeBranchResult, MergeBranchStatus, MergePreview, RowCountSummary,
92};
93pub use chrono::Duration;
94pub use delete::DeleteResult;
95use futures::future::join_all;
96pub use lance::dataset::refs::{BranchContents, Ref, TagContents, Tags as LanceTags};
97pub use lance::dataset::scanner::DatasetRecordBatchStream;
98pub use lance_index::optimize::OptimizeOptions;
99pub use lsm_stats::{BucketStats, GenerationStats, LsmStats, MemtableStats};
100pub use optimize::{CompactionOptions, OptimizeAction, OptimizeStats};
101pub use schema_evolution::{
102    AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate,
103    UpdateFieldMetadataResult,
104};
105use serde_with::skip_serializing_none;
106pub use update::{UpdateBuilder, UpdateResult};
107
108/// Walk a boxed error chain to find the innermost `NamespaceError`.
109///
110/// Callers like `DatasetBuilder::from_namespace` re-wrap their inner namespace error
111/// inside a fresh `lance::Error::Namespace`, so a single downcast at the top level
112/// won't find it. This walks `.source()` to unwrap arbitrarily nested layers.
113fn find_namespace_error<'a>(
114    err: &'a (dyn std::error::Error + 'static),
115) -> Option<&'a NamespaceError> {
116    let mut current: Option<&(dyn std::error::Error + 'static)> = Some(err);
117    while let Some(e) = current {
118        if let Some(ns_err) = e.downcast_ref::<NamespaceError>() {
119            return Some(ns_err);
120        }
121        current = e.source();
122    }
123    None
124}
125
126/// Map a `lance::Error` coming from a `lance-namespace` call into a `lancedb::Error`,
127/// preserving the fine-grained namespace error code (e.g. `TableNotFound`,
128/// `TableAlreadyExists`). Errors that aren't recognized namespace error variants fall
129/// through to a generic runtime error rather than `TableNotFound`/`TableAlreadyExists`.
130pub(crate) fn map_namespace_lance_error(err: lance::Error, table_name: &str) -> Error {
131    if let Some(code) = find_namespace_error(&err).map(NamespaceError::code) {
132        match code {
133            lance_namespace::error::ErrorCode::TableNotFound => {
134                return Error::TableNotFound {
135                    name: table_name.to_string(),
136                    source: Box::new(err),
137                };
138            }
139            lance_namespace::error::ErrorCode::TableAlreadyExists => {
140                return Error::TableAlreadyExists {
141                    name: table_name.to_string(),
142                };
143            }
144            _ => {}
145        }
146    }
147    match err {
148        lance::Error::Namespace { source, .. } => Error::Runtime {
149            message: format!("Namespace error: {}", source),
150        },
151        other => other.into(),
152    }
153}
154
155/// Map a `lance::Error::DatasetNotFound` for the table at `uri` into a `lancedb::Error`.
156///
157/// Lance reports "there is nothing at this location" and "there is a table directory
158/// here but nothing loadable inside it" with the same error. Only the first is a
159/// `TableNotFound`: a `<name>.lance` directory left behind by an interrupted drop and
160/// re-create is still reported by `Connection::table_names`, so callers need to be able
161/// to tell "never existed" from "exists but is broken".
162///
163/// See <https://github.com/lancedb/lancedb/issues/3127>.
164async fn map_dataset_not_found(
165    uri: &str,
166    name: &str,
167    params: ReadParams,
168    err: lance::Error,
169) -> Error {
170    let name = name.to_string();
171    let source = Box::new(err);
172    if table_dir_exists(uri, params).await.unwrap_or(false) {
173        Error::TableCorrupted { name, source }
174    } else {
175        Error::TableNotFound { name, source }
176    }
177}
178
179/// Whether a table directory is present at `uri`, even though no dataset could be
180/// loaded from it.
181///
182/// This looks for a `<name>.lance` entry in the parent directory, which is exactly what
183/// `ListingDatabase::table_names` lists, so the two APIs agree on whether a table is
184/// present. Probing `uri` itself would not work: object stores have no empty
185/// directories to probe, and on a local filesystem the interesting case is precisely an
186/// empty directory.
187async fn table_dir_exists(uri: &str, params: ReadParams) -> Result<bool> {
188    let (object_store, path, _) = DatasetBuilder::from_uri(uri)
189        .with_read_params(params)
190        .build_object_store()
191        .await?;
192    // Only `*.lance` entries are ever reported as tables, so nothing else can produce
193    // the list-then-open mismatch this guards against.
194    if path.extension() != Some(LANCE_FILE_EXTENSION) {
195        return Ok(false);
196    }
197    let (Some(parent), Some(dir_name)) = (path.parent(), path.filename()) else {
198        return Ok(false);
199    };
200    let entries = object_store.read_dir(parent).await?;
201    Ok(entries.iter().any(|entry| entry.as_str() == dir_name))
202}
203
204/// Defines the type of column
205#[derive(Debug, Clone, Serialize, Deserialize)]
206pub enum ColumnKind {
207    /// Columns populated by data from the user (this is the most common case)
208    Physical,
209    /// Columns populated by applying an embedding function to the input
210    Embedding(EmbeddingDefinition),
211}
212
213/// Defines a column in a table
214#[derive(Debug, Clone, Serialize, Deserialize)]
215pub struct ColumnDefinition {
216    /// The source of the column data
217    pub kind: ColumnKind,
218}
219
220#[derive(Debug, Clone)]
221pub struct TableDefinition {
222    pub column_definitions: Vec<ColumnDefinition>,
223    pub schema: SchemaRef,
224}
225
226impl TableDefinition {
227    pub fn new(schema: SchemaRef, column_definitions: Vec<ColumnDefinition>) -> Self {
228        Self {
229            column_definitions,
230            schema,
231        }
232    }
233
234    pub fn new_from_schema(schema: SchemaRef) -> Self {
235        let column_definitions = schema
236            .fields()
237            .iter()
238            .map(|_| ColumnDefinition {
239                kind: ColumnKind::Physical,
240            })
241            .collect();
242        Self::new(schema, column_definitions)
243    }
244
245    pub fn try_from_rich_schema(schema: SchemaRef) -> Result<Self> {
246        let column_definitions = schema.metadata.get("lancedb::column_definitions");
247        if let Some(column_definitions) = column_definitions {
248            let column_definitions: Vec<ColumnDefinition> =
249                serde_json::from_str(column_definitions).map_err(|e| Error::Runtime {
250                    message: format!("Failed to deserialize column definitions: {}", e),
251                })?;
252            Ok(Self::new(schema, column_definitions))
253        } else {
254            let column_definitions = schema
255                .fields()
256                .iter()
257                .map(|_| ColumnDefinition {
258                    kind: ColumnKind::Physical,
259                })
260                .collect();
261            Ok(Self::new(schema, column_definitions))
262        }
263    }
264
265    pub fn into_rich_schema(self) -> SchemaRef {
266        // We have full control over the structure of column definitions.  This should
267        // not fail, except for a bug
268        let lancedb_metadata = serde_json::to_string(&self.column_definitions).unwrap();
269        let mut schema_with_metadata = (*self.schema).clone();
270        schema_with_metadata
271            .metadata
272            .insert("lancedb::column_definitions".to_string(), lancedb_metadata);
273        Arc::new(schema_with_metadata)
274    }
275}
276
277/// Describes what happens when a vector either contains NaN or
278/// does not have enough values
279#[derive(Clone, Debug, Default)]
280#[allow(dead_code)] // https://github.com/lancedb/lancedb/issues/992
281enum BadVectorHandling {
282    /// An error is returned
283    #[default]
284    Error,
285    /// The offending row is droppped
286    Drop,
287    /// The invalid/missing items are replaced by fill_value
288    Fill(f32),
289    /// The invalid items are replaced by NULL
290    None,
291}
292
293/// Options to use when writing data
294#[derive(Clone, Debug, Default)]
295pub struct WriteOptions {
296    // Coming soon: https://github.com/lancedb/lancedb/issues/992
297    // /// What behavior to take if the data contains invalid vectors
298    // pub on_bad_vectors: BadVectorHandling,
299    /// Advanced parameters that can be used to customize table creation
300    ///
301    /// Overlapping `OpenTableBuilder` options (e.g. [AddDataBuilder::mode]) will take
302    /// precedence over their counterparts in `WriteOptions` (e.g. [WriteParams::mode]).
303    pub lance_write_params: Option<WriteParams>,
304}
305
306/// Filters that can be used to limit the rows returned by a query
307pub enum Filter {
308    /// A SQL filter string
309    Sql(String),
310    /// A Datafusion logical expression
311    Datafusion(Expr),
312}
313
314/// A predicate for filtering rows in delete operations.
315///
316/// Accepts either a SQL string or a DataFusion [`Expr`]. Use the [`From`]
317/// implementations to convert from `&str` or `&Expr` automatically.
318/// See [`Table::delete`] for usage examples.
319pub enum Predicate<'a> {
320    /// A SQL predicate string
321    String(&'a str),
322    /// A DataFusion logical expression
323    Expr(&'a Expr),
324}
325
326impl<'a> From<&'a str> for Predicate<'a> {
327    fn from(s: &'a str) -> Self {
328        Predicate::String(s)
329    }
330}
331
332impl<'a> From<&'a String> for Predicate<'a> {
333    fn from(s: &'a String) -> Self {
334        Predicate::String(s.as_str())
335    }
336}
337
338impl<'a> From<&'a Expr> for Predicate<'a> {
339    fn from(e: &'a Expr) -> Self {
340        Predicate::Expr(e)
341    }
342}
343
344#[async_trait]
345pub trait Tags: Send + Sync {
346    /// List the tags of the table.
347    async fn list(&self) -> Result<HashMap<String, TagContents>>;
348
349    /// Get the version of the table referenced by a tag.
350    async fn get_version(&self, tag: &str) -> Result<u64>;
351
352    /// Create a new tag for the given version of the table.
353    async fn create(&mut self, tag: &str, version: u64) -> Result<()>;
354
355    /// Delete a tag from the table.
356    async fn delete(&mut self, tag: &str) -> Result<()>;
357
358    /// Update an existing tag to point to a new version of the table.
359    async fn update(&mut self, tag: &str, version: u64) -> Result<()>;
360}
361
362pub use self::merge::MergeResult;
363
364/// Specification selecting Lance's MemWAL LSM-style write path for
365/// `merge_insert`.
366///
367/// Construct via [`LsmWriteSpec::bucket`], [`LsmWriteSpec::identity`], or
368/// [`LsmWriteSpec::unsharded`], then optionally chain
369/// [`LsmWriteSpec::with_maintained_indexes`] (indexes the MemWAL keeps up to
370/// date) and [`LsmWriteSpec::with_writer_config_defaults`] (default
371/// `ShardWriter` configuration recorded in the MemWAL index).
372///
373/// Install a spec with [`Table::set_lsm_write_spec`] and remove it with
374/// [`Table::unset_lsm_write_spec`]. The actual `merge_insert` dispatch
375/// onto the MemWAL writer is a follow-up.
376#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
377pub enum LsmWriteSpec {
378    /// Hash-bucket sharding by a scalar column.
379    ///
380    /// `column` must be a non-nested column with a supported scalar type.
381    /// `num_buckets` must be in `[1, 1024]`.
382    /// Iceberg-compatible Murmur3-x86-32 (seed 0) is used so each row's
383    /// `bucket(column, num_buckets)` value is stable across processes.
384    Bucket {
385        column: String,
386        num_buckets: u32,
387        /// Names of indexes (already created on the table) that the
388        /// MemWAL should maintain in-memory as rows are appended.
389        maintained_indexes: Vec<String>,
390        /// Default `ShardWriter` configuration recorded in the MemWAL index.
391        writer_config_defaults: HashMap<String, String>,
392    },
393    /// Identity sharding — shard by the raw value of `column`.
394    ///
395    /// Use this when the data is already partitioned by `column`; each
396    /// distinct value of `column` becomes its own shard.
397    Identity {
398        column: String,
399        /// Names of indexes (already created on the table) that the
400        /// MemWAL should maintain in-memory as rows are appended.
401        maintained_indexes: Vec<String>,
402        /// Default `ShardWriter` configuration recorded in the MemWAL index.
403        writer_config_defaults: HashMap<String, String>,
404    },
405    /// No sharding — every `merge_insert` call writes to a single MemWAL shard.
406    Unsharded {
407        /// Names of indexes (already created on the table) that the
408        /// MemWAL should maintain in-memory as rows are appended.
409        maintained_indexes: Vec<String>,
410        /// Default `ShardWriter` configuration recorded in the MemWAL index.
411        writer_config_defaults: HashMap<String, String>,
412    },
413}
414
415impl LsmWriteSpec {
416    /// Construct a hash-bucket sharding spec with no maintained indexes.
417    pub fn bucket(column: impl Into<String>, num_buckets: u32) -> Self {
418        Self::Bucket {
419            column: column.into(),
420            num_buckets,
421            maintained_indexes: Vec::new(),
422            writer_config_defaults: HashMap::new(),
423        }
424    }
425
426    /// Construct an identity-sharding spec (shard by the raw value of
427    /// `column`) with no maintained indexes.
428    ///
429    /// `column` must be a deterministic function of the unenforced primary
430    /// key: every row with a given primary key must always produce the same
431    /// `column` value. MemWAL dedups upserts by primary key but tracks
432    /// generations per shard, so if the same key is written with two
433    /// different `column` values its versions land in different shards and a
434    /// stale value can win. Typically `column` is the primary key itself, or
435    /// a stable attribute of it (e.g. a tenant id).
436    pub fn identity(column: impl Into<String>) -> Self {
437        Self::Identity {
438            column: column.into(),
439            maintained_indexes: Vec::new(),
440            writer_config_defaults: HashMap::new(),
441        }
442    }
443
444    /// Construct an unsharded spec with no maintained indexes.
445    pub fn unsharded() -> Self {
446        Self::Unsharded {
447            maintained_indexes: Vec::new(),
448            writer_config_defaults: HashMap::new(),
449        }
450    }
451
452    /// Replace the list of indexes the MemWAL should keep up to date as
453    /// rows are appended. Each name must reference an index that already
454    /// exists on the table at the time `set_lsm_write_spec` is called.
455    pub fn with_maintained_indexes<I, S>(mut self, indexes: I) -> Self
456    where
457        I: IntoIterator<Item = S>,
458        S: Into<String>,
459    {
460        let v: Vec<String> = indexes.into_iter().map(Into::into).collect();
461        match &mut self {
462            Self::Bucket {
463                maintained_indexes, ..
464            }
465            | Self::Identity {
466                maintained_indexes, ..
467            }
468            | Self::Unsharded {
469                maintained_indexes, ..
470            } => *maintained_indexes = v,
471        }
472        self
473    }
474
475    /// Replace the default `ShardWriter` configuration recorded in the MemWAL
476    /// index, so every writer starts from the same defaults. Keys are
477    /// `ShardWriter` config field names (`Duration` knobs use a `_ms` suffix);
478    /// values are their string encodings.
479    pub fn with_writer_config_defaults<I, K, V>(mut self, defaults: I) -> Self
480    where
481        I: IntoIterator<Item = (K, V)>,
482        K: Into<String>,
483        V: Into<String>,
484    {
485        let m: HashMap<String, String> = defaults
486            .into_iter()
487            .map(|(k, v)| (k.into(), v.into()))
488            .collect();
489        match &mut self {
490            Self::Bucket {
491                writer_config_defaults,
492                ..
493            }
494            | Self::Identity {
495                writer_config_defaults,
496                ..
497            }
498            | Self::Unsharded {
499                writer_config_defaults,
500                ..
501            } => *writer_config_defaults = m,
502        }
503        self
504    }
505
506    /// Borrow the list of index names this spec asks MemWAL to maintain.
507    pub fn maintained_indexes(&self) -> &[String] {
508        match self {
509            Self::Bucket {
510                maintained_indexes, ..
511            }
512            | Self::Identity {
513                maintained_indexes, ..
514            }
515            | Self::Unsharded {
516                maintained_indexes, ..
517            } => maintained_indexes,
518        }
519    }
520
521    /// Borrow the default `ShardWriter` configuration recorded by this spec.
522    pub fn writer_config_defaults(&self) -> &HashMap<String, String> {
523        match self {
524            Self::Bucket {
525                writer_config_defaults,
526                ..
527            }
528            | Self::Identity {
529                writer_config_defaults,
530                ..
531            }
532            | Self::Unsharded {
533                writer_config_defaults,
534                ..
535            } => writer_config_defaults,
536        }
537    }
538}
539
540/// A token produced by the tokenizer configured on a full-text search index.
541#[derive(Debug, Clone, PartialEq, Eq)]
542pub struct FtsToken {
543    /// The token text after the index tokenizer has applied its filters.
544    pub text: String,
545    /// The token position used by full-text query matching.
546    pub position: u32,
547}
548
549/// Tokenize a full-text search query using an explicit FTS tokenizer configuration.
550///
551/// This does not require a table or FTS index. Use
552/// [`crate::index::scalar::FtsIndexBuilder`] to supply the same tokenizer
553/// options used when creating an FTS index.
554pub fn tokenize(query: &str, params: &InvertedIndexParams) -> Result<Vec<FtsToken>> {
555    let mut tokenizer = params.build().map_err(|err| Error::InvalidInput {
556        message: format!("Failed to build tokenizer: {}", err),
557    })?;
558    let tokens = collect_query_tokens(query, &mut tokenizer);
559    Ok((0..tokens.len())
560        .map(|idx| FtsToken {
561            text: tokens.get_token(idx).to_string(),
562            position: tokens.position(idx),
563        })
564        .collect())
565}
566
567/// A trait for anything "table-like".  This is used for both native tables (which target
568/// Lance datasets) and remote tables (which target LanceDB cloud)
569///
570/// This trait is still EXPERIMENTAL and subject to change in the future
571#[async_trait]
572pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
573    /// Get a reference to std::any::Any
574    fn as_any(&self) -> &dyn std::any::Any;
575    /// Get the name of the table.
576    fn name(&self) -> &str;
577    /// Get the namespace of the table.
578    fn namespace(&self) -> &[String];
579    /// Get the id of the table
580    ///
581    /// This is the namespace of the table concatenated with the name
582    /// separated by $
583    fn id(&self) -> &str;
584    /// Get the arrow [Schema] of the table.
585    async fn schema(&self) -> Result<SchemaRef>;
586    /// Count the number of rows in this table.
587    async fn count_rows(&self, filter: Option<Filter>) -> Result<usize>;
588    /// Create a physical plan for the query.
589    async fn create_plan(
590        &self,
591        query: &AnyQuery,
592        options: QueryExecutionOptions,
593    ) -> Result<Arc<dyn ExecutionPlan>>;
594    /// Execute a query and return the results as a stream of RecordBatches.
595    async fn query(
596        &self,
597        query: &AnyQuery,
598        options: QueryExecutionOptions,
599    ) -> Result<DatasetRecordBatchStream>;
600    /// Explain the plan for a query.
601    async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result<String> {
602        let plan = self.create_plan(query, Default::default()).await?;
603        let display = DisplayableExecutionPlan::new(plan.as_ref());
604
605        Ok(format!("{}", display.indent(verbose)))
606    }
607    async fn analyze_plan(
608        &self,
609        query: &AnyQuery,
610        options: QueryExecutionOptions,
611    ) -> Result<String>;
612
613    /// Add new records to the table.
614    async fn add(&self, add: AddDataBuilder) -> Result<AddResult>;
615    /// Delete rows from the table matching the given [`Predicate`].
616    async fn delete(&self, predicate: Predicate<'_>) -> Result<DeleteResult>;
617    /// Update rows in the table.
618    async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult>;
619    /// Create an index on the provided column(s).
620    async fn create_index(&self, index: IndexBuilder) -> Result<()>;
621
622    /// Starts index creation, returning a handle to the resulting job.
623    async fn create_index_async(&self, index: IndexBuilder) -> Result<Job>;
624    /// List the indices on the table.
625    async fn list_indices(&self) -> Result<Vec<IndexConfig>>;
626    /// Drop an index from the table.
627    async fn drop_index(&self, name: &str) -> Result<()>;
628    /// Prewarm an index in the table.
629    async fn prewarm_index(&self, name: &str) -> Result<()>;
630    /// Prewarm data for the table.
631    ///
632    /// Currently only supported on remote tables.
633    /// If `columns` is `None`, all columns are prewarmed.
634    async fn prewarm_data(&self, columns: Option<Vec<String>>) -> Result<()>;
635    /// Get statistics about the index.
636    async fn index_stats(&self, index_name: &str) -> Result<Option<IndexStatistics>>;
637    /// Merge insert new records into the table.
638    async fn merge_insert(
639        &self,
640        params: MergeInsertBuilder,
641        new_data: Box<dyn RecordBatchReader + Send>,
642    ) -> Result<MergeResult>;
643    /// Set the unenforced primary key for the table to a single column.
644    ///
645    /// "Unenforced" means LanceDB does not check uniqueness on writes; the
646    /// column is recorded in the schema as the primary key for use by
647    /// features such as `merge_insert`. Only single-column primary keys are
648    /// supported, and the key cannot be changed once set.
649    ///
650    /// The default implementation returns `NotSupported`; table types
651    /// backed by a Lance dataset override it.
652    async fn set_unenforced_primary_key(&self, _columns: &[&str]) -> Result<()> {
653        Err(Error::NotSupported {
654            message: "set_unenforced_primary_key is not supported on this table type".into(),
655        })
656    }
657    /// Install an [`LsmWriteSpec`] on this table.
658    ///
659    /// The spec selects Lance's MemWAL LSM-style write path for future
660    /// `merge_insert` calls.
661    ///
662    /// The default implementation returns `NotSupported`. Implementations
663    /// that support the MemWAL LSM write path must override this.
664    async fn set_lsm_write_spec(&self, _spec: LsmWriteSpec) -> Result<()> {
665        Err(Error::NotSupported {
666            message: "set_lsm_write_spec is not supported on this table type".into(),
667        })
668    }
669    /// Remove the [`LsmWriteSpec`] from this table.
670    ///
671    /// This is a no-op if no spec is currently set.
672    ///
673    /// The default implementation returns `NotSupported`. Implementations
674    /// that support the MemWAL LSM write path must override this.
675    async fn unset_lsm_write_spec(&self) -> Result<()> {
676        Err(Error::NotSupported {
677            message: "unset_lsm_write_spec is not supported on this table type".into(),
678        })
679    }
680    /// Read the [`LsmWriteSpec`] currently installed on this table, returning
681    /// `None` when the MemWAL LSM write path is not enabled.
682    ///
683    /// The default implementation returns `NotSupported`. Implementations that
684    /// support the MemWAL LSM write path must override this.
685    async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
686        Err(Error::NotSupported {
687            message: "get_lsm_write_spec is not supported on this table type".into(),
688        })
689    }
690    /// Seal every bucket's active memtable into L0.
691    ///
692    /// The default implementation returns `NotSupported`.
693    async fn flush_lsm(&self) -> Result<()> {
694        Err(Error::NotSupported {
695            message: "flush_lsm is not supported on this table type".into(),
696        })
697    }
698    /// Trigger a background L0 → base compaction pass per bucket.
699    ///
700    /// The default implementation returns `NotSupported`.
701    async fn compact_lsm(&self) -> Result<()> {
702        Err(Error::NotSupported {
703            message: "compact_lsm is not supported on this table type".into(),
704        })
705    }
706    /// Read live LSM state, or `None` when the LSM write path is not
707    /// enabled for this table.
708    ///
709    /// The default implementation returns `NotSupported`.
710    async fn get_lsm_stats(&self, _include_generation_rows: bool) -> Result<Option<LsmStats>> {
711        Err(Error::NotSupported {
712            message: "get_lsm_stats is not supported on this table type".into(),
713        })
714    }
715    /// Drain and close any cached MemWAL shard writers for this table.
716    ///
717    /// The default implementation is a no-op; table types that maintain
718    /// MemWAL shard writers override it.
719    async fn close_lsm_writers(&self) -> Result<()> {
720        Ok(())
721    }
722    /// Names of the blob v2 columns in this table, in declaration order.
723    async fn blob_columns(&self) -> Result<Vec<String>> {
724        Err(Error::NotSupported {
725            message: "blob_columns is not supported on this table type".into(),
726        })
727    }
728    /// Materialize blob bytes for the given row ids. See [`Table::fetch_blobs`].
729    async fn fetch_blobs(&self, _column: &str, _row_ids: &[u64]) -> Result<LargeBinaryArray> {
730        Err(Error::NotSupported {
731            message: "fetch_blobs is not supported on this table type".into(),
732        })
733    }
734    /// Materialize blob-local ranges. See [`Table::fetch_blob_ranges`].
735    async fn fetch_blob_ranges(
736        &self,
737        _column: &str,
738        _requests: &[BlobRangeRequest],
739    ) -> Result<LargeBinaryArray> {
740        Err(Error::NotSupported {
741            message: "fetch_blob_ranges is not supported on this table type".into(),
742        })
743    }
744    /// Open lazy blob handles for the given row ids. See [`Table::fetch_blob_files`].
745    async fn fetch_blob_files(
746        &self,
747        _column: &str,
748        _row_ids: &[u64],
749    ) -> Result<Vec<Option<BlobFile>>> {
750        Err(Error::NotSupported {
751            message: "fetch_blob_files is not supported on this table type".into(),
752        })
753    }
754    /// Gets the table tag manager.
755    async fn tags(&self) -> Result<Box<dyn Tags + '_>>;
756    /// Optimize the dataset.
757    async fn optimize(&self, action: OptimizeAction) -> Result<OptimizeStats>;
758    /// Add columns to the table.
759    async fn add_columns(
760        &self,
761        transforms: NewColumnTransform,
762        read_columns: Option<Vec<String>>,
763    ) -> Result<AddColumnsResult>;
764    /// Alter columns in the table.
765    async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult>;
766    /// Drop columns from the table.
767    async fn drop_columns(&self, columns: &[&str]) -> Result<DropColumnsResult>;
768    /// Get the version of the table.
769    async fn version(&self) -> Result<u64>;
770    /// Checkout a specific version of the table.
771    async fn checkout(&self, version: u64) -> Result<()>;
772    /// Checkout a table version referenced by a tag.
773    /// Tags provide a human-readable way to reference specific versions of the table.
774    async fn checkout_tag(&self, tag: &str) -> Result<()>;
775    /// Checkout the latest version of the table.
776    async fn checkout_latest(&self) -> Result<()>;
777    /// Restore the table to the currently checked out version.
778    async fn restore(&self) -> Result<()>;
779    /// List the versions of the table.
780    async fn list_versions(&self) -> Result<Vec<Version>>;
781    /// Create a new branch from `from` and return a handle scoped to it.
782    async fn create_branch(
783        &self,
784        name: &str,
785        from: lance::dataset::refs::Ref,
786    ) -> Result<Arc<dyn BaseTable>>;
787    /// Check out an existing branch and return a handle scoped to it.
788    async fn checkout_branch(&self, name: &str) -> Result<Arc<dyn BaseTable>>;
789    /// Check out an existing branch at an optional version, returning a handle.
790    ///
791    /// `None` tracks the branch's latest; `Some(v)` pins it to that version
792    /// (read-only). The default implementation composes [`Self::checkout_branch`]
793    /// and [`Self::checkout`]; implementations may override it to resolve the
794    /// `(branch, version)` coordinate in a single manifest read.
795    async fn checkout_branch_version(
796        &self,
797        name: &str,
798        version: Option<u64>,
799    ) -> Result<Arc<dyn BaseTable>> {
800        let branch = self.checkout_branch(name).await?;
801        if let Some(version) = version {
802            branch.checkout(version).await?;
803        }
804        Ok(branch)
805    }
806    /// List the branches of the table.
807    async fn list_branches(&self) -> Result<HashMap<String, BranchContents>>;
808    /// Delete a branch.
809    async fn delete_branch(&self, name: &str) -> Result<()>;
810    /// Diff a branch against main. Remote only.
811    async fn diff_branch(&self, _from_branch: &str) -> Result<BranchDiff> {
812        Err(Error::NotSupported {
813            message: "diff_branch is only supported on remote tables".into(),
814        })
815    }
816    /// Merge a branch into main, or dry-run. Remote only.
817    /// HTTP 409 still returns [`Ok`] with [`MergeBranchStatus::Rejected`].
818    async fn merge_branch(&self, _from_branch: &str, _dry_run: bool) -> Result<MergeBranchResult> {
819        Err(Error::NotSupported {
820            message: "merge_branch is only supported on remote tables".into(),
821        })
822    }
823    /// The branch this handle is scoped to, or `None` for `main`.
824    fn current_branch(&self) -> Option<String>;
825    /// Get the table definition.
826    async fn table_definition(&self) -> Result<TableDefinition>;
827    /// Get the table URI (storage location)
828    async fn uri(&self) -> Result<String>;
829    /// Get the storage options used when opening this table, if any.
830    #[deprecated(since = "0.25.0", note = "Use initial_storage_options() instead")]
831    async fn storage_options(&self) -> Option<HashMap<String, String>>;
832    /// Get the initial storage options that were passed in when opening this table.
833    ///
834    /// For dynamically refreshed options (e.g., credential vending), use [`Self::latest_storage_options`].
835    async fn initial_storage_options(&self) -> Option<HashMap<String, String>>;
836    /// Get the latest storage options, refreshing from provider if configured.
837    ///
838    /// Returns `Ok(Some(options))` if storage options are available (static or refreshed),
839    /// `Ok(None)` if no storage options were configured, or `Err(...)` if refresh failed.
840    async fn latest_storage_options(&self) -> Result<Option<HashMap<String, String>>>;
841    /// Poll until the columns are fully indexed. Will return Error::Timeout if the columns
842    /// are not fully indexed within the timeout.
843    async fn wait_for_index(
844        &self,
845        index_names: &[&str],
846        timeout: std::time::Duration,
847    ) -> Result<()>;
848    /// Get statistics on the table
849    async fn stats(&self) -> Result<TableStatistics>;
850    /// Create an ExecutionPlan for inserting data into the table.
851    ///
852    /// This is used by the DataFusion TableProvider implementation to support
853    /// INSERT INTO statements.
854    async fn create_insert_exec(
855        &self,
856        _input: Arc<dyn datafusion_physical_plan::ExecutionPlan>,
857        _write_params: WriteParams,
858    ) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
859        Err(Error::NotSupported {
860            message: "create_insert_exec not implemented".to_string(),
861        })
862    }
863    /// Update per-field metadata. Merges into existing metadata by default;
864    /// [`FieldMetadataUpdate::remove`] deletes a key and
865    /// [`FieldMetadataUpdate::replace`] swaps the field's whole map.
866    ///
867    /// The default returns `NotSupported`; Lance-backed and remote tables override it.
868    async fn update_field_metadata(
869        &self,
870        _updates: &[FieldMetadataUpdate],
871    ) -> Result<UpdateFieldMetadataResult> {
872        Err(Error::NotSupported {
873            message: "update_field_metadata is not supported on this table type".into(),
874        })
875    }
876}
877
878/// A Table is a collection of strong typed Rows.
879///
880/// The type of the each row is defined in Apache Arrow [Schema].
881#[derive(Clone, Debug)]
882pub struct Table {
883    inner: Arc<dyn BaseTable>,
884    database: Option<Arc<dyn Database>>,
885    embedding_registry: Arc<dyn EmbeddingRegistry>,
886}
887
888#[cfg(all(test, feature = "remote"))]
889mod test_utils {
890    use super::*;
891
892    impl Table {
893        pub fn new_with_handler<T>(
894            name: impl Into<String>,
895            handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
896        ) -> Self
897        where
898            T: Into<reqwest::Body>,
899        {
900            let inner = Arc::new(crate::remote::table::RemoteTable::new_mock(
901                name.into(),
902                handler.clone(),
903                None,
904            ));
905            let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock(handler));
906            Self {
907                inner,
908                database: Some(database),
909                // Registry is unused.
910                embedding_registry: Arc::new(MemoryRegistry::new()),
911            }
912        }
913
914        pub fn new_with_handler_and_interval<T>(
915            name: impl Into<String>,
916            handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
917            read_consistency_interval: Option<std::time::Duration>,
918        ) -> Self
919        where
920            T: Into<reqwest::Body>,
921        {
922            let inner = Arc::new(
923                crate::remote::table::RemoteTable::new_mock_with_consistency_interval(
924                    name.into(),
925                    handler.clone(),
926                    read_consistency_interval,
927                ),
928            );
929            let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock(handler));
930            Self {
931                inner,
932                database: Some(database),
933                // Registry is unused.
934                embedding_registry: Arc::new(MemoryRegistry::new()),
935            }
936        }
937
938        pub fn new_with_handler_version<T>(
939            name: impl Into<String>,
940            version: semver::Version,
941            handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
942        ) -> Self
943        where
944            T: Into<reqwest::Body>,
945        {
946            let inner = Arc::new(crate::remote::table::RemoteTable::new_mock(
947                name.into(),
948                handler.clone(),
949                Some(version),
950            ));
951            let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock(handler));
952            Self {
953                inner,
954                database: Some(database),
955                // Registry is unused.
956                embedding_registry: Arc::new(MemoryRegistry::new()),
957            }
958        }
959
960        pub fn new_with_handler_and_config<T>(
961            name: impl Into<String>,
962            handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
963            config: crate::remote::ClientConfig,
964        ) -> Self
965        where
966            T: Into<reqwest::Body>,
967        {
968            let inner = Arc::new(crate::remote::table::RemoteTable::new_mock_with_config(
969                name.into(),
970                handler.clone(),
971                config.clone(),
972            ));
973            let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock_with_config(
974                handler, config,
975            ));
976            Self {
977                inner,
978                database: Some(database),
979                // Registry is unused.
980                embedding_registry: Arc::new(MemoryRegistry::new()),
981            }
982        }
983
984        pub fn new_with_handler_version_and_config<T>(
985            name: impl Into<String>,
986            version: semver::Version,
987            handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
988            config: crate::remote::ClientConfig,
989        ) -> Self
990        where
991            T: Into<reqwest::Body>,
992        {
993            let inner = Arc::new(
994                crate::remote::table::RemoteTable::new_mock_with_version_and_config(
995                    name.into(),
996                    handler.clone(),
997                    Some(version),
998                    config.clone(),
999                ),
1000            );
1001            let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock_with_config(
1002                handler, config,
1003            ));
1004            Self {
1005                inner,
1006                database: Some(database),
1007                // Registry is unused.
1008                embedding_registry: Arc::new(MemoryRegistry::new()),
1009            }
1010        }
1011    }
1012}
1013
1014impl std::fmt::Display for Table {
1015    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1016        write!(f, "{}", self.inner)
1017    }
1018}
1019
1020impl From<Arc<dyn BaseTable>> for Table {
1021    fn from(inner: Arc<dyn BaseTable>) -> Self {
1022        Self {
1023            inner,
1024            database: None,
1025            embedding_registry: Arc::new(MemoryRegistry::new()),
1026        }
1027    }
1028}
1029
1030impl Table {
1031    pub fn new(inner: Arc<dyn BaseTable>, database: Arc<dyn Database>) -> Self {
1032        Self {
1033            inner,
1034            database: Some(database),
1035            embedding_registry: Arc::new(MemoryRegistry::new()),
1036        }
1037    }
1038
1039    pub fn base_table(&self) -> &Arc<dyn BaseTable> {
1040        &self.inner
1041    }
1042
1043    pub fn database(&self) -> &Arc<dyn Database> {
1044        self.database.as_ref().unwrap()
1045    }
1046
1047    pub fn embedding_registry(&self) -> &Arc<dyn EmbeddingRegistry> {
1048        &self.embedding_registry
1049    }
1050
1051    pub(crate) fn new_with_embedding_registry(
1052        inner: Arc<dyn BaseTable>,
1053        database: Arc<dyn Database>,
1054        embedding_registry: Arc<dyn EmbeddingRegistry>,
1055    ) -> Self {
1056        Self {
1057            inner,
1058            database: Some(database),
1059            embedding_registry,
1060        }
1061    }
1062
1063    /// Cast as [`NativeTable`], or return None it if is not a [`NativeTable`].
1064    ///
1065    /// Warning: This function will be removed soon (features exclusive to NativeTable
1066    ///          will be added to Table)
1067    pub fn as_native(&self) -> Option<&NativeTable> {
1068        self.inner.as_native()
1069    }
1070
1071    /// Get the name of the table.
1072    pub fn name(&self) -> &str {
1073        self.inner.name()
1074    }
1075
1076    /// Get the namespace of the table.
1077    pub fn namespace(&self) -> &[String] {
1078        self.inner.namespace()
1079    }
1080
1081    /// Get the ID of the table (namespace + name joined by '$').
1082    pub fn id(&self) -> &str {
1083        self.inner.id()
1084    }
1085
1086    /// Get the dataset of the table if it is a native table
1087    ///
1088    /// Returns None otherwise
1089    pub fn dataset(&self) -> Option<&dataset::DatasetConsistencyWrapper> {
1090        self.inner.as_native().map(|t| &t.dataset)
1091    }
1092
1093    /// Get the arrow [Schema] of the table.
1094    pub async fn schema(&self) -> Result<SchemaRef> {
1095        self.inner.schema().await
1096    }
1097
1098    /// Count the number of rows in this dataset.
1099    ///
1100    /// # Arguments
1101    ///
1102    /// * `filter` if present, only count rows matching the filter
1103    pub async fn count_rows(&self, filter: Option<String>) -> Result<usize> {
1104        self.inner.count_rows(filter.map(Filter::Sql)).await
1105    }
1106
1107    /// Names of the blob v2 columns in this table, in declaration order.
1108    ///
1109    /// Nested blobs use dotted paths (e.g. `info.blob`). Returns
1110    /// [`Error::NotSupported`] on table types without blob support.
1111    pub async fn blob_columns(&self) -> Result<Vec<String>> {
1112        self.inner.blob_columns().await
1113    }
1114
1115    /// Materialize blob bytes for the given row ids.
1116    ///
1117    /// Output matches `row_ids` in length and order. Null blobs are null;
1118    /// valid empty blobs contain empty byte strings. Prefer
1119    /// [`Self::fetch_blob_files`] for large selections.
1120    ///
1121    /// ```
1122    /// use arrow_array::UInt64Array;
1123    /// use futures::TryStreamExt;
1124    /// use lancedb::query::{ExecutableQuery, QueryBase};
1125    ///
1126    /// # use lancedb::Table;
1127    /// # async fn materialize(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
1128    /// let mut stream = table.query().with_row_id().limit(10).execute().await?;
1129    /// while let Some(batch) = stream.try_next().await? {
1130    ///     let row_ids = batch
1131    ///         .column_by_name("_rowid")
1132    ///         .unwrap()
1133    ///         .as_any()
1134    ///         .downcast_ref::<UInt64Array>()
1135    ///         .unwrap();
1136    ///     let images = table.fetch_blobs("image", row_ids.values()).await?;
1137    ///     let _ = images;
1138    /// }
1139    /// # Ok(())
1140    /// # }
1141    /// ```
1142    ///
1143    /// Returns [`Error::InvalidInput`] when the column does not exist or is
1144    /// not a blob v2 column, and [`Error::NotSupported`] on table types
1145    /// without blob support.
1146    pub async fn fetch_blobs(
1147        &self,
1148        column: impl AsRef<str>,
1149        row_ids: &[u64],
1150    ) -> Result<LargeBinaryArray> {
1151        self.inner.fetch_blobs(column.as_ref(), row_ids).await
1152    }
1153
1154    /// Materialize row-specific ranges from a blob v2 column.
1155    ///
1156    /// Each request contains a row id and a blob-local offset and length.
1157    /// Requests may be duplicated or reordered, including multiple
1158    /// ranges for the same blob. The output has the same length and order as
1159    /// the requests. Null blobs produce null output slots; empty ranges on
1160    /// non-null blobs produce empty byte strings.
1161    ///
1162    /// ```
1163    /// use lancedb::blob::BlobRangeRequest;
1164    ///
1165    /// # use lancedb::Table;
1166    /// # async fn read_ranges(table: &Table, row_id: u64) -> Result<(), Box<dyn std::error::Error>> {
1167    /// let ranges = table
1168    ///     .fetch_blob_ranges(
1169    ///         "image",
1170    ///         [
1171    ///             BlobRangeRequest::new(row_id, 0, 1024),
1172    ///             BlobRangeRequest::new(row_id, 4096, 1024),
1173    ///         ],
1174    ///     )
1175    ///     .await?;
1176    /// # let _ = ranges;
1177    /// # Ok(())
1178    /// # }
1179    /// ```
1180    ///
1181    /// Returns an error when a range is invalid, a requested row id does not
1182    /// exist, or the column is not a blob v2 column. Returns
1183    /// [`Error::NotSupported`] on table types without blob support.
1184    pub async fn fetch_blob_ranges(
1185        &self,
1186        column: impl AsRef<str>,
1187        requests: impl IntoIterator<Item = BlobRangeRequest>,
1188    ) -> Result<LargeBinaryArray> {
1189        let requests = requests.into_iter().collect::<Vec<_>>();
1190        self.inner
1191            .fetch_blob_ranges(column.as_ref(), &requests)
1192            .await
1193    }
1194
1195    /// Open lazy [`BlobFile`] handles for the given row ids.
1196    ///
1197    /// Same length and order as `row_ids`. Null rows are `None`. Bytes are not
1198    /// read from disk until a call to [`BlobFile::read`].
1199    ///
1200    /// ```
1201    /// # use lancedb::Table;
1202    /// # async fn lazy_read(table: &Table, row_ids: &[u64]) -> Result<(), Box<dyn std::error::Error>> {
1203    /// let handles = table.fetch_blob_files("image", row_ids).await?;
1204    /// if let Some(Some(first)) = handles.first() {
1205    ///     let bytes = first.read().await?;
1206    ///     println!("first blob is {} bytes", bytes.len());
1207    /// }
1208    /// # Ok(())
1209    /// # }
1210    /// ```
1211    pub async fn fetch_blob_files(
1212        &self,
1213        column: impl AsRef<str>,
1214        row_ids: &[u64],
1215    ) -> Result<Vec<Option<BlobFile>>> {
1216        self.inner.fetch_blob_files(column.as_ref(), row_ids).await
1217    }
1218
1219    /// Insert new records into this Table
1220    ///
1221    /// # Arguments
1222    ///
1223    /// * `data` data to be added to the Table
1224    /// * `options` options to control how data is added
1225    pub fn add<T: Scannable + 'static>(&self, data: T) -> AddDataBuilder {
1226        AddDataBuilder::new(
1227            self.inner.clone(),
1228            Box::new(data),
1229            Some(self.embedding_registry.clone()),
1230        )
1231    }
1232
1233    /// Update existing records in the Table
1234    ///
1235    /// An update operation can be used to adjust existing values.  Use the
1236    /// returned builder to specify which columns to update.  The new value
1237    /// can be a literal value (e.g. replacing nulls with some default value)
1238    /// or an expression applied to the old value (e.g. incrementing a value)
1239    ///
1240    /// An optional condition can be specified (e.g. "only update if the old
1241    /// value is 0")
1242    ///
1243    /// Note: if your condition is something like "some_id_column == 7" and
1244    /// you are updating many rows (with different ids) then you will get
1245    /// better performance with a single [`merge_insert`] call instead of
1246    /// repeatedly calilng this method.
1247    pub fn update(&self) -> UpdateBuilder {
1248        UpdateBuilder::new(self.inner.clone())
1249    }
1250
1251    /// Delete the rows from table that match the predicate.
1252    ///
1253    /// # Arguments
1254    /// - `predicate` - A SQL string (`&str`) or DataFusion expression (`&Expr`)
1255    ///   that selects the rows to delete.
1256    ///
1257    /// # Example
1258    ///
1259    /// ```no_run
1260    /// # use std::sync::Arc;
1261    /// # use arrow_array::{FixedSizeListArray, types::Float32Type, RecordBatch,
1262    /// #   RecordBatchIterator, Int32Array};
1263    /// # use arrow_schema::{Schema, Field, DataType};
1264    /// use datafusion_expr::{col, lit};
1265    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
1266    /// let tmpdir = tempfile::tempdir().unwrap();
1267    /// let db = lancedb::connect(tmpdir.path().to_str().unwrap())
1268    ///     .execute()
1269    ///     .await
1270    ///     .unwrap();
1271    /// let schema = Arc::new(Schema::new(vec![
1272    ///     Field::new("id", DataType::Int32, false),
1273    ///     Field::new("vector", DataType::FixedSizeList(
1274    ///         Arc::new(Field::new("item", DataType::Float32, true)), 128), true),
1275    /// ]));
1276    /// let data = RecordBatch::try_new(
1277    ///     schema.clone(),
1278    ///     vec![
1279    ///         Arc::new(Int32Array::from_iter_values(0..10)),
1280    ///         Arc::new(
1281    ///             FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
1282    ///                 (0..10).map(|_| Some(vec![Some(1.0); 128])),
1283    ///                 128,
1284    ///             ),
1285    ///         ),
1286    ///     ],
1287    /// )
1288    /// .unwrap();
1289    /// let tbl = db
1290    ///     .create_table("delete_test", data)
1291    ///     .execute()
1292    ///     .await
1293    ///     .unwrap();
1294    ///
1295    /// // Using a SQL string:
1296    /// tbl.delete("id > 5").await.unwrap();
1297    ///
1298    /// // Using a DataFusion expression:
1299    /// let expr = col("id").lt(lit(4));
1300    /// tbl.delete(&expr).await.unwrap();
1301    /// # });
1302    /// ```
1303    pub async fn delete(&self, predicate: impl Into<Predicate<'_>>) -> Result<DeleteResult> {
1304        self.inner.delete(predicate.into()).await
1305    }
1306
1307    /// Create an index on the provided column(s).
1308    ///
1309    /// Indices are used to speed up searches and are often needed when the size of the table
1310    /// becomes large (the exact size depends on many factors but somewhere between 100K rows
1311    /// and 1M rows is a good rule of thumb)
1312    ///
1313    /// There are a variety of indices available.  They are described more in
1314    /// [`crate::index::Index`].  The simplest thing to do is to use `index::Index::Auto` which
1315    /// will attempt to create the most useful index based on the column type and column
1316    /// statistics. `BTree` index is created by default for numeric, temporal, and
1317    /// string columns.
1318    ///
1319    /// Once an index is created it will remain until the data is overwritten (e.g. an
1320    /// add operation with mode overwrite) or the indexed column is dropped.
1321    ///
1322    /// Indices are not automatically updated with new data.  If you add new data to the
1323    /// table then the index will not include the new rows.  However, a table search will
1324    /// still consider the unindexed rows.  Searches will issue both an indexed search (on
1325    /// the data covered by the index) and a flat search (on the unindexed data) and the
1326    /// results will be combined.
1327    ///
1328    /// If there is enough unindexed data then the flat search will become slow and the index
1329    /// should be optimized.  Optimizing an index will add any unindexed data to the existing
1330    /// index without rerunning the full index creation process.  For more details see
1331    /// [Table::optimize].
1332    ///
1333    /// Note: Multi-column (composite) indices are not currently supported.  However, they will
1334    /// be supported in the future and the API is designed to be compatible with them.
1335    ///
1336    /// # Examples
1337    ///
1338    /// ```no_run
1339    /// # use std::sync::Arc;
1340    /// # use arrow_array::{FixedSizeListArray, types::Float32Type, RecordBatch,
1341    /// #   RecordBatchIterator, Int32Array};
1342    /// # use arrow_schema::{Schema, Field, DataType};
1343    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
1344    /// use lancedb::index::Index;
1345    /// let tmpdir = tempfile::tempdir().unwrap();
1346    /// let db = lancedb::connect(tmpdir.path().to_str().unwrap())
1347    ///     .execute()
1348    ///     .await
1349    ///     .unwrap();
1350    /// # let tbl = db.open_table("idx_test").execute().await.unwrap();
1351    /// // Create IVF PQ index on the "vector" column by default.
1352    /// tbl.create_index(&["vector"], Index::Auto)
1353    ///    .execute()
1354    ///    .await
1355    ///    .unwrap();
1356    /// // Create a BTree index on the "id" column.
1357    /// tbl.create_index(&["id"], Index::Auto)
1358    ///     .execute()
1359    ///     .await
1360    ///     .unwrap();
1361    /// // Create a LabelList index on the "tags" column.
1362    /// tbl.create_index(&["tags"], Index::LabelList(Default::default()))
1363    ///     .execute()
1364    ///     .await
1365    ///     .unwrap();
1366    /// # });
1367    /// ```
1368    pub fn create_index(&self, columns: &[impl AsRef<str>], index: Index) -> IndexBuilder {
1369        IndexBuilder::new(
1370            self.inner.clone(),
1371            columns
1372                .iter()
1373                .map(|val| val.as_ref().to_string())
1374                .collect::<Vec<_>>(),
1375            index,
1376        )
1377    }
1378
1379    /// See [Table::create_index]
1380    /// For remote tables, this allows an optional wait_timeout to poll until asynchronous indexing is complete
1381    pub fn create_index_with_timeout(
1382        &self,
1383        columns: &[impl AsRef<str>],
1384        index: Index,
1385        wait_timeout: Option<std::time::Duration>,
1386    ) -> IndexBuilder {
1387        let mut builder = IndexBuilder::new(
1388            self.inner.clone(),
1389            columns
1390                .iter()
1391                .map(|val| val.as_ref().to_string())
1392                .collect::<Vec<_>>(),
1393            index,
1394        );
1395        if let Some(timeout) = wait_timeout {
1396            builder = builder.wait_timeout(timeout);
1397        }
1398        builder
1399    }
1400
1401    /// Create a builder for a merge insert operation
1402    ///
1403    /// This operation can add rows, update rows, and remove rows all in a single
1404    /// transaction. It is a very generic tool that can be used to create
1405    /// behaviors like "insert if not exists", "update or insert (i.e. upsert)",
1406    /// or even replace a portion of existing data with new data (e.g. replace
1407    /// all data where month="january")
1408    ///
1409    /// The merge insert operation works by combining new data from a
1410    /// **source table** with existing data in a **target table** by using a
1411    /// join.  There are three categories of records.
1412    ///
1413    /// "Matched" records are records that exist in both the source table and
1414    /// the target table. "Not matched" records exist only in the source table
1415    /// (e.g. these are new data) "Not matched by source" records exist only
1416    /// in the target table (this is old data)
1417    ///
1418    /// The builder returned by this method can be used to customize what
1419    /// should happen for each category of data.
1420    ///
1421    /// Please note that the data may appear to be reordered as part of this
1422    /// operation.  This is because updated rows will be deleted from the
1423    /// dataset and then reinserted at the end with the new values.
1424    ///
1425    /// # Arguments
1426    ///
1427    /// * `on` One or more columns to join on.  This is how records from the
1428    ///   source table and target table are matched.  Typically this is some
1429    ///   kind of key or id column.
1430    ///
1431    /// # Examples
1432    ///
1433    /// ```no_run
1434    /// # use std::sync::Arc;
1435    /// # use arrow_array::{FixedSizeListArray, types::Float32Type, RecordBatch,
1436    /// #   RecordBatchIterator, Int32Array};
1437    /// # use arrow_schema::{Schema, Field, DataType};
1438    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
1439    /// let tmpdir = tempfile::tempdir().unwrap();
1440    /// let db = lancedb::connect(tmpdir.path().to_str().unwrap())
1441    ///     .execute()
1442    ///     .await
1443    ///     .unwrap();
1444    /// # let tbl = db.open_table("idx_test").execute().await.unwrap();
1445    /// # let schema = Arc::new(Schema::new(vec![
1446    /// #  Field::new("id", DataType::Int32, false),
1447    /// #  Field::new("vector", DataType::FixedSizeList(
1448    /// #    Arc::new(Field::new("item", DataType::Float32, true)), 128), true),
1449    /// # ]));
1450    /// let new_data = RecordBatchIterator::new(
1451    ///     vec![RecordBatch::try_new(
1452    ///         schema.clone(),
1453    ///         vec![
1454    ///             Arc::new(Int32Array::from_iter_values(0..10)),
1455    ///             Arc::new(
1456    ///                 FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
1457    ///                     (0..10).map(|_| Some(vec![Some(1.0); 128])),
1458    ///                     128,
1459    ///                 ),
1460    ///             ),
1461    ///         ],
1462    ///     )
1463    ///     .unwrap()]
1464    ///     .into_iter()
1465    ///     .map(Ok),
1466    ///     schema.clone(),
1467    /// );
1468    /// // Perform an upsert operation
1469    /// let mut merge_insert = tbl.merge_insert(&["id"]);
1470    /// merge_insert
1471    ///     .when_matched_update_all(None)
1472    ///     .when_not_matched_insert_all();
1473    /// merge_insert.execute(Box::new(new_data)).await.unwrap();
1474    /// # });
1475    /// ```
1476    pub fn merge_insert(&self, on: &[&str]) -> MergeInsertBuilder {
1477        MergeInsertBuilder::new(
1478            self.inner.clone(),
1479            on.iter().map(|s| s.to_string()).collect(),
1480        )
1481    }
1482
1483    /// Create a [`Query`] Builder.
1484    ///
1485    /// Queries allow you to search your existing data.  By default the query will
1486    /// return all the data in the table in no particular order.  The builder
1487    /// returned by this method can be used to control the query using filtering,
1488    /// vector similarity, sorting, and more.
1489    ///
1490    /// Note: By default, all columns are returned.  For best performance, you should
1491    /// only fetch the columns you need.  See [`Query::select_with_projection`] for
1492    /// more details.
1493    ///
1494    /// When appropriate, various indices and statistics will be used to accelerate
1495    /// the query.
1496    ///
1497    /// # Examples
1498    ///
1499    /// ## Vector search
1500    ///
1501    /// This example will find the 10 rows whose value in the "vector" column are
1502    /// closest to the query vector [1.0, 2.0, 3.0].  If an index has been created
1503    /// on the "vector" column then this will perform an ANN search.
1504    ///
1505    /// The [`Query::refine_factor`] and [`Query::nprobes`] methods are used to
1506    /// control the recall / latency tradeoff of the search.
1507    ///
1508    /// ```no_run
1509    /// # use arrow_array::RecordBatch;
1510    /// # use futures::TryStreamExt;
1511    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
1512    /// # let conn = lancedb::connect("/tmp").execute().await.unwrap();
1513    /// # let tbl = conn.open_table("tbl").execute().await.unwrap();
1514    /// use crate::lancedb::Table;
1515    /// use crate::lancedb::query::ExecutableQuery;
1516    /// let stream = tbl
1517    ///     .query()
1518    ///     .nearest_to(&[1.0, 2.0, 3.0])
1519    ///     .unwrap()
1520    ///     .refine_factor(5)
1521    ///     .nprobes(10)
1522    ///     .execute()
1523    ///     .await
1524    ///     .unwrap();
1525    /// let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
1526    /// # });
1527    /// ```
1528    ///
1529    /// ## SQL-style filter
1530    ///
1531    /// This query will return up to 1000 rows whose value in the `id` column
1532    /// is greater than 5.  LanceDb supports a broad set of filtering functions.
1533    ///
1534    /// ```no_run
1535    /// # use arrow_array::RecordBatch;
1536    /// # use futures::TryStreamExt;
1537    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
1538    /// # let conn = lancedb::connect("/tmp").execute().await.unwrap();
1539    /// # let tbl = conn.open_table("tbl").execute().await.unwrap();
1540    /// use crate::lancedb::Table;
1541    /// use crate::lancedb::query::{ExecutableQuery, QueryBase};
1542    /// let stream = tbl
1543    ///     .query()
1544    ///     .only_if("id > 5")
1545    ///     .limit(1000)
1546    ///     .execute()
1547    ///     .await
1548    ///     .unwrap();
1549    /// let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
1550    /// # });
1551    /// ```
1552    ///
1553    /// ## Full scan
1554    ///
1555    /// This query will return everything in the table in no particular
1556    /// order.
1557    ///
1558    /// ```no_run
1559    /// # use arrow_array::RecordBatch;
1560    /// # use futures::TryStreamExt;
1561    /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
1562    /// # let conn = lancedb::connect("/tmp").execute().await.unwrap();
1563    /// # let tbl = conn.open_table("tbl").execute().await.unwrap();
1564    /// use crate::lancedb::Table;
1565    /// use crate::lancedb::query::ExecutableQuery;
1566    /// let stream = tbl.query().execute().await.unwrap();
1567    /// let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
1568    /// # });
1569    /// ```
1570    pub fn query(&self) -> Query {
1571        Query::new(self.inner.clone())
1572    }
1573
1574    /// Extract rows from the dataset using dataset offsets.
1575    ///
1576    /// Dataset offsets are 0-indexed and relative to the current version of the table.
1577    /// They are not stable.  A row with an offset of N may have a different offset in a
1578    /// different version of the table (e.g. if an earlier row is deleted).
1579    ///
1580    /// Offsets are useful for sampling as the set of all valid offsets is easily
1581    /// known in advance to be [0, len(table)).
1582    ///
1583    /// No guarantees are made regarding the order in which results are returned.  If you
1584    /// desire an output order that matches the order of the given offsets, you will need
1585    /// to add the row offset column to the output and align it yourself.
1586    ///
1587    /// Parameters
1588    /// ----------
1589    /// offsets: list[int]
1590    ///     The offsets to take.
1591    ///
1592    /// Returns
1593    /// -------
1594    /// pa.RecordBatch
1595    ///     A record batch containing the rows at the given offsets.
1596    pub fn take_offsets(&self, offsets: Vec<u64>) -> TakeQuery {
1597        TakeQuery::from_offsets(self.inner.clone(), offsets)
1598    }
1599
1600    /// Extract rows from the dataset using row ids.
1601    ///
1602    /// Row ids are not stable and are relative to the current version of the table.
1603    /// They can change due to compaction and updates.
1604    ///
1605    /// Even so, row ids are more stable than offsets and can be useful in some situations.
1606    ///
1607    /// There is an ongoing effort to make row ids stable which is tracked at
1608    /// https://github.com/lancedb/lancedb/issues/1120
1609    ///
1610    /// No guarantees are made regarding the order in which results are returned.  If you
1611    /// desire an output order that matches the order of the given ids, you will need
1612    /// to add the row id column to the output and align it yourself.
1613    /// Parameters
1614    /// ----------
1615    /// row_ids: list[int]
1616    ///     The row ids to take.
1617    ///
1618    pub fn take_row_ids(&self, row_ids: Vec<u64>) -> TakeQuery {
1619        TakeQuery::from_row_ids(self.inner.clone(), row_ids)
1620    }
1621
1622    /// Search the table with a given query vector.
1623    ///
1624    /// This is a convenience method for preparing a vector query and
1625    /// is the same thing as calling `nearest_to` on the builder returned
1626    /// by `query`.  See [`Query::nearest_to`] for more details.
1627    pub fn vector_search(&self, query: impl IntoQueryVector) -> Result<VectorQuery> {
1628        self.query().nearest_to(query)
1629    }
1630
1631    /// Optimize the on-disk data and indices for better performance.
1632    ///
1633    /// Modeled after ``VACUUM`` in PostgreSQL.
1634    ///
1635    /// Optimization is discussed in more detail in the [OptimizeAction] documentation
1636    /// and covers three operations:
1637    ///
1638    ///  * Compaction: Merges small files into larger ones
1639    ///  * Prune: Removes old versions of the dataset
1640    ///  * Index: Optimizes the indices, adding new data to existing indices
1641    ///
1642    /// The frequency an application should call optimize is based on the frequency of
1643    /// data modifications.  If data is frequently added, deleted, or updated then
1644    /// optimize should be run frequently.  A good rule of thumb is to run optimize if
1645    /// you have added or modified 100,000 or more records or run more than 20 data
1646    /// modification operations.
1647    pub async fn optimize(&self, action: OptimizeAction) -> Result<OptimizeStats> {
1648        self.inner.optimize(action).await
1649    }
1650
1651    /// Add new columns to the table, providing values to fill in.
1652    pub fn add_columns(&self) -> AddColumnsBuilder {
1653        AddColumnsBuilder::new(self.inner.clone())
1654    }
1655
1656    /// Change a column's name or nullability.
1657    pub async fn alter_columns(
1658        &self,
1659        alterations: &[ColumnAlteration],
1660    ) -> Result<AlterColumnsResult> {
1661        self.inner.alter_columns(alterations).await
1662    }
1663
1664    /// Update per-field metadata (merges by default).
1665    pub async fn update_field_metadata(
1666        &self,
1667        updates: &[FieldMetadataUpdate],
1668    ) -> Result<UpdateFieldMetadataResult> {
1669        self.inner.update_field_metadata(updates).await
1670    }
1671
1672    /// Remove columns from the table.
1673    pub async fn drop_columns(&self, columns: &[&str]) -> Result<DropColumnsResult> {
1674        self.inner.drop_columns(columns).await
1675    }
1676
1677    /// Set the unenforced primary key for this table to a single column.
1678    ///
1679    /// "Unenforced" means LanceDB does not check uniqueness on writes; the
1680    /// column is recorded in the schema as the primary key so that features
1681    /// such as `merge_insert` can use it.
1682    ///
1683    /// Only single-column primary keys are supported, and the key cannot be
1684    /// changed once set — calling this on a table that already has an
1685    /// unenforced primary key fails. `columns` is an iterable for binding
1686    /// ergonomics but must yield exactly one column:
1687    ///
1688    /// - `table.set_unenforced_primary_key(["id"])`
1689    pub async fn set_unenforced_primary_key<I, S>(&self, columns: I) -> Result<()>
1690    where
1691        I: IntoIterator<Item = S>,
1692        S: Into<String>,
1693    {
1694        let owned: Vec<String> = columns.into_iter().map(Into::into).collect();
1695        let borrowed: Vec<&str> = owned.iter().map(String::as_str).collect();
1696        self.inner.set_unenforced_primary_key(&borrowed).await
1697    }
1698
1699    /// Install an [`LsmWriteSpec`] on this table, selecting Lance's MemWAL
1700    /// LSM-style write path for future `merge_insert` calls.
1701    ///
1702    /// [`LsmWriteSpec`] chooses one of three sharding strategies:
1703    ///
1704    /// - [`LsmWriteSpec::bucket`] — hash-bucket writes by a scalar column.
1705    /// - [`LsmWriteSpec::identity`] — shard by the raw value of a scalar column.
1706    /// - [`LsmWriteSpec::unsharded`] — route every write to a single shard.
1707    ///
1708    /// # Example
1709    ///
1710    /// ```
1711    /// # use lancedb::table::{LsmWriteSpec, Table};
1712    /// # async fn example(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
1713    /// table
1714    ///     .set_lsm_write_spec(
1715    ///         LsmWriteSpec::bucket("id", 16).with_maintained_indexes(["id_idx"]),
1716    ///     )
1717    ///     .await?;
1718    /// # Ok(())
1719    /// # }
1720    /// ```
1721    pub async fn set_lsm_write_spec(&self, spec: LsmWriteSpec) -> Result<()> {
1722        self.inner.set_lsm_write_spec(spec).await
1723    }
1724
1725    /// Remove the [`LsmWriteSpec`] from this table, reverting to the standard
1726    /// `merge_insert` write path.
1727    ///
1728    /// Errors if no spec is currently set.
1729    pub async fn unset_lsm_write_spec(&self) -> Result<()> {
1730        self.inner.unset_lsm_write_spec().await
1731    }
1732
1733    /// Read the [`LsmWriteSpec`] currently installed on this table.
1734    ///
1735    /// Returns `Ok(None)` when the MemWAL LSM write path is not enabled (no
1736    /// spec has been set, or it was removed with [`Table::unset_lsm_write_spec`]).
1737    /// The returned spec — including its [`LsmWriteSpec::maintained_indexes`] and
1738    /// [`LsmWriteSpec::writer_config_defaults`] — mirrors what was passed to
1739    /// [`Table::set_lsm_write_spec`].
1740    ///
1741    /// # Example
1742    ///
1743    /// ```
1744    /// # use lancedb::table::Table;
1745    /// # async fn example(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
1746    /// if let Some(spec) = table.get_lsm_write_spec().await? {
1747    ///     println!("LSM write path enabled: {:?}", spec);
1748    /// }
1749    /// # Ok(())
1750    /// # }
1751    /// ```
1752    pub async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
1753        self.inner.get_lsm_write_spec().await
1754    }
1755
1756    /// Converge this table's LSM write path into its base table.
1757    ///
1758    /// One `flush` to seal every memtable into L0, then compaction triggers
1759    /// until every generation that existed at that moment has reached base.
1760    /// The loop runs client-side, reading progress from `get_lsm_stats`, so
1761    /// there is no held socket and nothing to reconcile if you drop this
1762    /// future partway through.
1763    ///
1764    /// **Best-effort.** Generations created *after* the opening flush are
1765    /// deliberately not waited on — that is what lets this terminate on a
1766    /// table taking writes. Idempotent and safe on a cadence: an
1767    /// already-converged table costs two round trips and triggers nothing.
1768    ///
1769    /// **No deadline, and the caller owns that.** It returns when the target
1770    /// generations are gone, propagates a terminal server fault, and
1771    /// otherwise waits however long the server takes. A slow table and a
1772    /// stuck one are the same picture from here: the compactor pool is shared
1773    /// across every table on the node, so a checkpoint queued behind
1774    /// unrelated work is indistinguishable from one that is merging. Wrap
1775    /// this in `tokio::time::timeout` for a wall-clock bound; abandoning it
1776    /// partway costs nothing.
1777    ///
1778    /// # Example
1779    ///
1780    /// ```no_run
1781    /// # use lancedb::Table;
1782    /// # async fn example(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
1783    /// let before = table.get_lsm_stats(false).await?;
1784    /// table.checkpoint_lsm().await?;
1785    /// let after = table.get_lsm_stats(false).await?;
1786    /// # Ok(())
1787    /// # }
1788    /// ```
1789    pub async fn checkpoint_lsm(&self) -> Result<()> {
1790        checkpoint::checkpoint_lsm(self).await
1791    }
1792
1793    /// Seal every bucket's active memtable into L0 without touching the
1794    /// base table.
1795    ///
1796    /// Independently useful: flushing makes memtable rows readable from L0 at
1797    /// a lower per-query cost. On a node that has not claimed this table it
1798    /// claims it and replays the WAL log first — reporting "nothing to flush"
1799    /// without replaying would lie about durable data.
1800    pub async fn flush_lsm(&self) -> Result<()> {
1801        self.inner.flush_lsm().await
1802    }
1803
1804    /// Run one bounded L0 → base compaction pass per bucket, reporting what
1805    /// it merged and what is left.
1806    ///
1807    /// One pass, not convergence: that bounds each request's cost and gives a
1808    /// caller driving its own cadence a progress signal per round trip.
1809    pub async fn compact_lsm(&self) -> Result<()> {
1810        self.inner.compact_lsm().await
1811    }
1812
1813    /// Read live per-bucket LSM state.
1814    ///
1815    /// Answers "how far behind is my fresh tier", "which bucket is hot", and
1816    /// "why is my fresh-tier vector search brute-force". Mutates no table
1817    /// state, though on a node that has not claimed this table it claims it,
1818    /// exactly as a read would.
1819    ///
1820    /// `include_generation_rows` reports a row count per L0 generation. Off by
1821    /// default: each count opens an uncached Lance dataset, and
1822    /// `checkpoint_lsm` polls this needing only generation numbers.
1823    ///
1824    /// `Ok(None)` only when the LSM write path is not enabled, matching
1825    /// [`Table::get_lsm_write_spec`]. Stats is fresh-tier only, so with the
1826    /// WAL off there is no manifest to report and a struct of zeros would
1827    /// read as measurements.
1828    ///
1829    /// Do not build a checkpoint's termination on this: the completion
1830    /// predicate lives in the `flush` and `compact` responses.
1831    pub async fn get_lsm_stats(&self, include_generation_rows: bool) -> Result<Option<LsmStats>> {
1832        self.inner.get_lsm_stats(include_generation_rows).await
1833    }
1834
1835    /// Drain and close any cached MemWAL shard writers held for this table.
1836    ///
1837    /// When an [`LsmWriteSpec`] is installed, `merge_insert` opens MemWAL shard
1838    /// writers and caches them for reuse across calls. This closes them,
1839    /// flushing pending data; writers reopen lazily on the next `merge_insert`.
1840    /// It is a no-op when no writers are cached.
1841    pub async fn close_lsm_writers(&self) -> Result<()> {
1842        self.inner.close_lsm_writers().await
1843    }
1844
1845    /// Retrieve the version of the table
1846    ///
1847    /// LanceDb supports versioning.  Every operation that modifies the table increases
1848    /// version.  As long as a version hasn't been deleted you can `[Self::checkout]` that
1849    /// version to view the data at that point.  In addition, you can `[Self::restore]` the
1850    /// version to replace the current table with a previous version.
1851    pub async fn version(&self) -> Result<u64> {
1852        self.inner.version().await
1853    }
1854
1855    /// Checks out a specific version of the Table
1856    ///
1857    /// Any read operation on the table will now access the data at the checked out version.
1858    /// As a consequence, calling this method will disable any read consistency interval
1859    /// that was previously set.
1860    ///
1861    /// This is a read-only operation that turns the table into a sort of "view"
1862    /// or "detached head".  Other table instances will not be affected.  To make the change
1863    /// permanent you can use the `[Self::restore]` method.
1864    ///
1865    /// Any operation that modifies the table will fail while the table is in a checked
1866    /// out state.
1867    ///
1868    /// To return the table to a normal state use `[Self::checkout_latest]`
1869    pub async fn checkout(&self, version: u64) -> Result<()> {
1870        self.inner.checkout(version).await
1871    }
1872
1873    /// Checks out a specific version of the Table by tag
1874    ///
1875    /// Any read operation on the table will now access the data at the version referenced by the tag.
1876    /// As a consequence, calling this method will disable any read consistency interval
1877    /// that was previously set.
1878    ///
1879    /// This is a read-only operation that turns the table into a sort of "view"
1880    /// or "detached head".  Other table instances will not be affected.  To make the change
1881    /// permanent you can use the `[Self::restore]` method.
1882    ///
1883    /// Any operation that modifies the table will fail while the table is in a checked
1884    /// out state.
1885    ///
1886    /// To return the table to a normal state use `[Self::checkout_latest]`
1887    pub async fn checkout_tag(&self, tag: &str) -> Result<()> {
1888        self.inner.checkout_tag(tag).await
1889    }
1890
1891    /// Ensures the table is pointing at the latest version
1892    ///
1893    /// This can be used to manually update a table when the read_consistency_interval is None
1894    /// It can also be used to undo a `[Self::checkout]` operation
1895    pub async fn checkout_latest(&self) -> Result<()> {
1896        self.inner.checkout_latest().await
1897    }
1898
1899    /// Restore the table to the currently checked out version
1900    ///
1901    /// This operation will fail if checkout has not been called previously
1902    ///
1903    /// This operation will overwrite the latest version of the table with a
1904    /// previous version.  Any changes made since the checked out version will
1905    /// no longer be visible.
1906    ///
1907    /// Once the operation concludes the table will no longer be in a checked
1908    /// out state and the read_consistency_interval, if any, will apply.
1909    pub async fn restore(&self) -> Result<()> {
1910        self.inner.restore().await
1911    }
1912
1913    /// List all the versions of the table
1914    pub async fn list_versions(&self) -> Result<Vec<Version>> {
1915        self.inner.list_versions().await
1916    }
1917
1918    /// List all indices that have been created with [`Self::create_index`]
1919    pub async fn list_indices(&self) -> Result<Vec<IndexConfig>> {
1920        self.inner.list_indices().await
1921    }
1922
1923    /// Tokenize a full-text search query using the tokenizer configured on an FTS index.
1924    ///
1925    /// Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in
1926    /// the client process from index metadata. For remote tables, this means the
1927    /// same tokenizer model files must also exist locally.
1928    pub async fn tokenize(&self, query: &str, index_name: &str) -> Result<Vec<FtsToken>> {
1929        let indices = self.inner.list_indices().await?;
1930        let matches = indices
1931            .iter()
1932            .filter(|idx| idx.name == index_name)
1933            .collect::<Vec<_>>();
1934        let index = match matches.as_slice() {
1935            [index] => *index,
1936            [] => {
1937                return Err(Error::InvalidInput {
1938                    message: format!("No index named '{}'", index_name),
1939                });
1940            }
1941            _ => {
1942                return Err(Error::InvalidInput {
1943                    message: format!("Index name '{}' is ambiguous", index_name),
1944                });
1945            }
1946        };
1947        if index.index_type != IndexType::FTS {
1948            return Err(Error::InvalidInput {
1949                message: format!("Index '{}' is not a full text search index", index_name),
1950            });
1951        }
1952        self.tokenize_with_index(query, index, index_name)
1953    }
1954
1955    /// Tokenize a full-text search query using the tokenizer configured on the
1956    /// FTS index for a column.
1957    ///
1958    /// The column must have exactly one FTS index. Model-backed tokenizers such
1959    /// as `jieba/*` and `lindera/*` are rebuilt in the client process from
1960    /// index metadata. For remote tables, this means the same tokenizer model
1961    /// files must also exist locally.
1962    pub async fn tokenize_with_column(&self, query: &str, column: &str) -> Result<Vec<FtsToken>> {
1963        let schema = self.inner.schema().await?;
1964        let (column, _) = resolve_arrow_field_path(schema.as_ref(), column)?;
1965        let indices = self.inner.list_indices().await?;
1966        let matches = indices
1967            .iter()
1968            .filter(|idx| {
1969                idx.index_type == IndexType::FTS
1970                    && idx.columns.len() == 1
1971                    && idx.columns[0] == column
1972            })
1973            .collect::<Vec<_>>();
1974        let index = match matches.as_slice() {
1975            [index] => *index,
1976            [] => {
1977                return Err(Error::InvalidInput {
1978                    message: format!("Column '{}' does not have a full text search index", column),
1979                });
1980            }
1981            _ => {
1982                return Err(Error::InvalidInput {
1983                    message: format!(
1984                        "Column '{}' has multiple full text search indexes; tokenization by column is ambiguous",
1985                        column
1986                    ),
1987                });
1988            }
1989        };
1990        self.tokenize(query, &index.name).await
1991    }
1992
1993    fn tokenize_with_index(
1994        &self,
1995        query: &str,
1996        index: &IndexConfig,
1997        index_name: &str,
1998    ) -> Result<Vec<FtsToken>> {
1999        let selector_description = format!("index name '{}'", index_name);
2000        let details = index
2001            .index_details
2002            .as_deref()
2003            .ok_or_else(|| Error::InvalidInput {
2004                message: format!(
2005                    "Full text search index '{}' for {} does not include tokenizer details",
2006                    index.name, selector_description
2007                ),
2008            })?;
2009        let params = serde_json::from_str::<InvertedIndexParams>(details).map_err(|err| {
2010            Error::InvalidInput {
2011                message: format!(
2012                    "Failed to parse tokenizer details for full text search index '{}' for {}: {}",
2013                    index.name, selector_description, err
2014                ),
2015            }
2016        })?;
2017        tokenize(query, &params).map_err(|err| match err {
2018            Error::InvalidInput { message } => Error::InvalidInput {
2019                message: format!(
2020                    "{} for full text search index '{}' for {}",
2021                    message, index.name, selector_description
2022                ),
2023            },
2024            err => err,
2025        })
2026    }
2027
2028    /// Get the table URI (storage location)
2029    ///
2030    /// Returns the full storage location of the table (e.g., S3/GCS path).
2031    /// For remote tables, this fetches the location from the server via describe.
2032    pub async fn uri(&self) -> Result<String> {
2033        self.inner.uri().await
2034    }
2035
2036    /// Get the storage options used when opening this table, if any.
2037    ///
2038    /// Warning: This is an internal API and the return value is subject to change.
2039    #[deprecated(since = "0.25.0", note = "Use initial_storage_options() instead")]
2040    pub async fn storage_options(&self) -> Option<HashMap<String, String>> {
2041        #[allow(deprecated)]
2042        self.inner.storage_options().await
2043    }
2044
2045    /// Get the initial storage options that were passed in when opening this table.
2046    ///
2047    /// For dynamically refreshed options (e.g., credential vending), use [`Self::latest_storage_options`].
2048    ///
2049    /// Warning: This is an internal API and the return value is subject to change.
2050    pub async fn initial_storage_options(&self) -> Option<HashMap<String, String>> {
2051        self.inner.initial_storage_options().await
2052    }
2053
2054    /// Get the latest storage options, refreshing from provider if configured.
2055    ///
2056    /// This method is useful for credential vending scenarios where storage options
2057    /// may be refreshed dynamically. If no dynamic provider is configured, this
2058    /// returns the initial static options.
2059    ///
2060    /// Warning: This is an internal API and the return value is subject to change.
2061    pub async fn latest_storage_options(&self) -> Result<Option<HashMap<String, String>>> {
2062        self.inner.latest_storage_options().await
2063    }
2064
2065    /// Get statistics about an index.
2066    /// Returns None if the index does not exist.
2067    pub async fn index_stats(
2068        &self,
2069        index_name: impl AsRef<str>,
2070    ) -> Result<Option<IndexStatistics>> {
2071        self.inner.index_stats(index_name.as_ref()).await
2072    }
2073
2074    /// Drop an index from the table.
2075    ///
2076    /// Note: This is not yet available in LanceDB cloud.
2077    ///
2078    /// This does not delete the index from disk, it just removes it from the table.
2079    /// To delete the index, run [`Self::optimize()`] after dropping the index.
2080    ///
2081    /// Use [`Self::list_indices()`] to find the names of the indices.
2082    pub async fn drop_index(&self, name: &str) -> Result<()> {
2083        self.inner.drop_index(name).await
2084    }
2085
2086    /// Prewarm an index in the table.
2087    ///
2088    /// This is a hint to the database that the index will be accessed in the
2089    /// future and should be loaded into memory if possible.  This can reduce
2090    /// cold-start latency for subsequent queries.
2091    ///
2092    /// This call initiates prewarming and returns once the request is accepted.
2093    /// It is idempotent and safe to call from multiple clients concurrently.
2094    ///
2095    /// It is generally wasteful to call this if the index does not fit into the
2096    /// available cache.  Not all index types support prewarming; unsupported
2097    /// indices will silently ignore the request.
2098    ///
2099    /// Use [`Self::list_indices()`] to find the names of the indices.
2100    pub async fn prewarm_index(&self, name: &str) -> Result<()> {
2101        self.inner.prewarm_index(name).await
2102    }
2103
2104    /// Prewarm data for the table.
2105    ///
2106    /// This is a hint to the database that the given columns will be accessed in
2107    /// the future and the database should prefetch the data if possible.  This
2108    /// can reduce cold-start latency for subsequent queries.  Currently only
2109    /// supported on remote tables.
2110    ///
2111    /// This call initiates prewarming and returns once the request is accepted.
2112    /// It is idempotent and safe to call from multiple clients concurrently —
2113    /// calling it on already-prewarmed columns is a no-op on the server.
2114    ///
2115    /// This operation has a large upfront cost but can speed up future queries
2116    /// that need to fetch the given columns.  Large columns such as embeddings
2117    /// or binary data may not be practical to prewarm.  This feature is intended
2118    /// for workloads that issue many queries against the same columns.
2119    ///
2120    /// If `columns` is `None`, all columns are prewarmed.
2121    pub async fn prewarm_data(&self, columns: Option<Vec<String>>) -> Result<()> {
2122        self.inner.prewarm_data(columns).await
2123    }
2124
2125    /// Poll until the columns are fully indexed. Will return Error::Timeout if the columns
2126    /// are not fully indexed within the timeout.
2127    pub async fn wait_for_index(
2128        &self,
2129        index_names: &[&str],
2130        timeout: std::time::Duration,
2131    ) -> Result<()> {
2132        self.inner.wait_for_index(index_names, timeout).await
2133    }
2134
2135    /// Get the tags manager.
2136    pub async fn tags(&self) -> Result<Box<dyn Tags + '_>> {
2137        self.inner.tags().await
2138    }
2139
2140    /// Create a new branch from `from` (a version, tag, or branch)
2141    pub async fn create_branch(
2142        &self,
2143        name: &str,
2144        from: impl Into<lance::dataset::refs::Ref>,
2145    ) -> Result<Self> {
2146        let inner = self.inner.create_branch(name, from.into()).await?;
2147        Ok(Self {
2148            inner,
2149            database: self.database.clone(),
2150            embedding_registry: self.embedding_registry.clone(),
2151        })
2152    }
2153
2154    /// Check out an existing branch and return a handle scoped to it.
2155    ///
2156    /// With `version` set, the returned handle is pinned to that version of the
2157    /// branch: a read-only, detached view (as with [`Self::checkout`]). With
2158    /// `version` as `None` it tracks the branch's latest and stays writable.
2159    ///
2160    /// ```
2161    /// # use lancedb::Table;
2162    /// # async fn f(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
2163    /// let exp_at_v3 = table.checkout_branch("exp", Some(3)).await?;
2164    /// # Ok(())
2165    /// # }
2166    /// ```
2167    pub async fn checkout_branch(&self, name: &str, version: Option<u64>) -> Result<Self> {
2168        let inner = self.inner.checkout_branch_version(name, version).await?;
2169        Ok(Self {
2170            inner,
2171            database: self.database.clone(),
2172            embedding_registry: self.embedding_registry.clone(),
2173        })
2174    }
2175
2176    /// List the branches of the table.
2177    pub async fn list_branches(&self) -> Result<HashMap<String, BranchContents>> {
2178        self.inner.list_branches().await
2179    }
2180
2181    /// Delete a branch.
2182    pub async fn delete_branch(&self, name: &str) -> Result<()> {
2183        self.inner.delete_branch(name).await
2184    }
2185
2186    /// Diff a branch against main. Remote only.
2187    pub async fn diff_branch(&self, from_branch: &str) -> Result<BranchDiff> {
2188        self.inner.diff_branch(from_branch).await
2189    }
2190
2191    /// Merge a branch into main, or dry-run. Remote only.
2192    /// HTTP 409 still returns [`Ok`] with [`MergeBranchStatus::Rejected`].
2193    pub async fn merge_branch(
2194        &self,
2195        from_branch: &str,
2196        dry_run: bool,
2197    ) -> Result<MergeBranchResult> {
2198        self.inner.merge_branch(from_branch, dry_run).await
2199    }
2200
2201    /// The branch this handle is scoped to, or `None` for `main`.
2202    pub fn current_branch(&self) -> Option<String> {
2203        self.inner.current_branch()
2204    }
2205
2206    /// Retrieve statistics on the table
2207    pub async fn stats(&self) -> Result<TableStatistics> {
2208        self.inner.stats().await
2209    }
2210}
2211
2212pub struct NativeTags {
2213    dataset: dataset::DatasetConsistencyWrapper,
2214}
2215#[async_trait]
2216impl Tags for NativeTags {
2217    async fn list(&self) -> Result<HashMap<String, TagContents>> {
2218        let dataset = self.dataset.get().await?;
2219        Ok(dataset.tags().list().await?)
2220    }
2221
2222    async fn get_version(&self, tag: &str) -> Result<u64> {
2223        let dataset = self.dataset.get().await?;
2224        Ok(dataset.tags().get_version(tag).await?)
2225    }
2226
2227    async fn create(&mut self, tag: &str, version: u64) -> Result<()> {
2228        let dataset = self.dataset.get().await?;
2229        dataset.tags().create(tag, version).await?;
2230        Ok(())
2231    }
2232
2233    async fn delete(&mut self, tag: &str) -> Result<()> {
2234        let dataset = self.dataset.get().await?;
2235        dataset.tags().delete(tag).await?;
2236        Ok(())
2237    }
2238
2239    async fn update(&mut self, tag: &str, version: u64) -> Result<()> {
2240        let dataset = self.dataset.get().await?;
2241        dataset.tags().update(tag, version).await?;
2242        Ok(())
2243    }
2244}
2245
2246pub trait NativeTableExt {
2247    /// Cast as [`NativeTable`], or return None it if is not a [`NativeTable`].
2248    fn as_native(&self) -> Option<&NativeTable>;
2249}
2250
2251impl NativeTableExt for Arc<dyn BaseTable> {
2252    fn as_native(&self) -> Option<&NativeTable> {
2253        self.as_any().downcast_ref::<NativeTable>()
2254    }
2255}
2256
2257/// A table in a LanceDB database.
2258#[derive(Clone)]
2259pub struct NativeTable {
2260    name: String,
2261    namespace: Vec<String>,
2262    id: String,
2263    uri: String,
2264    pub(crate) dataset: dataset::DatasetConsistencyWrapper,
2265    // This comes from the connection options. We store here so we can pass down
2266    // to the dataset when we recreate it (for example, in checkout_latest).
2267    read_consistency_interval: Option<std::time::Duration>,
2268    // Optional namespace client for namespace operations (e.g., managed versioning).
2269    // pub(crate) so query.rs can access the field for server-side query execution.
2270    pub(crate) namespace_client: Option<Arc<dyn LanceNamespace>>,
2271    // Operations to push down to the namespace server.
2272    // pub(crate) so query.rs can access the field for server-side query execution.
2273    pub(crate) pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2274    // Read-freshness baseline; `Some` only for namespace-backed tables.
2275    freshness: Option<TableFreshness>,
2276}
2277
2278impl std::fmt::Debug for NativeTable {
2279    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2280        f.debug_struct("NativeTable")
2281            .field("name", &self.name)
2282            .field("namespace", &self.namespace)
2283            .field("id", &self.id)
2284            .field("uri", &self.uri)
2285            .field("read_consistency_interval", &self.read_consistency_interval)
2286            .field("namespace_client", &self.namespace_client)
2287            .field("pushdown_operations", &self.pushdown_operations)
2288            .finish()
2289    }
2290}
2291
2292impl std::fmt::Display for NativeTable {
2293    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2294        write!(
2295            f,
2296            "NativeTable({}, uri={}, read_consistency_interval={})",
2297            self.name,
2298            self.uri,
2299            match self.read_consistency_interval {
2300                None => {
2301                    "None".to_string()
2302                }
2303                Some(duration) => {
2304                    format!("{}s", duration.as_secs_f64())
2305                }
2306            }
2307        )
2308    }
2309}
2310
2311impl NativeTable {
2312    /// Opens an existing Table
2313    ///
2314    /// # Arguments
2315    ///
2316    /// * `uri` - The uri to a [NativeTable]
2317    /// * `name` - The table name
2318    ///
2319    /// # Returns
2320    ///
2321    /// * A [NativeTable] object.
2322    pub async fn open(uri: &str) -> Result<Self> {
2323        let name = Self::get_table_name(uri)?;
2324        Self::open_with_params(
2325            uri,
2326            &name,
2327            vec![],
2328            None,
2329            None,
2330            None,
2331            None,
2332            HashSet::new(),
2333            None,
2334        )
2335        .await
2336    }
2337
2338    /// Opens an existing Table
2339    ///
2340    /// # Arguments
2341    ///
2342    /// * `base_path` - The base path where the table is located
2343    /// * `name` The Table name
2344    /// * `params` The [ReadParams] to use when opening the table
2345    /// * `namespace_client` - Optional namespace client for namespace operations
2346    /// * `pushdown_operations` - Operations to push down to the namespace server
2347    /// * `managed_versioning` - Whether managed versioning is enabled. If None and namespace_client
2348    ///   is provided, the value will be fetched via describe_table.
2349    ///
2350    /// # Returns
2351    ///
2352    /// * A [NativeTable] object.
2353    #[allow(clippy::too_many_arguments)]
2354    pub async fn open_with_params(
2355        uri: &str,
2356        name: &str,
2357        namespace: Vec<String>,
2358        write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2359        params: Option<ReadParams>,
2360        read_consistency_interval: Option<std::time::Duration>,
2361        namespace_client: Option<Arc<dyn LanceNamespace>>,
2362        pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2363        managed_versioning: Option<bool>,
2364    ) -> Result<Self> {
2365        let params = params.unwrap_or_default();
2366        // patch the params if we have a write store wrapper
2367        let params = match write_store_wrapper.clone() {
2368            Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
2369            None => params,
2370        };
2371
2372        // Build table_id from namespace + name
2373        let mut table_id = namespace.clone();
2374        table_id.push(name.to_string());
2375
2376        // Determine if managed_versioning is enabled
2377        // Use the provided value if available, otherwise query the namespace
2378        let managed_versioning = match managed_versioning {
2379            Some(value) => value,
2380            None if namespace_client.is_some() => {
2381                let ns_client = namespace_client.as_ref().unwrap();
2382                let describe_request = DescribeTableRequest {
2383                    id: Some(table_id.clone()),
2384                    ..Default::default()
2385                };
2386                let response = ns_client
2387                    .describe_table(describe_request)
2388                    .await
2389                    .map_err(|e| Error::Runtime {
2390                        message: format!(
2391                            "Failed to describe table via namespace client: {}. \
2392                             If you don't need managed versioning, don't pass namespace_client.",
2393                            e
2394                        ),
2395                    })?;
2396                response.managed_versioning == Some(true)
2397            }
2398            None => false,
2399        };
2400
2401        // Kept so that a `DatasetNotFound` can be re-checked against storage below.
2402        let recovery_params = params.clone();
2403        let mut builder = DatasetBuilder::from_uri(uri).with_read_params(params);
2404
2405        // Set up commit handler when managed_versioning is enabled
2406        if managed_versioning && let Some(ref ns_client) = namespace_client {
2407            let external_store = LanceNamespaceExternalManifestStore::for_table_uri(
2408                ns_client.clone(),
2409                table_id.clone(),
2410                uri,
2411            )?;
2412            let commit_handler: Arc<dyn CommitHandler> = Arc::new(ExternalManifestCommitHandler {
2413                external_manifest_store: Arc::new(external_store),
2414            });
2415            builder = builder.with_commit_handler(commit_handler);
2416        }
2417
2418        let dataset = match builder.load().await {
2419            Ok(dataset) => dataset,
2420            Err(e @ lance::Error::DatasetNotFound { .. }) => {
2421                return Err(map_dataset_not_found(uri, name, recovery_params, e).await);
2422            }
2423            Err(e) => return Err(e.into()),
2424        };
2425
2426        let dataset = DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval);
2427        let id = Self::build_id(&namespace, name);
2428
2429        Ok(Self {
2430            name: name.to_string(),
2431            namespace,
2432            id,
2433            uri: uri.to_string(),
2434            dataset,
2435            read_consistency_interval,
2436            namespace_client,
2437            pushdown_operations,
2438            freshness: None,
2439        })
2440    }
2441
2442    /// Set the namespace client for server-side query execution.
2443    ///
2444    /// When set, queries will be executed on the namespace server instead of locally.
2445    pub fn with_namespace_client(mut self, namespace_client: Arc<dyn LanceNamespace>) -> Self {
2446        self.namespace_client = Some(namespace_client);
2447        self
2448    }
2449
2450    /// Attach the read-freshness baseline handle (namespace connections only).
2451    pub(crate) fn with_freshness(mut self, freshness: TableFreshness) -> Self {
2452        self.freshness = Some(freshness);
2453        self
2454    }
2455
2456    /// Build a sibling `NativeTable` with the same identity but a different
2457    /// (independent) dataset wrapper — used to hand out branch-scoped handles.
2458    fn with_dataset(&self, dataset: dataset::DatasetConsistencyWrapper) -> Self {
2459        Self {
2460            name: self.name.clone(),
2461            namespace: self.namespace.clone(),
2462            id: self.id.clone(),
2463            uri: self.uri.clone(),
2464            dataset,
2465            read_consistency_interval: self.read_consistency_interval,
2466            namespace_client: self.namespace_client.clone(),
2467            pushdown_operations: self.pushdown_operations.clone(),
2468            freshness: self.freshness.clone(),
2469        }
2470    }
2471
2472    /// Bump the read-freshness baseline; no-op for non-namespace tables.
2473    fn bump_freshness(&self) {
2474        if let Some(freshness) = &self.freshness {
2475            freshness.bump();
2476        }
2477    }
2478
2479    fn validate_branch_name(name: &str, field: &str) -> Result<()> {
2480        if name.is_empty() {
2481            return Err(Error::InvalidInput {
2482                message: format!("{field} must be a non-empty string"),
2483            });
2484        }
2485        Ok(())
2486    }
2487
2488    /// Opens an existing Table using a namespace client.
2489    ///
2490    /// This method uses `DatasetBuilder::from_namespace` to open the table, which
2491    /// automatically fetches the table location and storage options from the namespace.
2492    /// This eliminates the need to pre-fetch and merge storage options before opening.
2493    ///
2494    /// # Arguments
2495    ///
2496    /// * `namespace_client` - The namespace client to use for fetching table metadata
2497    /// * `name` - The table name
2498    /// * `namespace` - The namespace path (e.g., vec!["parent", "child"])
2499    /// * `write_store_wrapper` - Optional wrapper for the object store on write path
2500    /// * `params` - Optional read parameters
2501    /// * `read_consistency_interval` - Optional interval for read consistency
2502    /// * `pushdown_operations` - Operations to push down to the namespace server.
2503    ///   When `QueryTable` is included, queries will be executed on the namespace server.
2504    /// * `session` - Optional session for object stores and caching
2505    ///
2506    /// # Returns
2507    ///
2508    /// * A [NativeTable] object.
2509    #[allow(clippy::too_many_arguments)]
2510    pub async fn open_from_namespace(
2511        namespace_client: Arc<dyn LanceNamespace>,
2512        name: &str,
2513        namespace: Vec<String>,
2514        write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2515        params: Option<ReadParams>,
2516        read_consistency_interval: Option<std::time::Duration>,
2517        pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2518        session: Option<Arc<lance::session::Session>>,
2519    ) -> Result<Self> {
2520        let mut params = params.unwrap_or_default();
2521
2522        // Set the session in read params
2523        if let Some(sess) = session {
2524            params.session(sess);
2525        }
2526
2527        // patch the params if we have a write store wrapper
2528        let params = match write_store_wrapper.clone() {
2529            Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
2530            None => params,
2531        };
2532
2533        // Build table_id from namespace + name
2534        let mut table_id = namespace.clone();
2535        table_id.push(name.to_string());
2536
2537        // Use DatasetBuilder::from_namespace which automatically fetches location
2538        // and storage options from the namespace
2539        let builder = DatasetBuilder::from_namespace(namespace_client.clone(), table_id)
2540            .await
2541            .map_err(|e| map_namespace_lance_error(e, name))?;
2542
2543        let dataset = builder
2544            .with_read_params(params)
2545            .load()
2546            .await
2547            .map_err(|e| match e {
2548                lance::Error::DatasetNotFound { .. } => Error::TableNotFound {
2549                    name: name.to_string(),
2550                    source: Box::new(e),
2551                },
2552                e => e.into(),
2553            })?;
2554
2555        let uri = dataset.uri().to_string();
2556        let dataset = DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval);
2557        let id = Self::build_id(&namespace, name);
2558
2559        let stored_namespace_client =
2560            if pushdown_operations.contains(&NamespaceClientPushdownOperation::QueryTable) {
2561                Some(namespace_client)
2562            } else {
2563                None
2564            };
2565
2566        Ok(Self {
2567            name: name.to_string(),
2568            namespace,
2569            id,
2570            uri,
2571            dataset,
2572            read_consistency_interval,
2573            namespace_client: stored_namespace_client,
2574            pushdown_operations,
2575            freshness: None,
2576        })
2577    }
2578
2579    fn get_table_name(uri: &str) -> Result<String> {
2580        let path = Path::new(uri);
2581        let name = path
2582            .file_stem()
2583            .ok_or(Error::TableNotFound {
2584                name: uri.to_string(),
2585                source: format!("Could not extract table name from URI: '{}'", uri).into(),
2586            })?
2587            .to_str()
2588            .ok_or(Error::InvalidTableName {
2589                name: uri.to_string(),
2590                reason: "Table name is not valid URL".to_string(),
2591            })?;
2592        Ok(name.to_string())
2593    }
2594
2595    fn build_id(namespace: &[String], name: &str) -> String {
2596        if namespace.is_empty() {
2597            name.to_string()
2598        } else {
2599            let mut parts = namespace.to_vec();
2600            parts.push(name.to_string());
2601            parts.join("$")
2602        }
2603    }
2604
2605    /// Creates a new Table
2606    ///
2607    /// # Arguments
2608    ///
2609    /// * `uri` - The URI to the table. When namespace is not empty, the caller must
2610    ///   provide an explicit URI (location) rather than deriving it from the table name.
2611    /// * `name` The Table name
2612    /// * `namespace` - The namespace path. When non-empty, an explicit URI must be provided.
2613    /// * `batches` RecordBatch to be saved in the database.
2614    /// * `params` - Write parameters.
2615    /// * `namespace_client` - Optional namespace client for namespace operations
2616    /// * `pushdown_operations` - Operations to push down to the namespace server
2617    ///
2618    /// # Returns
2619    ///
2620    /// * A [TableImpl] object.
2621    #[allow(clippy::too_many_arguments)]
2622    pub async fn create(
2623        uri: &str,
2624        name: &str,
2625        namespace: Vec<String>,
2626        batches: impl StreamingWriteSource,
2627        write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2628        params: Option<WriteParams>,
2629        read_consistency_interval: Option<std::time::Duration>,
2630        namespace_client: Option<Arc<dyn LanceNamespace>>,
2631        pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2632    ) -> Result<Self> {
2633        // Default params uses format v1.
2634        let params = params.unwrap_or(WriteParams {
2635            ..Default::default()
2636        });
2637        // patch the params if we have a write store wrapper
2638        let params = match write_store_wrapper.clone() {
2639            Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
2640            None => params,
2641        };
2642
2643        let insert_builder = InsertBuilder::new(uri).with_params(&params);
2644        let dataset = insert_builder
2645            .execute_stream(batches)
2646            .await
2647            .map_err(|e| match e {
2648                lance::Error::DatasetAlreadyExists { .. } => Error::TableAlreadyExists {
2649                    name: name.to_string(),
2650                },
2651                e => e.into(),
2652            })?;
2653
2654        let id = Self::build_id(&namespace, name);
2655
2656        Ok(Self {
2657            name: name.to_string(),
2658            namespace,
2659            id,
2660            uri: uri.to_string(),
2661            dataset: DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval),
2662            read_consistency_interval,
2663            namespace_client,
2664            pushdown_operations,
2665            freshness: None,
2666        })
2667    }
2668
2669    #[allow(clippy::too_many_arguments)]
2670    pub async fn create_empty(
2671        uri: &str,
2672        name: &str,
2673        namespace: Vec<String>,
2674        schema: SchemaRef,
2675        write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2676        params: Option<WriteParams>,
2677        read_consistency_interval: Option<std::time::Duration>,
2678        namespace_client: Option<Arc<dyn LanceNamespace>>,
2679        pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2680    ) -> Result<Self> {
2681        let data: Box<dyn Scannable> = Box::new(RecordBatch::new_empty(schema));
2682        Self::create(
2683            uri,
2684            name,
2685            namespace,
2686            data,
2687            write_store_wrapper,
2688            params,
2689            read_consistency_interval,
2690            namespace_client,
2691            pushdown_operations,
2692        )
2693        .await
2694    }
2695
2696    /// Creates a new Table using a namespace client for storage options.
2697    ///
2698    /// This method sets up a `StorageOptionsProvider` from the namespace client,
2699    /// enabling automatic credential refresh for cloud storage. The namespace
2700    /// is used for:
2701    /// 1. Setting up storage options provider for credential vending
2702    /// 2. Optionally enabling server-side query execution
2703    ///
2704    /// # Arguments
2705    ///
2706    /// * `namespace_client` - The namespace client to use for storage options
2707    /// * `uri` - The URI to the table (obtained from create_empty_table response)
2708    /// * `name` - The table name
2709    /// * `namespace` - The namespace path (e.g., vec!["parent", "child"])
2710    /// * `batches` - RecordBatch to be saved in the database
2711    /// * `write_store_wrapper` - Optional wrapper for the object store on write path
2712    /// * `params` - Optional write parameters
2713    /// * `read_consistency_interval` - Optional interval for read consistency
2714    /// * `pushdown_operations` - Operations to push down to the namespace server
2715    ///
2716    /// # Returns
2717    ///
2718    /// * A [NativeTable] object.
2719    #[allow(clippy::too_many_arguments)]
2720    pub async fn create_from_namespace(
2721        namespace_client: Arc<dyn LanceNamespace>,
2722        uri: &str,
2723        name: &str,
2724        namespace: Vec<String>,
2725        batches: impl StreamingWriteSource,
2726        write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2727        params: Option<WriteParams>,
2728        read_consistency_interval: Option<std::time::Duration>,
2729        pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2730        session: Option<Arc<lance::session::Session>>,
2731    ) -> Result<Self> {
2732        // Build table_id from namespace + name for the storage options provider
2733        let mut table_id = namespace.clone();
2734        table_id.push(name.to_string());
2735
2736        // Set up storage options provider from namespace
2737        let storage_options_provider = Arc::new(LanceNamespaceStorageOptionsProvider::new(
2738            namespace_client.clone(),
2739            table_id,
2740        ));
2741
2742        // Start with provided params or defaults
2743        let mut params = params.unwrap_or_default();
2744
2745        // Set the session in write params
2746        if let Some(sess) = session {
2747            params.session = Some(sess);
2748        }
2749
2750        // Ensure store_params exists and set the storage options provider
2751        let store_params = params
2752            .store_params
2753            .get_or_insert_with(ObjectStoreParams::default);
2754        let accessor = match store_params.storage_options().cloned() {
2755            Some(options) => {
2756                StorageOptionsAccessor::with_initial_and_provider(options, storage_options_provider)
2757            }
2758            None => StorageOptionsAccessor::with_provider(storage_options_provider),
2759        };
2760        store_params.storage_options_accessor = Some(Arc::new(accessor));
2761
2762        // Patch the params if we have a write store wrapper
2763        let params = match write_store_wrapper.clone() {
2764            Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
2765            None => params,
2766        };
2767
2768        let insert_builder = InsertBuilder::new(uri).with_params(&params);
2769        let dataset = insert_builder
2770            .execute_stream(batches)
2771            .await
2772            .map_err(|e| match e {
2773                lance::Error::DatasetAlreadyExists { .. } => Error::TableAlreadyExists {
2774                    name: name.to_string(),
2775                },
2776                e => e.into(),
2777            })?;
2778
2779        let id = Self::build_id(&namespace, name);
2780
2781        let stored_namespace_client =
2782            if pushdown_operations.contains(&NamespaceClientPushdownOperation::QueryTable) {
2783                Some(namespace_client)
2784            } else {
2785                None
2786            };
2787
2788        Ok(Self {
2789            name: name.to_string(),
2790            namespace,
2791            id,
2792            uri: uri.to_string(),
2793            dataset: DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval),
2794            read_consistency_interval,
2795            namespace_client: stored_namespace_client,
2796            pushdown_operations,
2797            freshness: None,
2798        })
2799    }
2800
2801    /// Merge new data into this table.
2802    pub async fn merge(
2803        &mut self,
2804        batches: impl RecordBatchReader + Send + 'static,
2805        left_on: &str,
2806        right_on: &str,
2807    ) -> Result<()> {
2808        self.dataset.ensure_mutable()?;
2809        let mut dataset = (*self.dataset.get().await?).clone();
2810        dataset.merge(batches, left_on, right_on).await?;
2811        self.dataset.update(dataset);
2812        Ok(())
2813    }
2814
2815    // TODO: why are these individual methods and not some single "get_stats" method?
2816    pub async fn count_fragments(&self) -> Result<usize> {
2817        Ok(self.dataset.get().await?.count_fragments())
2818    }
2819
2820    pub async fn count_deleted_rows(&self) -> Result<usize> {
2821        Ok(self.dataset.get().await?.count_deleted_rows().await?)
2822    }
2823
2824    pub async fn num_small_files(&self, max_rows_per_group: usize) -> Result<usize> {
2825        Ok(self
2826            .dataset
2827            .get()
2828            .await?
2829            .num_small_files(max_rows_per_group)
2830            .await)
2831    }
2832    /// Check whether the table uses V2 manifest paths.
2833    ///
2834    /// See [Self::migrate_manifest_paths_v2] and [ManifestNamingScheme] for
2835    /// more information.
2836    pub async fn uses_v2_manifest_paths(&self) -> Result<bool> {
2837        let dataset = self.dataset.get().await?;
2838        Ok(dataset.manifest_location().naming_scheme == ManifestNamingScheme::V2)
2839    }
2840
2841    /// Migrate the table to use the new manifest path scheme.
2842    ///
2843    /// This function will rename all V1 manifests to V2 manifest paths.
2844    /// These paths provide more efficient opening of datasets with many versions
2845    /// on object stores.
2846    ///
2847    /// This function is idempotent, and can be run multiple times without
2848    /// changing the state of the object store.
2849    ///
2850    /// However, it should not be run while other concurrent operations are happening.
2851    /// And it should also run until completion before resuming other operations.
2852    ///
2853    /// You can use [Self::uses_v2_manifest_paths] to check if the table is already
2854    /// using V2 manifest paths.
2855    pub async fn migrate_manifest_paths_v2(&self) -> Result<()> {
2856        self.dataset.ensure_mutable()?;
2857        let mut dataset = (*self.dataset.get().await?).clone();
2858        dataset.migrate_manifest_paths_v2().await?;
2859        self.dataset.update(dataset);
2860        Ok(())
2861    }
2862
2863    /// Get the table manifest
2864    pub async fn manifest(&self) -> Result<Manifest> {
2865        let dataset = self.dataset.get().await?;
2866        Ok(dataset.manifest().clone())
2867    }
2868
2869    /// Update key-value pairs in config.
2870    pub async fn update_config(
2871        &self,
2872        upsert_values: impl IntoIterator<Item = (String, String)>,
2873    ) -> Result<()> {
2874        self.dataset.ensure_mutable()?;
2875        let mut dataset = (*self.dataset.get().await?).clone();
2876        dataset.update_config(upsert_values).await?;
2877        self.dataset.update(dataset);
2878        Ok(())
2879    }
2880
2881    /// Delete keys from the config
2882    pub async fn delete_config_keys(&self, delete_keys: &[&str]) -> Result<()> {
2883        self.dataset.ensure_mutable()?;
2884        let mut dataset = (*self.dataset.get().await?).clone();
2885        // TODO: update this when we implement metadata APIs
2886        #[allow(deprecated)]
2887        dataset.delete_config_keys(delete_keys).await?;
2888        self.dataset.update(dataset);
2889        Ok(())
2890    }
2891
2892    /// Update schema metadata
2893    pub async fn replace_schema_metadata(
2894        &self,
2895        upsert_values: impl IntoIterator<Item = (String, String)>,
2896    ) -> Result<()> {
2897        self.dataset.ensure_mutable()?;
2898        let mut dataset = (*self.dataset.get().await?).clone();
2899        // TODO: update this when we implement metadata APIs
2900        #[allow(deprecated)]
2901        dataset.replace_schema_metadata(upsert_values).await?;
2902        self.dataset.update(dataset);
2903        Ok(())
2904    }
2905
2906    /// Update field metadata
2907    ///
2908    /// # Arguments:
2909    /// * `new_values` - An iterator of tuples where the first element is the
2910    ///   field id and the second element is a hashmap of metadata key-value
2911    ///   pairs.
2912    ///
2913    #[deprecated(since = "0.33.1", note = "Use `update_field_metadata` instead")]
2914    pub async fn replace_field_metadata(
2915        &self,
2916        new_values: impl IntoIterator<Item = (u32, HashMap<String, String>)>,
2917    ) -> Result<()> {
2918        self.dataset.ensure_mutable()?;
2919        let mut dataset = (*self.dataset.get().await?).clone();
2920        dataset.replace_field_metadata(new_values).await?;
2921        self.dataset.update(dataset);
2922        Ok(())
2923    }
2924}
2925
2926#[async_trait::async_trait]
2927impl BaseTable for NativeTable {
2928    fn as_any(&self) -> &dyn std::any::Any {
2929        self
2930    }
2931
2932    fn name(&self) -> &str {
2933        self.name.as_str()
2934    }
2935
2936    fn namespace(&self) -> &[String] {
2937        &self.namespace
2938    }
2939
2940    fn id(&self) -> &str {
2941        &self.id
2942    }
2943
2944    async fn version(&self) -> Result<u64> {
2945        Ok(self.dataset.get().await?.version().version)
2946    }
2947
2948    async fn checkout(&self, version: u64) -> Result<()> {
2949        self.dataset.as_time_travel(version).await
2950    }
2951
2952    async fn checkout_tag(&self, tag: &str) -> Result<()> {
2953        self.dataset.as_time_travel(tag).await
2954    }
2955
2956    async fn checkout_latest(&self) -> Result<()> {
2957        // Bump before resolving "latest" so that request carries the floor.
2958        self.bump_freshness();
2959        self.dataset.as_latest().await?;
2960        self.dataset.reload().await
2961    }
2962
2963    async fn create_branch(
2964        &self,
2965        name: &str,
2966        from: lance::dataset::refs::Ref,
2967    ) -> Result<Arc<dyn BaseTable>> {
2968        Self::validate_branch_name(name, "branch name")?;
2969        if let lance::dataset::refs::Ref::Version(Some(from_branch), _) = &from {
2970            Self::validate_branch_name(from_branch, "from_ref")?;
2971        }
2972        let mut ds = (*self.dataset.get().await?).clone();
2973        let branch_ds = ds.create_branch(name, from, None).await?;
2974        let dataset = dataset::DatasetConsistencyWrapper::new_latest(
2975            branch_ds,
2976            self.read_consistency_interval,
2977        );
2978        Ok(Arc::new(self.with_dataset(dataset)))
2979    }
2980
2981    async fn checkout_branch(&self, name: &str) -> Result<Arc<dyn BaseTable>> {
2982        Self::validate_branch_name(name, "branch name")?;
2983        let branch_ds = self.dataset.get().await?.checkout_branch(name).await?;
2984        let dataset = dataset::DatasetConsistencyWrapper::new_latest(
2985            branch_ds,
2986            self.read_consistency_interval,
2987        );
2988        Ok(Arc::new(self.with_dataset(dataset)))
2989    }
2990
2991    async fn checkout_branch_version(
2992        &self,
2993        name: &str,
2994        version: Option<u64>,
2995    ) -> Result<Arc<dyn BaseTable>> {
2996        let Some(version) = version else {
2997            return self.checkout_branch(name).await;
2998        };
2999        Self::validate_branch_name(name, "branch name")?;
3000        // Resolve (branch, version) in a single manifest read.
3001        let branch_ds = self
3002            .dataset
3003            .get()
3004            .await?
3005            .checkout_version((name, version))
3006            .await?;
3007        let dataset = dataset::DatasetConsistencyWrapper::new_time_travel(
3008            branch_ds,
3009            self.read_consistency_interval,
3010        );
3011        Ok(Arc::new(self.with_dataset(dataset)))
3012    }
3013
3014    async fn list_branches(&self) -> Result<HashMap<String, BranchContents>> {
3015        Ok(self.dataset.get().await?.list_branches().await?)
3016    }
3017
3018    async fn delete_branch(&self, name: &str) -> Result<()> {
3019        Self::validate_branch_name(name, "branch name")?;
3020        let mut ds = (*self.dataset.get().await?).clone();
3021        ds.delete_branch(name).await?;
3022        Ok(())
3023    }
3024
3025    fn current_branch(&self) -> Option<String> {
3026        self.dataset.current_branch()
3027    }
3028
3029    async fn list_versions(&self) -> Result<Vec<Version>> {
3030        Ok(self.dataset.get().await?.versions().await?)
3031    }
3032
3033    async fn restore(&self) -> Result<()> {
3034        let version = self
3035            .dataset
3036            .time_travel_version()
3037            .ok_or_else(|| Error::InvalidInput {
3038                message: "you must run checkout before running restore".to_string(),
3039            })?;
3040        {
3041            // restore is the only "write" operation allowed in time travel mode
3042            let mut dataset = (*self.dataset.get().await?).clone();
3043            debug_assert_eq!(dataset.version().version, version);
3044            dataset.restore().await?;
3045        }
3046        // Restore moves "latest", so bump before resolving it (as RemoteTable does).
3047        self.bump_freshness();
3048        self.dataset.as_latest().await?;
3049        Ok(())
3050    }
3051
3052    async fn schema(&self) -> Result<SchemaRef> {
3053        let lance_schema = self.dataset.get().await?.schema().clone();
3054        Ok(Arc::new(Schema::from(&lance_schema)))
3055    }
3056
3057    async fn table_definition(&self) -> Result<TableDefinition> {
3058        let schema = self.schema().await?;
3059        TableDefinition::try_from_rich_schema(schema)
3060    }
3061
3062    async fn count_rows(&self, filter: Option<Filter>) -> Result<usize> {
3063        let dataset = self.dataset.get().await?;
3064        match filter {
3065            None => Ok(dataset.count_rows(None).await?),
3066            Some(Filter::Sql(sql)) => Ok(dataset.count_rows(Some(sql)).await?),
3067            Some(Filter::Datafusion(_)) => Err(Error::NotSupported {
3068                message: "Datafusion filters are not yet supported".to_string(),
3069            }),
3070        }
3071    }
3072
3073    async fn add(&self, mut add: AddDataBuilder) -> Result<AddResult> {
3074        let table_def = self.table_definition().await?;
3075
3076        self.dataset.ensure_mutable()?;
3077        let ds_wrapper = self.dataset.clone();
3078        let ds = self.dataset.get().await?;
3079
3080        let table_schema = Schema::from(&ds.schema().clone());
3081
3082        let num_partitions = if let Some(parallelism) = add.write_parallelism {
3083            parallelism
3084        } else {
3085            // Peek at the first batch to estimate a good partition count for
3086            // write parallelism.
3087            let mut peeked = PeekedScannable::new(add.data);
3088            let n = if let Some(first_batch) = peeked.peek().await {
3089                let max_partitions = lance_core::utils::tokio::get_num_compute_intensive_cpus();
3090                estimate_write_partitions(
3091                    first_batch.get_array_memory_size(),
3092                    first_batch.num_rows(),
3093                    peeked.num_rows(),
3094                    max_partitions,
3095                )
3096            } else {
3097                1
3098            };
3099            add.data = Box::new(peeked);
3100            n
3101        };
3102
3103        let output = add.into_plan(&table_schema, &table_def)?;
3104
3105        let lance_params = output
3106            .write_options
3107            .lance_write_params
3108            .unwrap_or(WriteParams {
3109                mode: match output.mode {
3110                    AddDataMode::Append => WriteMode::Append,
3111                    AddDataMode::Overwrite => WriteMode::Overwrite,
3112                },
3113                ..Default::default()
3114            });
3115
3116        // Repartition for write parallelism if beneficial.
3117        let plan = if num_partitions > 1 {
3118            Arc::new(
3119                datafusion_physical_plan::repartition::RepartitionExec::try_new(
3120                    output.plan,
3121                    datafusion_physical_plan::Partitioning::RoundRobinBatch(num_partitions),
3122                )?,
3123            ) as Arc<dyn ExecutionPlan>
3124        } else {
3125            output.plan
3126        };
3127
3128        let insert_exec = Arc::new(InsertExec::new_with_tracker(
3129            ds_wrapper.clone(),
3130            ds,
3131            plan,
3132            lance_params,
3133            output.tracker.clone(),
3134        ));
3135
3136        let tracker_for_tasks = output.tracker.clone();
3137        if let Some(ref t) = tracker_for_tasks {
3138            t.set_total_tasks(num_partitions);
3139        }
3140        let _finish = write_progress::FinishOnDrop(output.tracker);
3141
3142        // Execute all partitions in parallel.
3143        let task_ctx = Arc::new(TaskContext::default());
3144        let handles = FuturesUnordered::new();
3145        for partition in 0..num_partitions {
3146            let exec = insert_exec.clone();
3147            let ctx = task_ctx.clone();
3148            let tracker = tracker_for_tasks.clone();
3149            handles.push(tokio::spawn(async move {
3150                let _guard = tracker.as_ref().map(|t| t.track_task());
3151                let mut stream = exec
3152                    .execute(partition, ctx)
3153                    .map_err(|e| -> Error { e.into() })?;
3154                while let Some(batch) = stream.next().await {
3155                    batch.map_err(|e| -> Error { e.into() })?;
3156                }
3157                Ok::<_, Error>(())
3158            }));
3159        }
3160        for handle in handles {
3161            handle.await.map_err(|e| Error::Runtime {
3162                message: format!("Insert task panicked: {}", e),
3163            })??;
3164        }
3165
3166        let version = ds_wrapper.get().await?.manifest().version;
3167        self.bump_freshness();
3168        Ok(AddResult { version })
3169    }
3170
3171    async fn create_index(&self, opts: IndexBuilder) -> Result<()> {
3172        let prepared = self.prepare_index(&opts).await?;
3173        self.build_index(opts, prepared).await
3174    }
3175
3176    async fn create_index_async(&self, opts: IndexBuilder) -> Result<Job> {
3177        // Prepare before spawning so bad input is reported by this call rather
3178        // than only by the job.
3179        let prepared = self.prepare_index(&opts).await?;
3180        let table = self.clone();
3181        Ok(Job::spawned(tokio::spawn(async move {
3182            table.build_index(opts, prepared).await
3183        })))
3184    }
3185
3186    async fn drop_index(&self, index_name: &str) -> Result<()> {
3187        self.dataset.ensure_mutable()?;
3188        let mut dataset = (*self.dataset.get().await?).clone();
3189        dataset.drop_index(index_name).await?;
3190        self.dataset.update(dataset);
3191        Ok(())
3192    }
3193
3194    async fn prewarm_index(&self, index_name: &str) -> Result<()> {
3195        let dataset = self.dataset.get().await?;
3196        Ok(dataset.prewarm_index(index_name).await?)
3197    }
3198
3199    async fn prewarm_data(&self, _columns: Option<Vec<String>>) -> Result<()> {
3200        Err(Error::NotSupported {
3201            message: "prewarm_data is currently only supported on remote tables.".into(),
3202        })
3203    }
3204
3205    async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult> {
3206        // Delegate to the submodule implementation
3207        let result = update::execute_update(self, update).await?;
3208        self.bump_freshness();
3209        Ok(result)
3210    }
3211
3212    async fn create_plan(
3213        &self,
3214        query: &AnyQuery,
3215        options: QueryExecutionOptions,
3216    ) -> Result<Arc<dyn ExecutionPlan>> {
3217        query::create_plan(self, query, options).await
3218    }
3219
3220    async fn query(
3221        &self,
3222        query: &AnyQuery,
3223        options: QueryExecutionOptions,
3224    ) -> Result<DatasetRecordBatchStream> {
3225        query::execute_query(self, query, options).await
3226    }
3227
3228    async fn analyze_plan(
3229        &self,
3230        query: &AnyQuery,
3231        options: QueryExecutionOptions,
3232    ) -> Result<String> {
3233        query::analyze_query_plan(self, query, options).await
3234    }
3235
3236    async fn merge_insert(
3237        &self,
3238        params: MergeInsertBuilder,
3239        new_data: Box<dyn RecordBatchReader + Send>,
3240    ) -> Result<MergeResult> {
3241        let result = merge::execute_merge_insert(self, params, new_data).await?;
3242        self.bump_freshness();
3243        Ok(result)
3244    }
3245
3246    async fn set_unenforced_primary_key(&self, columns: &[&str]) -> Result<()> {
3247        primary_key::set_unenforced_primary_key(self, columns).await
3248    }
3249
3250    async fn set_lsm_write_spec(&self, spec: LsmWriteSpec) -> Result<()> {
3251        merge::lsm::set_lsm_write_spec(self, spec).await
3252    }
3253
3254    async fn unset_lsm_write_spec(&self) -> Result<()> {
3255        merge::lsm::unset_lsm_write_spec(self).await
3256    }
3257
3258    async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
3259        merge::lsm::get_lsm_write_spec(self).await
3260    }
3261
3262    async fn close_lsm_writers(&self) -> Result<()> {
3263        merge::lsm::close_lsm_writers(self).await
3264    }
3265
3266    async fn blob_columns(&self) -> Result<Vec<String>> {
3267        let schema = self.schema().await?;
3268        Ok(crate::blob::blob_column_names(schema.as_ref()))
3269    }
3270
3271    async fn fetch_blobs(&self, column: &str, row_ids: &[u64]) -> Result<LargeBinaryArray> {
3272        let dataset = self.dataset.get().await?;
3273        crate::blob::take_blobs_aligned(&dataset, column, row_ids).await
3274    }
3275
3276    async fn fetch_blob_ranges(
3277        &self,
3278        column: &str,
3279        requests: &[BlobRangeRequest],
3280    ) -> Result<LargeBinaryArray> {
3281        let dataset = self.dataset.get().await?;
3282        crate::blob::take_blob_ranges_aligned(&dataset, column, requests).await
3283    }
3284
3285    async fn fetch_blob_files(
3286        &self,
3287        column: &str,
3288        row_ids: &[u64],
3289    ) -> Result<Vec<Option<BlobFile>>> {
3290        let dataset = self.dataset.get().await?;
3291        crate::blob::take_blob_files_aligned(&dataset, column, row_ids).await
3292    }
3293
3294    /// Delete rows from the table
3295    async fn delete(&self, predicate: Predicate<'_>) -> Result<DeleteResult> {
3296        let result = delete::execute_delete(self, predicate).await?;
3297        self.bump_freshness();
3298        Ok(result)
3299    }
3300
3301    async fn tags(&self) -> Result<Box<dyn Tags + '_>> {
3302        Ok(Box::new(NativeTags {
3303            dataset: self.dataset.clone(),
3304        }))
3305    }
3306
3307    async fn optimize(&self, action: OptimizeAction) -> Result<OptimizeStats> {
3308        // Delegate to the submodule implementation
3309        optimize::execute_optimize(self, action).await
3310    }
3311
3312    async fn add_columns(
3313        &self,
3314        transforms: NewColumnTransform,
3315        read_columns: Option<Vec<String>>,
3316    ) -> Result<AddColumnsResult> {
3317        let result = schema_evolution::execute_add_columns(self, transforms, read_columns).await?;
3318        self.bump_freshness();
3319        Ok(result)
3320    }
3321
3322    async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult> {
3323        let result = schema_evolution::execute_alter_columns(self, alterations).await?;
3324        self.bump_freshness();
3325        Ok(result)
3326    }
3327
3328    async fn update_field_metadata(
3329        &self,
3330        updates: &[FieldMetadataUpdate],
3331    ) -> Result<UpdateFieldMetadataResult> {
3332        let result = schema_evolution::execute_update_field_metadata(self, updates).await?;
3333        self.bump_freshness();
3334        Ok(result)
3335    }
3336
3337    async fn drop_columns(&self, columns: &[&str]) -> Result<DropColumnsResult> {
3338        let result = schema_evolution::execute_drop_columns(self, columns).await?;
3339        self.bump_freshness();
3340        Ok(result)
3341    }
3342
3343    async fn list_indices(&self) -> Result<Vec<IndexConfig>> {
3344        let dataset = self.dataset.get().await?;
3345        let total_rows = dataset.count_rows(None).await? as u64;
3346        let descriptions = dataset.describe_indices(None).await?;
3347        let mut indices: Vec<IndexConfig> = descriptions
3348            .iter()
3349            .filter_map(|idx_desc| {
3350                let index_type: crate::index::IndexType = idx_desc
3351                    .index_type()
3352                    .parse()
3353                    .unwrap_or(crate::index::IndexType::Unknown);
3354                if index_type == crate::index::IndexType::Unknown {
3355                    // Internal or future index types that this version doesn't recognize
3356                    // (e.g. Lance's internal FragReuseIndex) are silently excluded from
3357                    // the user-visible index listing.
3358                    log::debug!(
3359                        "Skipping unrecognized index '{}' (type '{}') in list_indices",
3360                        idx_desc.name(),
3361                        idx_desc.index_type(),
3362                    );
3363                    return None;
3364                }
3365
3366                let field_ids = idx_desc.field_ids();
3367                let mut columns = Vec::with_capacity(field_ids.len());
3368                for field_id in field_ids {
3369                    let field_path = match dataset.schema().field_path(*field_id as i32) {
3370                        Ok(field_path) => field_path,
3371                        Err(e) => {
3372                            log::warn!(
3373                                "Failed to resolve field path for index {} field id {}: {}",
3374                                idx_desc.name(),
3375                                field_id,
3376                                e
3377                            );
3378                            return None;
3379                        }
3380                    };
3381                    columns.push(field_path);
3382                }
3383
3384                let segments = idx_desc.segments();
3385                let index_uuid = segments.first().map(|seg| seg.uuid.to_string());
3386                let created_at = segments.iter().filter_map(|seg| seg.created_at).min();
3387                let index_version = segments.first().map(|seg| seg.index_version);
3388                let num_indexed_rows = idx_desc.rows_indexed();
3389
3390                Some(IndexConfig {
3391                    name: idx_desc.name().to_string(),
3392                    index_type,
3393                    columns,
3394                    index_uuid,
3395                    type_url: Some(idx_desc.type_url().to_string()),
3396                    created_at,
3397                    num_indexed_rows: Some(num_indexed_rows),
3398                    num_unindexed_rows: Some(total_rows.saturating_sub(num_indexed_rows)),
3399                    size_bytes: idx_desc.total_size_bytes(),
3400                    num_segments: Some(segments.len() as u32),
3401                    index_version,
3402                    index_details: idx_desc.details().ok(),
3403                })
3404            })
3405            .collect();
3406
3407        for index in indices
3408            .iter_mut()
3409            .filter(|index| index.index_type == crate::index::IndexType::FTS)
3410        {
3411            let Some(description) = descriptions
3412                .iter()
3413                .find(|description| description.name() == index.name)
3414            else {
3415                continue;
3416            };
3417            let segments = description.segments();
3418            let Some(segment) = segments.first() else {
3419                continue;
3420            };
3421            let params = load_segment_params(&dataset, segment).await?;
3422            let details = serde_json::to_string(&params).map_err(|source| Error::Other {
3423                message: format!(
3424                    "Failed to serialize full text search configuration for index '{}'",
3425                    index.name
3426                ),
3427                source: Some(Box::new(source)),
3428            })?;
3429            index.index_details = Some(details);
3430        }
3431        Ok(indices)
3432    }
3433
3434    async fn uri(&self) -> Result<String> {
3435        Ok(self.uri.clone())
3436    }
3437
3438    async fn storage_options(&self) -> Option<HashMap<String, String>> {
3439        self.initial_storage_options().await
3440    }
3441
3442    async fn initial_storage_options(&self) -> Option<HashMap<String, String>> {
3443        self.dataset
3444            .get()
3445            .await
3446            .ok()
3447            .and_then(|dataset| dataset.initial_storage_options().cloned())
3448    }
3449
3450    async fn latest_storage_options(&self) -> Result<Option<HashMap<String, String>>> {
3451        let dataset = self.dataset.get().await?;
3452        Ok(dataset.latest_storage_options().await?.map(|o| o.0))
3453    }
3454
3455    async fn index_stats(&self, index_name: &str) -> Result<Option<IndexStatistics>> {
3456        // describe_indices() reads only manifest-level metadata (no index file I/O).
3457        // VectorIndexDetails in the manifest carries distance_type for indices written
3458        // by recent Lance versions. For older datasets that didn't write those details
3459        // we fall back to index_statistics() for vector index types.
3460        let dataset = self.dataset.get().await?;
3461
3462        let mut descriptions = dataset
3463            .describe_indices(Some(IndexCriteria::default().with_name(index_name)))
3464            .await?;
3465        let Some(description) = descriptions.pop() else {
3466            return Ok(None);
3467        };
3468
3469        let index_type: crate::index::IndexType = description
3470            .index_type()
3471            .parse()
3472            .unwrap_or(crate::index::IndexType::Unknown);
3473
3474        let is_vector = matches!(
3475            index_type,
3476            crate::index::IndexType::IvfFlat
3477                | crate::index::IndexType::IvfSq
3478                | crate::index::IndexType::IvfPq
3479                | crate::index::IndexType::IvfRq
3480                | crate::index::IndexType::IvfHnswPq
3481                | crate::index::IndexType::IvfHnswSq
3482                | crate::index::IndexType::IvfHnswFlat
3483        );
3484
3485        // details() serializes VectorIndexDetails to JSON with an uppercase "metric_type"
3486        // field (e.g. "L2", "COSINE"). Parse it with a case-insensitive match.
3487        let distance_type = description.details().ok().and_then(|json| {
3488            #[derive(serde::Deserialize)]
3489            struct Details {
3490                metric_type: Option<String>,
3491            }
3492            serde_json::from_str::<Details>(&json)
3493                .ok()
3494                .and_then(|d| d.metric_type)
3495                .and_then(|m| match m.to_uppercase().as_str() {
3496                    "L2" => Some(DistanceType::L2),
3497                    "COSINE" => Some(DistanceType::Cosine),
3498                    "DOT" => Some(DistanceType::Dot),
3499                    "HAMMING" => Some(DistanceType::Hamming),
3500                    _ => None,
3501                })
3502        });
3503
3504        // Older Lance datasets didn't write VectorIndexDetails, so distance_type won't
3505        // be in the manifest. Fall back to index_statistics() only in that case.
3506        if is_vector && distance_type.is_none() {
3507            let stats = dataset.index_statistics(index_name).await?;
3508            let mut stats: IndexStatisticsImpl =
3509                serde_json::from_str(&stats).map_err(|e| Error::InvalidInput {
3510                    message: format!("error deserializing index statistics: {}", e),
3511                })?;
3512            let first_index = stats.indices.pop().ok_or_else(|| Error::InvalidInput {
3513                message: "index statistics is empty".to_string(),
3514            })?;
3515            return Ok(Some(IndexStatistics {
3516                num_indexed_rows: stats.num_indexed_rows,
3517                num_unindexed_rows: stats.num_unindexed_rows,
3518                index_type,
3519                distance_type: first_index.metric_type,
3520                num_indices: stats.num_indices,
3521            }));
3522        }
3523
3524        let num_indexed_rows = description.rows_indexed() as usize;
3525        let total_rows = dataset.count_rows(None).await?;
3526        let num_unindexed_rows = total_rows.saturating_sub(num_indexed_rows);
3527        Ok(Some(IndexStatistics {
3528            num_indexed_rows,
3529            num_unindexed_rows,
3530            index_type,
3531            distance_type,
3532            num_indices: Some(description.metadata().len() as u32),
3533        }))
3534    }
3535
3536    /// Poll until the columns are fully indexed. Will return Error::Timeout if the columns
3537    /// are not fully indexed within the timeout.
3538    async fn wait_for_index(
3539        &self,
3540        index_names: &[&str],
3541        timeout: std::time::Duration,
3542    ) -> Result<()> {
3543        wait_for_index(self, index_names, timeout).await
3544    }
3545
3546    async fn stats(&self) -> Result<TableStatistics> {
3547        let num_rows = self.count_rows(None).await?;
3548        let num_indices = self.list_indices().await?.len();
3549        let ds = self.dataset.get().await?;
3550        // Sizes come from the manifest. Summing per-field `bytes_on_disk` instead
3551        // would open every data file to read its column metadata, which costs one
3552        // IO per fragment and reports 0 for legacy v1 storage.
3553        //
3554        // The manifest summary covers only the fragments' base data files, so
3555        // overlay files (recorded on each fragment) and index files (recorded in
3556        // the manifest's index section) are added separately.
3557        let mut total_bytes = ds.manifest().summary().total_files_size as usize;
3558        for frag in ds.manifest().fragments.iter() {
3559            for overlay in &frag.overlays {
3560                if let Some(size) = overlay.data_file.file_size_bytes.get() {
3561                    total_bytes += size.get() as usize;
3562                }
3563            }
3564        }
3565        for index in ds.load_indices().await?.iter() {
3566            total_bytes += index.total_size_bytes().unwrap_or(0) as usize;
3567        }
3568
3569        let frags = ds.get_fragments();
3570        let mut sorted_sizes = join_all(
3571            frags
3572                .iter()
3573                .map(|frag| async move { frag.physical_rows().await.unwrap_or(0) }),
3574        )
3575        .await;
3576        sorted_sizes.sort();
3577
3578        let small_frag_threshold = 100000;
3579        let num_fragments = sorted_sizes.len();
3580        let num_small_fragments = sorted_sizes
3581            .iter()
3582            .filter(|&&size| size < small_frag_threshold)
3583            .count();
3584
3585        let p25 = *sorted_sizes.get(num_fragments / 4).unwrap_or(&0);
3586        let p50 = *sorted_sizes.get(num_fragments / 2).unwrap_or(&0);
3587        let p75 = *sorted_sizes.get(num_fragments * 3 / 4).unwrap_or(&0);
3588        let p99 = *sorted_sizes.get(num_fragments * 99 / 100).unwrap_or(&0);
3589        let min = sorted_sizes.first().copied().unwrap_or(0);
3590        let max = sorted_sizes.last().copied().unwrap_or(0);
3591        let mean = sorted_sizes
3592            .iter()
3593            .copied()
3594            .sum::<usize>()
3595            .checked_div(num_fragments)
3596            .unwrap_or(0);
3597
3598        let frag_stats = FragmentStatistics {
3599            num_fragments,
3600            num_small_fragments,
3601            lengths: FragmentSummaryStats {
3602                min,
3603                max,
3604                mean,
3605                p25,
3606                p50,
3607                p75,
3608                p99,
3609            },
3610        };
3611        let stats = TableStatistics {
3612            total_bytes,
3613            num_rows,
3614            num_indices,
3615            fragment_stats: frag_stats,
3616        };
3617        Ok(stats)
3618    }
3619
3620    async fn create_insert_exec(
3621        &self,
3622        input: Arc<dyn datafusion_physical_plan::ExecutionPlan>,
3623        write_params: WriteParams,
3624    ) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
3625        let ds = self.dataset.get().await?;
3626        let dataset = Arc::new((*ds).clone());
3627        Ok(Arc::new(datafusion::insert::InsertExec::new(
3628            self.dataset.clone(),
3629            dataset,
3630            input,
3631            write_params,
3632        )))
3633    }
3634}
3635
3636#[skip_serializing_none]
3637#[derive(Debug, Deserialize, PartialEq)]
3638pub struct TableStatistics {
3639    /// The total size, in bytes, of the table's data files, index files, and
3640    /// overlay files
3641    ///
3642    /// Read from the manifest, so this excludes deletion files and manifests,
3643    /// and it excludes any file whose size the manifest does not record
3644    /// (tables and indices written before writers persisted file sizes).
3645    pub total_bytes: usize,
3646
3647    /// The number of rows in the table
3648    pub num_rows: usize,
3649
3650    /// The number of indices in the table
3651    pub num_indices: usize,
3652
3653    /// Statistics on table fragments
3654    pub fragment_stats: FragmentStatistics,
3655}
3656
3657#[skip_serializing_none]
3658#[derive(Debug, Deserialize, PartialEq)]
3659pub struct FragmentStatistics {
3660    /// The number of fragments in the table
3661    pub num_fragments: usize,
3662
3663    /// The number of uncompacted fragments in the table
3664    pub num_small_fragments: usize,
3665
3666    /// Statistics on the number of rows in the table fragments
3667    pub lengths: FragmentSummaryStats,
3668    // todo: add size statistics
3669    // /// Statistics on the number of bytes in the table fragments
3670    // sizes: FragmentStats,
3671}
3672
3673#[skip_serializing_none]
3674#[derive(Debug, Deserialize, PartialEq)]
3675pub struct FragmentSummaryStats {
3676    pub min: usize,
3677    pub max: usize,
3678    pub mean: usize,
3679    pub p25: usize,
3680    pub p50: usize,
3681    pub p75: usize,
3682    pub p99: usize,
3683}
3684
3685#[cfg(test)]
3686#[allow(deprecated)]
3687mod tests {
3688    use std::sync::Arc;
3689    use std::sync::atomic::{AtomicBool, Ordering};
3690    use std::time::Duration;
3691
3692    use arrow_array::{
3693        Int32Array, RecordBatch, RecordBatchIterator, RecordBatchReader, StringArray,
3694    };
3695    use arrow_schema::{DataType, Field, Schema};
3696    use futures::TryStreamExt;
3697    use lance::Dataset;
3698    use lance::io::{ObjectStoreParams, WrappingObjectStore};
3699    use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION;
3700    use tempfile::tempdir;
3701
3702    use super::*;
3703    use crate::connect;
3704    use crate::connection::ConnectBuilder;
3705    use crate::io::object_store::io_tracking::IoTrackingStore;
3706    use crate::query::Select;
3707    use crate::query::{ExecutableQuery, QueryBase};
3708    use crate::test_utils::connection::new_test_connection;
3709
3710    #[test]
3711    fn test_tokenize_uses_explicit_simple_tokenizer() {
3712        let params =
3713            crate::index::scalar::FtsIndexBuilder::default().base_tokenizer("simple".to_string());
3714        let tokens = crate::tokenize("Running in cafés", &params).unwrap();
3715
3716        assert_eq!(
3717            tokens,
3718            vec![
3719                FtsToken {
3720                    text: "run".to_string(),
3721                    position: 0,
3722                },
3723                FtsToken {
3724                    text: "cafe".to_string(),
3725                    position: 2,
3726                },
3727            ]
3728        );
3729    }
3730
3731    #[tokio::test]
3732    async fn test_open() {
3733        let tmp_dir = tempdir().unwrap();
3734        let dataset_path = tmp_dir.path().join("test.lance");
3735
3736        let batch = make_test_batches();
3737        let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
3738        Dataset::write(reader, dataset_path.to_str().unwrap(), None)
3739            .await
3740            .unwrap();
3741
3742        let table = NativeTable::open(dataset_path.to_str().unwrap())
3743            .await
3744            .unwrap();
3745
3746        assert_eq!(table.name, "test")
3747    }
3748
3749    #[tokio::test]
3750    async fn test_open_not_found() {
3751        let tmp_dir = tempdir().unwrap();
3752        let uri = tmp_dir.path().to_str().unwrap();
3753        let table = NativeTable::open(uri).await;
3754        assert!(matches!(table.unwrap_err(), Error::TableNotFound { .. }));
3755    }
3756
3757    #[tokio::test]
3758    async fn test_open_not_found_missing_lance_dir() {
3759        let tmp_dir = tempdir().unwrap();
3760        let dataset_path = tmp_dir.path().join("test.lance");
3761
3762        let err = NativeTable::open(dataset_path.to_str().unwrap())
3763            .await
3764            .unwrap_err();
3765        assert!(
3766            matches!(&err, Error::TableNotFound { name, .. } if name == "test"),
3767            "got {err:?}"
3768        );
3769    }
3770
3771    /// Write a table and then break it, leaving the `<name>.lance` directory in place.
3772    ///
3773    /// `remove_all` reproduces an interrupted drop + re-create (the directory is left
3774    /// empty); otherwise only the manifests are removed, leaving the data files behind.
3775    async fn write_then_corrupt_table(dir: &std::path::Path, remove_all: bool) -> String {
3776        let dataset_path = dir.join("test.lance");
3777        let uri = dataset_path.to_str().unwrap().to_string();
3778
3779        let batch = make_test_batches();
3780        let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
3781        Dataset::write(reader, &uri, None).await.unwrap();
3782
3783        if remove_all {
3784            for entry in std::fs::read_dir(&dataset_path).unwrap() {
3785                let entry = entry.unwrap();
3786                if entry.file_type().unwrap().is_dir() {
3787                    std::fs::remove_dir_all(entry.path()).unwrap();
3788                } else {
3789                    std::fs::remove_file(entry.path()).unwrap();
3790                }
3791            }
3792            assert_eq!(std::fs::read_dir(&dataset_path).unwrap().count(), 0);
3793        } else {
3794            let versions = dataset_path.join("_versions");
3795            assert!(versions.is_dir(), "expected manifests under {versions:?}");
3796            std::fs::remove_dir_all(&versions).unwrap();
3797            assert!(std::fs::read_dir(&dataset_path).unwrap().count() > 0);
3798        }
3799
3800        uri
3801    }
3802
3803    #[tokio::test]
3804    async fn test_open_corrupt_empty_dir() {
3805        let tmp_dir = tempdir().unwrap();
3806        let uri = write_then_corrupt_table(tmp_dir.path(), true).await;
3807
3808        let err = NativeTable::open(&uri).await.unwrap_err();
3809        assert!(
3810            matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
3811            "got {err:?}"
3812        );
3813    }
3814
3815    #[tokio::test]
3816    async fn test_open_corrupt_missing_manifest() {
3817        let tmp_dir = tempdir().unwrap();
3818        let uri = write_then_corrupt_table(tmp_dir.path(), false).await;
3819
3820        let err = NativeTable::open(&uri).await.unwrap_err();
3821        assert!(
3822            matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
3823            "got {err:?}"
3824        );
3825    }
3826
3827    /// A table listed by `table_names()` must not be reported as missing by
3828    /// `open_table()`. See <https://github.com/lancedb/lancedb/issues/3127>.
3829    #[tokio::test]
3830    async fn test_open_table_corrupt_is_still_listed() {
3831        let tmp_dir = tempdir().unwrap();
3832        let db = connect(tmp_dir.path().to_str().unwrap())
3833            .execute()
3834            .await
3835            .unwrap();
3836
3837        write_then_corrupt_table(tmp_dir.path(), true).await;
3838
3839        assert_eq!(
3840            db.table_names().execute().await.unwrap(),
3841            vec!["test".to_string()]
3842        );
3843        let err = db.open_table("test").execute().await.unwrap_err();
3844        assert!(
3845            matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
3846            "got {err:?}"
3847        );
3848        assert!(
3849            err.to_string().contains("exists but could not be loaded"),
3850            "got {err}"
3851        );
3852    }
3853
3854    #[test]
3855    #[cfg(not(windows))]
3856    fn test_object_store_path() {
3857        use std::path::Path as StdPath;
3858        let p = StdPath::new("s3://bucket/path/to/file");
3859        let c = p.join("subfile");
3860        assert_eq!(c.to_str().unwrap(), "s3://bucket/path/to/file/subfile");
3861    }
3862
3863    #[tokio::test]
3864    async fn test_count_rows() {
3865        let tmp_dir = tempdir().unwrap();
3866        let uri = tmp_dir.path().to_str().unwrap();
3867
3868        let batch = make_test_batches();
3869        let reader: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
3870            vec![Ok(batch.clone())],
3871            batch.schema(),
3872        ));
3873        let table = NativeTable::create(
3874            uri,
3875            "test",
3876            vec![],
3877            reader,
3878            None,
3879            None,
3880            None,
3881            None,
3882            HashSet::new(),
3883        )
3884        .await
3885        .unwrap();
3886
3887        assert_eq!(table.count_rows(None).await.unwrap(), 10);
3888        assert_eq!(
3889            table
3890                .count_rows(Some(Filter::Sql("i >= 5".to_string())))
3891                .await
3892                .unwrap(),
3893            5
3894        );
3895    }
3896
3897    #[derive(Default, Debug)]
3898    struct NoOpCacheWrapper {
3899        called: AtomicBool,
3900    }
3901
3902    impl NoOpCacheWrapper {
3903        fn called(&self) -> bool {
3904            self.called.load(Ordering::Relaxed)
3905        }
3906    }
3907
3908    impl WrappingObjectStore for NoOpCacheWrapper {
3909        fn wrap(
3910            &self,
3911            _store_prefix: &str,
3912            original: Arc<dyn object_store::ObjectStore>,
3913        ) -> Arc<dyn object_store::ObjectStore> {
3914            self.called.store(true, Ordering::Relaxed);
3915            original
3916        }
3917    }
3918
3919    #[tokio::test]
3920    async fn test_open_table_options() {
3921        let tmp_dir = tempdir().unwrap();
3922        let dataset_path = tmp_dir.path().join("test.lance");
3923        let uri = dataset_path.to_str().unwrap();
3924        let conn = connect(uri).execute().await.unwrap();
3925
3926        let batches = make_test_batches();
3927
3928        conn.create_table("my_table", batches)
3929            .execute()
3930            .await
3931            .unwrap();
3932
3933        let wrapper = Arc::new(NoOpCacheWrapper::default());
3934
3935        let object_store_params = ObjectStoreParams {
3936            object_store_wrapper: Some(wrapper.clone()),
3937            ..Default::default()
3938        };
3939        let param = ReadParams {
3940            store_options: Some(object_store_params),
3941            ..Default::default()
3942        };
3943        assert!(!wrapper.called());
3944        conn.open_table("my_table")
3945            .lance_read_params(param)
3946            .execute()
3947            .await
3948            .unwrap();
3949        assert!(wrapper.called());
3950    }
3951
3952    fn make_test_batches() -> RecordBatch {
3953        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
3954        RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from_iter_values(0..10))]).unwrap()
3955    }
3956
3957    #[tokio::test]
3958    async fn test_tags() {
3959        let tmp_dir = tempdir().unwrap();
3960        let uri = tmp_dir.path().to_str().unwrap();
3961
3962        let conn = ConnectBuilder::new(uri)
3963            .read_consistency_interval(Duration::from_secs(0))
3964            .execute()
3965            .await
3966            .unwrap();
3967        let table = conn
3968            .create_table("my_table", some_sample_data())
3969            .execute()
3970            .await
3971            .unwrap();
3972        assert_eq!(table.version().await.unwrap(), 1);
3973        table.add(some_sample_data()).execute().await.unwrap();
3974        assert_eq!(table.version().await.unwrap(), 2);
3975        let mut tags_manager = table.tags().await.unwrap();
3976        let tags = tags_manager.list().await.unwrap();
3977        assert!(tags.is_empty(), "Tags should be empty initially");
3978        let tag1 = "tag1";
3979        tags_manager.create(tag1, 1).await.unwrap();
3980        assert_eq!(tags_manager.get_version(tag1).await.unwrap(), 1);
3981        let tags = tags_manager.list().await.unwrap();
3982        assert_eq!(tags.len(), 1);
3983        assert!(tags.contains_key(tag1));
3984        assert_eq!(tags.get(tag1).unwrap().version, 1);
3985        tags_manager.create("tag2", 2).await.unwrap();
3986        assert_eq!(tags_manager.get_version("tag2").await.unwrap(), 2);
3987        let tags = tags_manager.list().await.unwrap();
3988        assert_eq!(tags.len(), 2);
3989        assert!(tags.contains_key(tag1));
3990        assert_eq!(tags.get(tag1).unwrap().version, 1);
3991        assert!(tags.contains_key("tag2"));
3992        assert_eq!(tags.get("tag2").unwrap().version, 2);
3993        // Test update and delete
3994        table.add(some_sample_data()).execute().await.unwrap();
3995        tags_manager.update(tag1, 3).await.unwrap();
3996        assert_eq!(tags_manager.get_version(tag1).await.unwrap(), 3);
3997        tags_manager.delete("tag2").await.unwrap();
3998        let tags = tags_manager.list().await.unwrap();
3999        assert_eq!(tags.len(), 1);
4000        assert!(tags.contains_key(tag1));
4001        assert_eq!(tags.get(tag1).unwrap().version, 3);
4002        // Test checkout tag
4003        table.add(some_sample_data()).execute().await.unwrap();
4004        assert_eq!(table.version().await.unwrap(), 4);
4005        table.checkout_tag(tag1).await.unwrap();
4006        assert_eq!(table.version().await.unwrap(), 3);
4007        table.checkout_latest().await.unwrap();
4008        assert_eq!(table.version().await.unwrap(), 4);
4009    }
4010
4011    #[tokio::test]
4012    async fn test_branches() {
4013        let tmp_dir = tempdir().unwrap();
4014        let uri = tmp_dir.path().to_str().unwrap();
4015
4016        let conn = ConnectBuilder::new(uri)
4017            .read_consistency_interval(Duration::from_secs(0))
4018            .execute()
4019            .await
4020            .unwrap();
4021
4022        // main: one row at v1
4023        let table = conn
4024            .create_table("my_table", some_sample_data())
4025            .execute()
4026            .await
4027            .unwrap();
4028        assert_eq!(table.count_rows(None).await.unwrap(), 1);
4029        assert_eq!(table.current_branch(), None);
4030        let main_version = table.version().await.unwrap();
4031
4032        // branch off main's current version; it starts with main's data
4033        let branch = table.create_branch("exp", main_version).await.unwrap();
4034        assert_eq!(branch.current_branch().as_deref(), Some("exp"));
4035        assert_eq!(branch.count_rows(None).await.unwrap(), 1);
4036
4037        // writes on the branch are isolated from main
4038        branch.add(some_sample_data()).execute().await.unwrap();
4039        assert_eq!(branch.count_rows(None).await.unwrap(), 2);
4040        assert_eq!(
4041            table.count_rows(None).await.unwrap(),
4042            1,
4043            "main must be untouched by branch writes"
4044        );
4045
4046        // the branch shows up in the listing
4047        let branches = table.list_branches().await.unwrap();
4048        assert!(branches.contains_key("exp"));
4049
4050        // checking out the branch from the main handle sees the branch's latest data
4051        let checked_out = table.checkout_branch("exp", None).await.unwrap();
4052        assert_eq!(checked_out.current_branch().as_deref(), Some("exp"));
4053        assert_eq!(checked_out.count_rows(None).await.unwrap(), 2);
4054
4055        // open_table(...).branch(...) opens directly onto the branch
4056        let opened = conn
4057            .open_table("my_table")
4058            .branch("exp")
4059            .execute()
4060            .await
4061            .unwrap();
4062        assert_eq!(opened.current_branch().as_deref(), Some("exp"));
4063        assert_eq!(opened.count_rows(None).await.unwrap(), 2);
4064
4065        // delete removes it from the listing
4066        table.delete_branch("exp").await.unwrap();
4067        let branches = table.list_branches().await.unwrap();
4068        assert!(!branches.contains_key("exp"));
4069    }
4070
4071    #[tokio::test]
4072    async fn test_branch_version_checkout() {
4073        let tmp_dir = tempdir().unwrap();
4074        let uri = tmp_dir.path().to_str().unwrap();
4075
4076        let conn = ConnectBuilder::new(uri)
4077            .read_consistency_interval(Duration::from_secs(0))
4078            .execute()
4079            .await
4080            .unwrap();
4081
4082        // main: a single fork-point row (i = 0)
4083        let table = conn
4084            .create_table("my_table", sample_rows(vec![0]))
4085            .execute()
4086            .await
4087            .unwrap();
4088        let fork_point = table.version().await.unwrap();
4089
4090        // Fork "exp", then advance exp AND main independently past the fork so
4091        // they diverge while sharing version numbers.
4092        let branch = table.create_branch("exp", fork_point).await.unwrap();
4093        let exp_fork = branch.version().await.unwrap(); // exp's shallow-clone version
4094        branch.add(sample_rows(vec![1])).execute().await.unwrap(); // exp: {0, 1}
4095        let exp_v2 = branch.version().await.unwrap();
4096        branch.add(sample_rows(vec![2])).execute().await.unwrap(); // exp HEAD: {0, 1, 2}
4097
4098        // main's own commit reaches the SAME version number with different data
4099        table
4100            .add(sample_rows(vec![100, 101, 102]))
4101            .execute()
4102            .await
4103            .unwrap(); // main HEAD: {0, 100, 101, 102}
4104        let main_v2 = table.version().await.unwrap();
4105        assert_eq!(
4106            exp_v2, main_v2,
4107            "branch and main must share the version number for this test to mean anything"
4108        );
4109
4110        // Open exp at the shared version. The data must be exp's, not main's:
4111        // count alone cannot prove this (main@v2 differs), so assert provenance
4112        // by content.
4113        let pinned = conn
4114            .open_table("my_table")
4115            .branch("exp")
4116            .version(exp_v2)
4117            .execute()
4118            .await
4119            .unwrap();
4120        assert_eq!(pinned.current_branch().as_deref(), Some("exp"));
4121        // isolated from exp's HEAD (3 rows) and from main@v2 (4 rows)
4122        assert_eq!(pinned.count_rows(None).await.unwrap(), 2);
4123        // exp's post-fork row is visible; main's divergent rows are not
4124        assert_eq!(
4125            pinned.count_rows(Some("i = 1".to_string())).await.unwrap(),
4126            1
4127        );
4128        assert_eq!(
4129            pinned
4130                .count_rows(Some("i = 100".to_string()))
4131                .await
4132                .unwrap(),
4133            0
4134        );
4135
4136        // the same coordinate is reachable directly via checkout_branch(name, version)
4137        let pinned_direct = table.checkout_branch("exp", Some(exp_v2)).await.unwrap();
4138        assert_eq!(pinned_direct.current_branch().as_deref(), Some("exp"));
4139        assert_eq!(pinned_direct.count_rows(None).await.unwrap(), 2);
4140
4141        // the HEADs are unaffected
4142        let head = conn
4143            .open_table("my_table")
4144            .branch("exp")
4145            .execute()
4146            .await
4147            .unwrap();
4148        assert_eq!(head.count_rows(None).await.unwrap(), 3);
4149        assert_eq!(table.count_rows(None).await.unwrap(), 4);
4150
4151        // a pinned version is a detached head: writes are rejected
4152        assert!(pinned.add(sample_rows(vec![9])).execute().await.is_err());
4153
4154        // version-only (no branch) time-travels main itself: its fork-point
4155        // version holds only main's first row, and the shared version number
4156        // resolves to main's data, not the branch's ("opens main at the version")
4157        let old_main = conn
4158            .open_table("my_table")
4159            .version(fork_point)
4160            .execute()
4161            .await
4162            .unwrap();
4163        assert_eq!(old_main.current_branch(), None);
4164        assert_eq!(old_main.count_rows(None).await.unwrap(), 1);
4165        let shared_on_main = conn
4166            .open_table("my_table")
4167            .version(exp_v2)
4168            .execute()
4169            .await
4170            .unwrap();
4171        assert_eq!(shared_on_main.current_branch(), None);
4172        assert_eq!(shared_on_main.count_rows(None).await.unwrap(), 4);
4173
4174        // a nonexistent version is rejected
4175        assert!(
4176            conn.open_table("my_table")
4177                .version(9999)
4178                .execute()
4179                .await
4180                .is_err()
4181        );
4182
4183        // a nonexistent version on a branch is rejected too: this resolves on
4184        // the branch's path, a distinct miss from the main lookup above
4185        assert!(
4186            conn.open_table("my_table")
4187                .branch("exp")
4188                .version(9999)
4189                .execute()
4190                .await
4191                .is_err()
4192        );
4193
4194        // opening the branch at its fork point (the shallow-clone manifest)
4195        // shows just the cloned state: main's fork-point row
4196        let exp_at_fork = conn
4197            .open_table("my_table")
4198            .branch("exp")
4199            .version(exp_fork)
4200            .execute()
4201            .await
4202            .unwrap();
4203        assert_eq!(exp_at_fork.current_branch().as_deref(), Some("exp"));
4204        assert_eq!(exp_at_fork.count_rows(None).await.unwrap(), 1);
4205
4206        // checkout_latest re-attaches the pinned handle to the BRANCH's HEAD
4207        // (writable again), not main's HEAD, and not staying pinned
4208        pinned.checkout_latest().await.unwrap();
4209        assert_eq!(pinned.current_branch().as_deref(), Some("exp"));
4210        assert_eq!(pinned.count_rows(None).await.unwrap(), 3); // exp HEAD, not main's 4
4211        pinned.add(sample_rows(vec![3])).execute().await.unwrap();
4212        assert_eq!(pinned.count_rows(None).await.unwrap(), 4); // writable again
4213    }
4214
4215    #[tokio::test]
4216    async fn test_branch_version_two_branches() {
4217        let tmp_dir = tempdir().unwrap();
4218        let uri = tmp_dir.path().to_str().unwrap();
4219        let conn = ConnectBuilder::new(uri)
4220            .read_consistency_interval(Duration::from_secs(0))
4221            .execute()
4222            .await
4223            .unwrap();
4224
4225        let table = conn
4226            .create_table("my_table", sample_rows(vec![0]))
4227            .execute()
4228            .await
4229            .unwrap();
4230        let fork_point = table.version().await.unwrap();
4231
4232        // two branches off the same point, each advanced once so they reach the
4233        // SAME version number with divergent data
4234        let exp1 = table.create_branch("exp1", fork_point).await.unwrap();
4235        let exp2 = table.create_branch("exp2", fork_point).await.unwrap();
4236        exp1.add(sample_rows(vec![10])).execute().await.unwrap();
4237        exp2.add(sample_rows(vec![20])).execute().await.unwrap();
4238        let v1 = exp1.version().await.unwrap();
4239        let v2 = exp2.version().await.unwrap();
4240        assert_eq!(v1, v2, "both branches must reach the same version number");
4241
4242        // that shared version number resolves to each branch's own data
4243        let at1 = table.checkout_branch("exp1", Some(v1)).await.unwrap();
4244        assert_eq!(at1.count_rows(Some("i = 10".to_string())).await.unwrap(), 1);
4245        assert_eq!(at1.count_rows(Some("i = 20".to_string())).await.unwrap(), 0);
4246        let at2 = table.checkout_branch("exp2", Some(v2)).await.unwrap();
4247        assert_eq!(at2.count_rows(Some("i = 20".to_string())).await.unwrap(), 1);
4248        assert_eq!(at2.count_rows(Some("i = 10".to_string())).await.unwrap(), 0);
4249    }
4250
4251    #[tokio::test]
4252    async fn test_branch_name_validation() {
4253        let tmp_dir = tempdir().unwrap();
4254        let uri = tmp_dir.path().to_str().unwrap();
4255        let conn = ConnectBuilder::new(uri).execute().await.unwrap();
4256        let table = conn
4257            .create_table("my_table", some_sample_data())
4258            .execute()
4259            .await
4260            .unwrap();
4261
4262        // every entry point rejects an empty name instead of passing it down
4263        assert!(matches!(
4264            table.create_branch("", 1u64).await,
4265            Err(Error::InvalidInput { .. })
4266        ));
4267        assert!(matches!(
4268            table.checkout_branch("", None).await,
4269            Err(Error::InvalidInput { .. })
4270        ));
4271        assert!(matches!(
4272            table.delete_branch("").await,
4273            Err(Error::InvalidInput { .. })
4274        ));
4275        // an empty source branch is rejected too
4276        assert!(matches!(
4277            table
4278                .create_branch(
4279                    "ok",
4280                    lance::dataset::refs::Ref::Version(Some(String::new()), None)
4281                )
4282                .await,
4283            Err(Error::InvalidInput { .. })
4284        ));
4285    }
4286
4287    #[tokio::test]
4288    async fn test_branch_handle_tracks_concurrent_writes() {
4289        let tmp_dir = tempdir().unwrap();
4290        let uri = tmp_dir.path().to_str().unwrap();
4291
4292        // interval = 0 so every read checks storage for new commits
4293        let conn = ConnectBuilder::new(uri)
4294            .read_consistency_interval(Duration::from_secs(0))
4295            .execute()
4296            .await
4297            .unwrap();
4298        let table = conn
4299            .create_table("my_table", some_sample_data())
4300            .execute()
4301            .await
4302            .unwrap();
4303        let v1 = table.version().await.unwrap();
4304
4305        // two independent handles on the same branch
4306        let writer = table.create_branch("exp", v1).await.unwrap();
4307        let reader = conn
4308            .open_table("my_table")
4309            .branch("exp")
4310            .execute()
4311            .await
4312            .unwrap();
4313        assert_eq!(reader.count_rows(None).await.unwrap(), 1);
4314
4315        // a concurrent write on the branch is visible to the other handle, which
4316        // tracks the branch's HEAD (not main's)
4317        writer.add(some_sample_data()).execute().await.unwrap();
4318        assert_eq!(reader.count_rows(None).await.unwrap(), 2);
4319        // main is untouched
4320        assert_eq!(table.count_rows(None).await.unwrap(), 1);
4321    }
4322
4323    #[tokio::test]
4324    async fn test_branch_handle_without_consistency_interval_is_pinned() {
4325        let tmp_dir = tempdir().unwrap();
4326        let uri = tmp_dir.path().to_str().unwrap();
4327
4328        // default interval (None): handles do not auto-refresh
4329        let conn = ConnectBuilder::new(uri).execute().await.unwrap();
4330        let table = conn
4331            .create_table("my_table", some_sample_data())
4332            .execute()
4333            .await
4334            .unwrap();
4335        let v1 = table.version().await.unwrap();
4336
4337        let writer = table.create_branch("exp", v1).await.unwrap();
4338        let reader = conn
4339            .open_table("my_table")
4340            .branch("exp")
4341            .execute()
4342            .await
4343            .unwrap();
4344        assert_eq!(reader.count_rows(None).await.unwrap(), 1);
4345
4346        // without a consistency interval the reader stays on the version it
4347        // opened, exactly like a main-branch handle...
4348        writer.add(some_sample_data()).execute().await.unwrap();
4349        assert_eq!(reader.count_rows(None).await.unwrap(), 1);
4350
4351        // ...until it explicitly refreshes
4352        reader.checkout_latest().await.unwrap();
4353        assert_eq!(reader.count_rows(None).await.unwrap(), 2);
4354    }
4355
4356    #[tokio::test]
4357    async fn test_dynamic_select() {
4358        let tc = new_test_connection().await.unwrap();
4359        let db = tc.connection;
4360
4361        let table = db
4362            .create_table("test", some_sample_data())
4363            .execute()
4364            .await
4365            .unwrap();
4366
4367        let query = table.query().select(Select::dynamic(&[("i_alias", "i")]));
4368
4369        let result = query.execute().await;
4370        let batches = result
4371            .expect("should have result")
4372            .try_collect::<Vec<_>>()
4373            .await
4374            .unwrap();
4375
4376        for batch in batches {
4377            assert!(batch.column_by_name("i_alias").is_some());
4378        }
4379    }
4380
4381    fn some_sample_data() -> Box<dyn arrow_array::RecordBatchReader + Send> {
4382        let batch = RecordBatch::try_new(
4383            Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])),
4384            vec![Arc::new(Int32Array::from(vec![1]))],
4385        )
4386        .unwrap();
4387        let schema = batch.schema().clone();
4388        let batch = Ok(batch);
4389
4390        Box::new(RecordBatchIterator::new(vec![batch], schema))
4391    }
4392
4393    /// A single-batch reader holding the given `i` (Int32) values. Lets a test
4394    /// write distinguishable rows so it can assert data provenance, not row count.
4395    fn sample_rows(values: Vec<i32>) -> Box<dyn arrow_array::RecordBatchReader + Send> {
4396        let batch = RecordBatch::try_new(
4397            Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])),
4398            vec![Arc::new(Int32Array::from(values))],
4399        )
4400        .unwrap();
4401        let schema = batch.schema().clone();
4402
4403        Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema))
4404    }
4405
4406    #[tokio::test]
4407    async fn test_read_consistency_interval() {
4408        use crate::utils::background_cache::clock;
4409
4410        let intervals = vec![
4411            None,
4412            Some(0),
4413            Some(100), // 100 ms
4414        ];
4415
4416        for interval in intervals {
4417            let data = some_sample_data();
4418
4419            let tmp_dir = tempdir().unwrap();
4420            let uri = tmp_dir.path().to_str().unwrap();
4421
4422            let conn1 = ConnectBuilder::new(uri).execute().await.unwrap();
4423            let table1 = conn1
4424                .create_empty_table("my_table", RecordBatchReader::schema(&data))
4425                .execute()
4426                .await
4427                .unwrap();
4428
4429            let mut conn2 = ConnectBuilder::new(uri);
4430            if let Some(interval) = interval {
4431                conn2 = conn2.read_consistency_interval(std::time::Duration::from_millis(interval));
4432            }
4433            let conn2 = conn2.execute().await.unwrap();
4434            let table2 = conn2.open_table("my_table").execute().await.unwrap();
4435
4436            // Freeze the consistency clock now that `table2` has seeded its cache, so the
4437            // interval only elapses when this test advances it. Otherwise the write and
4438            // count_rows calls below race the real 100ms interval, which a loaded CI
4439            // runner loses. Must come after open_table: creating the cache clears the mock.
4440            clock::pin();
4441
4442            assert_eq!(table1.count_rows(None).await.unwrap(), 0);
4443            assert_eq!(table2.count_rows(None).await.unwrap(), 0);
4444
4445            table1.add(data).execute().await.unwrap();
4446            assert_eq!(table1.count_rows(None).await.unwrap(), 1);
4447
4448            match interval {
4449                None => {
4450                    assert_eq!(table2.count_rows(None).await.unwrap(), 0);
4451                    table2.checkout_latest().await.unwrap();
4452                    assert_eq!(table2.count_rows(None).await.unwrap(), 1);
4453                }
4454                Some(0) => {
4455                    assert_eq!(table2.count_rows(None).await.unwrap(), 1);
4456                }
4457                Some(100) => {
4458                    assert_eq!(table2.count_rows(None).await.unwrap(), 0);
4459                    clock::advance_by(Duration::from_millis(100));
4460                    assert_eq!(table2.count_rows(None).await.unwrap(), 1);
4461                }
4462                _ => unreachable!(),
4463            }
4464        }
4465    }
4466
4467    #[tokio::test]
4468    async fn test_time_travel_write() {
4469        let tmp_dir = tempdir().unwrap();
4470        let uri = tmp_dir.path().to_str().unwrap();
4471
4472        let conn = ConnectBuilder::new(uri)
4473            .read_consistency_interval(Duration::from_secs(0))
4474            .execute()
4475            .await
4476            .unwrap();
4477        let table = conn
4478            .create_table("my_table", some_sample_data())
4479            .execute()
4480            .await
4481            .unwrap();
4482        let version = table.version().await.unwrap();
4483        table.add(some_sample_data()).execute().await.unwrap();
4484        table.checkout(version).await.unwrap();
4485        assert!(table.add(some_sample_data()).execute().await.is_err())
4486    }
4487
4488    #[tokio::test]
4489    async fn test_update_dataset_config() {
4490        let tmp_dir = tempdir().unwrap();
4491        let uri = tmp_dir.path().to_str().unwrap();
4492
4493        let conn = ConnectBuilder::new(uri)
4494            .read_consistency_interval(Duration::from_secs(0))
4495            .execute()
4496            .await
4497            .unwrap();
4498
4499        let table = conn
4500            .create_table("my_table", some_sample_data())
4501            .execute()
4502            .await
4503            .unwrap();
4504        let native_tbl = table.as_native().unwrap();
4505
4506        let manifest = native_tbl.manifest().await.unwrap();
4507        let base_config_len = manifest.config.len();
4508
4509        native_tbl
4510            .update_config(vec![("test_key1".to_string(), "test_val1".to_string())])
4511            .await
4512            .unwrap();
4513
4514        let manifest = native_tbl.manifest().await.unwrap();
4515        assert_eq!(manifest.config.len(), 1 + base_config_len);
4516        assert_eq!(
4517            manifest.config.get("test_key1"),
4518            Some(&"test_val1".to_string())
4519        );
4520
4521        native_tbl
4522            .update_config(vec![("test_key2".to_string(), "test_val2".to_string())])
4523            .await
4524            .unwrap();
4525        let manifest = native_tbl.manifest().await.unwrap();
4526        assert_eq!(manifest.config.len(), 2 + base_config_len);
4527        assert_eq!(
4528            manifest.config.get("test_key1"),
4529            Some(&"test_val1".to_string())
4530        );
4531        assert_eq!(
4532            manifest.config.get("test_key2"),
4533            Some(&"test_val2".to_string())
4534        );
4535
4536        native_tbl
4537            .update_config(vec![(
4538                "test_key2".to_string(),
4539                "test_val2_update".to_string(),
4540            )])
4541            .await
4542            .unwrap();
4543        let manifest = native_tbl.manifest().await.unwrap();
4544        assert_eq!(manifest.config.len(), 2 + base_config_len);
4545        assert_eq!(
4546            manifest.config.get("test_key1"),
4547            Some(&"test_val1".to_string())
4548        );
4549        assert_eq!(
4550            manifest.config.get("test_key2"),
4551            Some(&"test_val2_update".to_string())
4552        );
4553
4554        native_tbl.delete_config_keys(&["test_key1"]).await.unwrap();
4555        let manifest = native_tbl.manifest().await.unwrap();
4556        assert_eq!(manifest.config.len(), 1 + base_config_len);
4557        assert_eq!(
4558            manifest.config.get("test_key2"),
4559            Some(&"test_val2_update".to_string())
4560        );
4561    }
4562
4563    #[tokio::test]
4564    async fn test_schema_metadata_config() {
4565        let tmp_dir = tempdir().unwrap();
4566        let uri = tmp_dir.path().to_str().unwrap();
4567
4568        let conn = ConnectBuilder::new(uri)
4569            .read_consistency_interval(Duration::from_secs(0))
4570            .execute()
4571            .await
4572            .unwrap();
4573        let table = conn
4574            .create_table("my_table", some_sample_data())
4575            .execute()
4576            .await
4577            .unwrap();
4578
4579        let native_tbl = table.as_native().unwrap();
4580        let schema = native_tbl.schema().await.unwrap();
4581        let metadata = schema.metadata();
4582        assert_eq!(metadata.len(), 0);
4583
4584        native_tbl
4585            .replace_schema_metadata(vec![("test_key1".to_string(), "test_val1".to_string())])
4586            .await
4587            .unwrap();
4588
4589        let schema = native_tbl.schema().await.unwrap();
4590        let metadata = schema.metadata();
4591        assert_eq!(metadata.len(), 1);
4592        assert_eq!(metadata.get("test_key1"), Some(&"test_val1".to_string()));
4593
4594        native_tbl
4595            .replace_schema_metadata(vec![
4596                ("test_key1".to_string(), "test_val1_update".to_string()),
4597                ("test_key2".to_string(), "test_val2".to_string()),
4598            ])
4599            .await
4600            .unwrap();
4601        let schema = native_tbl.schema().await.unwrap();
4602        let metadata = schema.metadata();
4603        assert_eq!(metadata.len(), 2);
4604        assert_eq!(
4605            metadata.get("test_key1"),
4606            Some(&"test_val1_update".to_string())
4607        );
4608        assert_eq!(metadata.get("test_key2"), Some(&"test_val2".to_string()));
4609
4610        native_tbl
4611            .replace_schema_metadata(vec![(
4612                "test_key2".to_string(),
4613                "test_val2_update".to_string(),
4614            )])
4615            .await
4616            .unwrap();
4617        let schema = native_tbl.schema().await.unwrap();
4618        let metadata = schema.metadata();
4619        assert_eq!(
4620            metadata.get("test_key2"),
4621            Some(&"test_val2_update".to_string())
4622        );
4623    }
4624
4625    #[tokio::test]
4626    pub async fn test_field_metadata_update() {
4627        let tmp_dir = tempdir().unwrap();
4628        let uri = tmp_dir.path().to_str().unwrap();
4629
4630        let conn = ConnectBuilder::new(uri)
4631            .read_consistency_interval(Duration::from_secs(0))
4632            .execute()
4633            .await
4634            .unwrap();
4635        let table = conn
4636            .create_table("my_table", some_sample_data())
4637            .execute()
4638            .await
4639            .unwrap();
4640
4641        let native_tbl = table.as_native().unwrap();
4642        let schema = native_tbl.manifest().await.unwrap().schema;
4643
4644        let field = schema.field("i").unwrap();
4645        assert_eq!(field.metadata.len(), 0);
4646
4647        native_tbl
4648            .replace_schema_metadata(vec![(
4649                "test_key2".to_string(),
4650                "test_val2_update".to_string(),
4651            )])
4652            .await
4653            .unwrap();
4654
4655        let schema = native_tbl.schema().await.unwrap();
4656        let metadata = schema.metadata();
4657        assert_eq!(metadata.len(), 1);
4658        assert_eq!(
4659            metadata.get("test_key2"),
4660            Some(&"test_val2_update".to_string())
4661        );
4662
4663        native_tbl
4664            .update_field_metadata(&[
4665                FieldMetadataUpdate::new("i").set("test_field_key1", "test_field_val1")
4666            ])
4667            .await
4668            .unwrap();
4669
4670        let schema = native_tbl.manifest().await.unwrap().schema;
4671        let field = schema.field("i").unwrap();
4672        assert_eq!(field.metadata.len(), 1);
4673        assert_eq!(
4674            field.metadata.get("test_field_key1"),
4675            Some(&"test_field_val1".to_string())
4676        );
4677    }
4678
4679    #[tokio::test]
4680    async fn test_set_unenforced_primary_key() {
4681        let tmp_dir = tempdir().unwrap();
4682        let uri = tmp_dir.path().to_str().unwrap();
4683
4684        let schema = Arc::new(Schema::new(vec![
4685            Field::new("id", DataType::Int64, false),
4686            Field::new("name", DataType::Utf8, true),
4687            Field::new("score", DataType::Float64, true),
4688        ]));
4689        let batch = RecordBatch::try_new(
4690            schema.clone(),
4691            vec![
4692                Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
4693                Arc::new(StringArray::from(vec!["a", "b", "c"])),
4694                Arc::new(arrow_array::Float64Array::from(vec![1.0, 2.0, 3.0])),
4695            ],
4696        )
4697        .unwrap();
4698        let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4699            Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4700
4701        let conn = ConnectBuilder::new(uri)
4702            .read_consistency_interval(Duration::from_secs(0))
4703            .execute()
4704            .await
4705            .unwrap();
4706        let table = conn.create_table("t", reader).execute().await.unwrap();
4707
4708        // Reject empty input.
4709        let err = table
4710            .set_unenforced_primary_key(Vec::<&str>::new())
4711            .await
4712            .expect_err("empty input should be rejected");
4713        assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4714
4715        // Reject compound (multi-column) input.
4716        let err = table
4717            .set_unenforced_primary_key(["id", "name"])
4718            .await
4719            .expect_err("compound primary key should be rejected");
4720        assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4721
4722        // Reject unknown column.
4723        let err = table
4724            .set_unenforced_primary_key(["nonexistent"])
4725            .await
4726            .expect_err("nonexistent column should be rejected");
4727        assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4728
4729        // Reject unsupported dtype (Float64).
4730        let err = table
4731            .set_unenforced_primary_key(["score"])
4732            .await
4733            .expect_err("Float64 should be rejected");
4734        assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4735
4736        // None of the rejected calls set a primary key.
4737        let lance_schema = table.as_native().unwrap().manifest().await.unwrap().schema;
4738        assert!(lance_schema.unenforced_primary_key().is_empty());
4739
4740        // Happy path: set the primary key to "id".
4741        table.set_unenforced_primary_key(["id"]).await.unwrap();
4742        let lance_schema = table.as_native().unwrap().manifest().await.unwrap().schema;
4743        let pk = lance_schema.unenforced_primary_key();
4744        assert_eq!(pk.len(), 1);
4745        assert_eq!(pk[0].name, "id");
4746        // Position metadata is 1-indexed.
4747        assert_eq!(
4748            pk[0].metadata.get(LANCE_UNENFORCED_PRIMARY_KEY_POSITION),
4749            Some(&"1".to_string())
4750        );
4751
4752        // The primary key is immutable: re-setting it is rejected, whether to
4753        // the same column or a different one.
4754        let err = table
4755            .set_unenforced_primary_key(["id"])
4756            .await
4757            .expect_err("re-setting the same primary key should be rejected");
4758        assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4759        let err = table
4760            .set_unenforced_primary_key(["name"])
4761            .await
4762            .expect_err("changing the primary key should be rejected");
4763        assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4764
4765        // The primary key is unchanged after the rejected calls.
4766        let lance_schema = table.as_native().unwrap().manifest().await.unwrap().schema;
4767        let pk = lance_schema.unenforced_primary_key();
4768        assert_eq!(pk.len(), 1);
4769        assert_eq!(pk[0].name, "id");
4770    }
4771
4772    #[tokio::test]
4773    async fn test_set_unenforced_primary_key_concurrent() {
4774        let tmp_dir = tempdir().unwrap();
4775        let uri = tmp_dir.path().to_str().unwrap();
4776
4777        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
4778        let batch = RecordBatch::try_new(
4779            schema.clone(),
4780            vec![Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3]))],
4781        )
4782        .unwrap();
4783        let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4784            Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4785
4786        // A long read-consistency interval keeps each handle pinned to the
4787        // version it opened, so the second handle commits against a stale
4788        // base — the same situation as two processes racing.
4789        let conn = ConnectBuilder::new(uri)
4790            .read_consistency_interval(Duration::from_secs(3600))
4791            .execute()
4792            .await
4793            .unwrap();
4794        conn.create_table("t", reader).execute().await.unwrap();
4795
4796        let table_a = conn.open_table("t").execute().await.unwrap();
4797        let table_b = conn.open_table("t").execute().await.unwrap();
4798
4799        // Handle A sets the primary key first.
4800        table_a.set_unenforced_primary_key(["id"]).await.unwrap();
4801
4802        // Handle B committed against a stale base that had no primary key, so
4803        // its own up-front check did not see A's key. The commit itself must
4804        // still fail rather than silently overriding A's primary key. (The
4805        // cross-process race on a *different* column is caught by the Lance
4806        // commit layer.)
4807        let err = table_b
4808            .set_unenforced_primary_key(["id"])
4809            .await
4810            .expect_err("concurrent primary key commit on a stale base should fail");
4811        assert!(
4812            !matches!(err, Error::InvalidInput { .. }),
4813            "expected a commit-time conflict, not an up-front input error: {:?}",
4814            err
4815        );
4816
4817        // The committed primary key is exactly what A set — no corruption.
4818        let fresh = conn.open_table("t").execute().await.unwrap();
4819        let lance_schema = fresh.as_native().unwrap().manifest().await.unwrap().schema;
4820        let pk = lance_schema.unenforced_primary_key();
4821        assert_eq!(pk.len(), 1);
4822        assert_eq!(pk[0].name, "id");
4823    }
4824
4825    #[tokio::test]
4826    async fn test_set_lsm_write_spec() {
4827        use arrow_array::StringArray;
4828        use lance::dataset::mem_wal::DatasetMemWalExt;
4829
4830        let tmp_dir = tempdir().unwrap();
4831        let uri = tmp_dir.path().to_str().unwrap();
4832
4833        let schema = Arc::new(Schema::new(vec![
4834            Field::new("id", DataType::Int64, false),
4835            Field::new("name", DataType::Utf8, true),
4836        ]));
4837        let batch = RecordBatch::try_new(
4838            schema.clone(),
4839            vec![
4840                Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
4841                Arc::new(StringArray::from(vec!["a", "b", "c"])),
4842            ],
4843        )
4844        .unwrap();
4845        let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4846            Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4847
4848        let conn = ConnectBuilder::new(uri)
4849            .read_consistency_interval(Duration::from_secs(0))
4850            .execute()
4851            .await
4852            .unwrap();
4853        let table = conn.create_table("t", reader).execute().await.unwrap();
4854
4855        // Reject num_buckets out of range.
4856        for bad in [0u32, 1025] {
4857            let err = table
4858                .set_lsm_write_spec(LsmWriteSpec::bucket("id", bad))
4859                .await
4860                .expect_err("should reject");
4861            assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4862        }
4863
4864        // Happy path: install spec; verify MemWAL details record it.
4865        table
4866            .set_lsm_write_spec(LsmWriteSpec::bucket("id", 4))
4867            .await
4868            .unwrap();
4869
4870        let native_tbl = table.as_native().unwrap();
4871        let dataset = native_tbl.dataset.get().await.unwrap();
4872        let details = dataset
4873            .mem_wal_index_details()
4874            .await
4875            .unwrap()
4876            .expect("MemWAL index should be initialized");
4877        assert_eq!(details.num_shards, 4);
4878        assert_eq!(details.sharding_specs.len(), 1);
4879        let installed = &details.sharding_specs[0];
4880        assert_eq!(installed.fields.len(), 1);
4881        let f = &installed.fields[0];
4882        assert_eq!(f.transform.as_deref(), Some("bucket"));
4883        assert_eq!(
4884            f.parameters.get("num_buckets").map(String::as_str),
4885            Some("4")
4886        );
4887        // Bucket parameters must hold only `num_buckets`.
4888        assert_eq!(f.parameters.len(), 1);
4889
4890        // Mutation rejected.
4891        let err = table
4892            .set_lsm_write_spec(LsmWriteSpec::bucket("id", 8))
4893            .await
4894            .expect_err("mutation should be rejected");
4895        assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4896    }
4897
4898    #[tokio::test]
4899    async fn test_set_lsm_write_spec_unsharded() {
4900        use lance::dataset::mem_wal::DatasetMemWalExt;
4901
4902        let tmp_dir = tempdir().unwrap();
4903        let uri = tmp_dir.path().to_str().unwrap();
4904
4905        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
4906        let batch = RecordBatch::try_new(
4907            schema.clone(),
4908            vec![Arc::new(arrow_array::Int64Array::from(vec![1]))],
4909        )
4910        .unwrap();
4911        let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4912            Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4913        let conn = ConnectBuilder::new(uri)
4914            .read_consistency_interval(Duration::from_secs(0))
4915            .execute()
4916            .await
4917            .unwrap();
4918        let table = conn.create_table("t", reader).execute().await.unwrap();
4919
4920        table
4921            .set_lsm_write_spec(LsmWriteSpec::unsharded())
4922            .await
4923            .unwrap();
4924
4925        let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
4926        let details = dataset
4927            .mem_wal_index_details()
4928            .await
4929            .unwrap()
4930            .expect("MemWAL index should be initialized");
4931        assert_eq!(details.num_shards, 1);
4932        assert_eq!(details.sharding_specs.len(), 1);
4933        let f = &details.sharding_specs[0].fields[0];
4934        assert_eq!(f.transform.as_deref(), Some("unsharded"));
4935        assert!(f.source_ids.is_empty());
4936    }
4937
4938    #[tokio::test]
4939    async fn test_set_lsm_write_spec_identity() {
4940        use lance::dataset::mem_wal::DatasetMemWalExt;
4941
4942        let tmp_dir = tempdir().unwrap();
4943        let uri = tmp_dir.path().to_str().unwrap();
4944
4945        let schema = Arc::new(Schema::new(vec![
4946            Field::new("id", DataType::Int64, false),
4947            Field::new("region", DataType::Utf8, true),
4948        ]));
4949        let batch = RecordBatch::try_new(
4950            schema.clone(),
4951            vec![
4952                Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
4953                Arc::new(StringArray::from(vec!["a", "b", "c"])),
4954            ],
4955        )
4956        .unwrap();
4957        let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4958            Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4959        let conn = ConnectBuilder::new(uri)
4960            .read_consistency_interval(Duration::from_secs(0))
4961            .execute()
4962            .await
4963            .unwrap();
4964        let table = conn.create_table("t", reader).execute().await.unwrap();
4965
4966        table
4967            .set_lsm_write_spec(
4968                LsmWriteSpec::identity("region")
4969                    .with_writer_config_defaults([("durable_write", "false")]),
4970            )
4971            .await
4972            .unwrap();
4973
4974        let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
4975        let details = dataset
4976            .mem_wal_index_details()
4977            .await
4978            .unwrap()
4979            .expect("MemWAL index should be initialized");
4980        // Identity sharding records an open-ended shard count.
4981        assert_eq!(details.num_shards, 0);
4982        assert_eq!(details.sharding_specs.len(), 1);
4983        let f = &details.sharding_specs[0].fields[0];
4984        assert_eq!(f.transform.as_deref(), Some("identity"));
4985        // Writer config defaults round-trip into the MemWAL index.
4986        assert_eq!(
4987            details
4988                .writer_config_defaults
4989                .get("durable_write")
4990                .map(String::as_str),
4991            Some("false")
4992        );
4993    }
4994
4995    #[tokio::test]
4996    async fn test_unset_lsm_write_spec() {
4997        use lance::dataset::mem_wal::DatasetMemWalExt;
4998
4999        let tmp_dir = tempdir().unwrap();
5000        let uri = tmp_dir.path().to_str().unwrap();
5001
5002        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
5003        let batch = RecordBatch::try_new(
5004            schema.clone(),
5005            vec![Arc::new(arrow_array::Int64Array::from(vec![1]))],
5006        )
5007        .unwrap();
5008        let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
5009            Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
5010        let conn = ConnectBuilder::new(uri)
5011            .read_consistency_interval(Duration::from_secs(0))
5012            .execute()
5013            .await
5014            .unwrap();
5015        let table = conn.create_table("t", reader).execute().await.unwrap();
5016
5017        // unset errors when no spec is set.
5018        table.unset_lsm_write_spec().await.unwrap_err();
5019
5020        // Install a spec, then unset it.
5021        table
5022            .set_lsm_write_spec(LsmWriteSpec::bucket("id", 4))
5023            .await
5024            .unwrap();
5025        {
5026            let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
5027            assert!(dataset.mem_wal_index_details().await.unwrap().is_some());
5028        }
5029
5030        table.unset_lsm_write_spec().await.unwrap();
5031        {
5032            let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
5033            assert!(dataset.mem_wal_index_details().await.unwrap().is_none());
5034        }
5035
5036        // A second unset errors; a fresh spec can still be installed afterwards.
5037        table.unset_lsm_write_spec().await.unwrap_err();
5038        table
5039            .set_lsm_write_spec(LsmWriteSpec::bucket("id", 8))
5040            .await
5041            .unwrap();
5042        {
5043            let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
5044            assert!(dataset.mem_wal_index_details().await.unwrap().is_some());
5045        }
5046    }
5047
5048    #[tokio::test]
5049    async fn test_get_lsm_write_spec() {
5050        let tmp_dir = tempdir().unwrap();
5051        let uri = tmp_dir.path().to_str().unwrap();
5052
5053        let schema = Arc::new(Schema::new(vec![
5054            Field::new("id", DataType::Int64, false),
5055            Field::new("region", DataType::Utf8, true),
5056        ]));
5057        let batch = RecordBatch::try_new(
5058            schema.clone(),
5059            vec![
5060                Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
5061                Arc::new(StringArray::from(vec!["a", "b", "c"])),
5062            ],
5063        )
5064        .unwrap();
5065        let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
5066            Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
5067        let conn = ConnectBuilder::new(uri)
5068            .read_consistency_interval(Duration::from_secs(0))
5069            .execute()
5070            .await
5071            .unwrap();
5072        let table = conn.create_table("t", reader).execute().await.unwrap();
5073
5074        // No spec installed yet.
5075        assert_eq!(table.get_lsm_write_spec().await.unwrap(), None);
5076
5077        // A real scalar index is needed to name it as a maintained index.
5078        table
5079            .create_index(&["id"], Index::Auto)
5080            .execute()
5081            .await
5082            .unwrap();
5083        let idx_name = table.list_indices().await.unwrap()[0].name.clone();
5084
5085        // Bucket spec round-trips exactly, including the routing column (recovered
5086        // from its field id), maintained indexes, and writer config defaults.
5087        let spec = LsmWriteSpec::bucket("id", 4)
5088            .with_maintained_indexes([idx_name])
5089            .with_writer_config_defaults([("durable_write", "false")]);
5090        table.set_lsm_write_spec(spec.clone()).await.unwrap();
5091        assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
5092
5093        // After unset, no spec is reported.
5094        table.unset_lsm_write_spec().await.unwrap();
5095        assert_eq!(table.get_lsm_write_spec().await.unwrap(), None);
5096
5097        // Identity sharding round-trips (column recovered from the schema).
5098        let spec = LsmWriteSpec::identity("region");
5099        table.set_lsm_write_spec(spec.clone()).await.unwrap();
5100        assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
5101        table.unset_lsm_write_spec().await.unwrap();
5102
5103        // Unsharded round-trips (no routing column).
5104        let spec = LsmWriteSpec::unsharded();
5105        table.set_lsm_write_spec(spec.clone()).await.unwrap();
5106        assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
5107    }
5108
5109    #[tokio::test]
5110    pub async fn test_stats() {
5111        let tmp_dir = tempdir().unwrap();
5112        let uri = tmp_dir.path().to_str().unwrap();
5113
5114        let conn = ConnectBuilder::new(uri).execute().await.unwrap();
5115
5116        let schema = Arc::new(Schema::new(vec![
5117            Field::new("id", DataType::Int32, false),
5118            Field::new("foo", DataType::Int32, true),
5119        ]));
5120        let batch = RecordBatch::try_new(
5121            schema.clone(),
5122            vec![
5123                Arc::new(Int32Array::from_iter_values(0..100)),
5124                Arc::new(Int32Array::from_iter_values(0..100)),
5125            ],
5126        )
5127        .unwrap();
5128
5129        let table = conn
5130            .create_table("test_stats", batch.clone())
5131            .execute()
5132            .await
5133            .unwrap();
5134        for _ in 0..10 {
5135            let batch = RecordBatch::try_new(
5136                schema.clone(),
5137                vec![
5138                    Arc::new(Int32Array::from_iter_values(0..15)),
5139                    Arc::new(Int32Array::from_iter_values(0..15)),
5140                ],
5141            )
5142            .unwrap();
5143            table.add(batch.clone()).execute().await.unwrap();
5144        }
5145
5146        let empty_table = conn
5147            .create_table("test_stats_empty", RecordBatch::new_empty(batch.schema()))
5148            .execute()
5149            .await
5150            .unwrap();
5151
5152        let res = table.stats().await.unwrap();
5153        println!("{:#?}", res);
5154        // `total_bytes` is the full on-disk size of the 11 data files (this table
5155        // has no index or overlay files), so it is well above the 2000 bytes of
5156        // column data these 250 int32 pairs hold: each file carries its own footer
5157        // and metadata.
5158        assert_eq!(
5159            res,
5160            TableStatistics {
5161                num_rows: 250,
5162                num_indices: 0,
5163                total_bytes: 8925,
5164                fragment_stats: FragmentStatistics {
5165                    num_fragments: 11,
5166                    num_small_fragments: 11,
5167                    lengths: FragmentSummaryStats {
5168                        min: 15,
5169                        max: 100,
5170                        mean: 22,
5171                        p25: 15,
5172                        p50: 15,
5173                        p75: 15,
5174                        p99: 100,
5175                    },
5176                },
5177            }
5178        );
5179        let res = empty_table.stats().await.unwrap();
5180        println!("{:#?}", res);
5181        assert_eq!(
5182            res,
5183            TableStatistics {
5184                num_rows: 0,
5185                num_indices: 0,
5186                total_bytes: 0,
5187                fragment_stats: FragmentStatistics {
5188                    num_fragments: 0,
5189                    num_small_fragments: 0,
5190                    lengths: FragmentSummaryStats {
5191                        min: 0,
5192                        max: 0,
5193                        mean: 0,
5194                        p25: 0,
5195                        p50: 0,
5196                        p75: 0,
5197                        p99: 0,
5198                    },
5199                },
5200            }
5201        )
5202    }
5203
5204    /// `total_bytes` counts more than the base data files: index files and
5205    /// overlay files recorded in the manifest are included too.
5206    #[tokio::test]
5207    pub async fn test_stats_includes_index_and_overlay_files() {
5208        use lance::dataset::WriteDestination;
5209        use lance::dataset::transaction::{DataOverlayGroup, Operation};
5210        use lance_file::version::{ConcreteFileVersion, LanceFileVersion};
5211        use lance_file::writer::{FileWriter, FileWriterOptions};
5212        use lance_io::utils::CachedFileSize;
5213        use lance_table::format::DataFile;
5214        use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage};
5215        use roaring::RoaringBitmap;
5216
5217        let tmp_dir = tempdir().unwrap();
5218        let uri = tmp_dir.path().to_str().unwrap();
5219        let conn = ConnectBuilder::new(uri)
5220            .read_consistency_interval(Duration::from_secs(0))
5221            .execute()
5222            .await
5223            .unwrap();
5224
5225        let schema = Arc::new(Schema::new(vec![
5226            Field::new("id", DataType::Int32, false),
5227            Field::new("foo", DataType::Int32, true),
5228        ]));
5229        let batch = RecordBatch::try_new(
5230            schema.clone(),
5231            vec![
5232                Arc::new(Int32Array::from_iter_values(0..100)),
5233                Arc::new(Int32Array::from_iter_values(0..100)),
5234            ],
5235        )
5236        .unwrap();
5237        let table = conn
5238            .create_table("test_stats_extra_files", batch)
5239            .execute()
5240            .await
5241            .unwrap();
5242
5243        let data_only = table.stats().await.unwrap().total_bytes;
5244        assert!(data_only > 0);
5245
5246        // A scalar index adds index files whose sizes are recorded in the
5247        // manifest's index section.
5248        table
5249            .create_index(&["id"], Index::Auto)
5250            .execute()
5251            .await
5252            .unwrap();
5253        let with_index = table.stats().await.unwrap().total_bytes;
5254        let dataset = {
5255            let native = table.as_native().unwrap();
5256            (*native.dataset.get().await.unwrap()).clone()
5257        };
5258        let index_bytes: usize = dataset
5259            .load_indices()
5260            .await
5261            .unwrap()
5262            .iter()
5263            .map(|idx| idx.total_size_bytes().unwrap_or(0) as usize)
5264            .sum();
5265        assert!(index_bytes > 0);
5266        assert_eq!(with_index, data_only + index_bytes);
5267
5268        // Commit an overlay file supplying new `foo` values for the first three
5269        // rows of fragment 0. There is no high-level API that writes overlays
5270        // yet, so write the overlay's data file and commit the `DataOverlay`
5271        // operation by hand.
5272        let read_version = dataset.version().version;
5273        let fragment_id = dataset.get_fragments()[0].id() as u64;
5274        let foo_field_id = dataset.schema().field("foo").unwrap().id;
5275        let overlay_schema = dataset.schema().project_by_ids(&[foo_field_id], true);
5276        let file_version = ConcreteFileVersion::from(LanceFileVersion::Stable);
5277
5278        let filename = "overlay.lance".to_string();
5279        let store = dataset.object_store(None).await.unwrap();
5280        let path = dataset.data_dir().child(filename.clone());
5281        let obj_writer = store.create(&path).await.unwrap();
5282        let mut writer = FileWriter::try_new(
5283            obj_writer,
5284            overlay_schema,
5285            FileWriterOptions {
5286                format_version: Some(file_version.into()),
5287                ..Default::default()
5288            },
5289        )
5290        .unwrap();
5291        writer
5292            .write_column(0, Arc::new(Int32Array::from(vec![1000, 1001, 1002])) as _)
5293            .await
5294            .unwrap();
5295        let summary = writer.finish().await.unwrap();
5296        let overlay_bytes = summary.size_bytes as usize;
5297        assert!(overlay_bytes > 0);
5298
5299        let mut data_file = DataFile::new_unstarted(filename, file_version);
5300        data_file.fields = writer
5301            .field_id_to_column_indices()
5302            .iter()
5303            .map(|(field_id, _)| *field_id as i32)
5304            .collect::<Vec<_>>()
5305            .into();
5306        data_file.column_indices = writer
5307            .field_id_to_column_indices()
5308            .iter()
5309            .map(|(_, column_index)| *column_index as i32)
5310            .collect::<Vec<_>>()
5311            .into();
5312        data_file.file_size_bytes = CachedFileSize::new(summary.size_bytes);
5313
5314        let overlay = DataOverlayFile {
5315            data_file,
5316            coverage: OverlayCoverage::dense(RoaringBitmap::from_iter(0..3)),
5317            committed_version: 0,
5318        };
5319        Dataset::commit(
5320            WriteDestination::Dataset(Arc::new(dataset)),
5321            Operation::DataOverlay {
5322                groups: vec![DataOverlayGroup {
5323                    fragment_id,
5324                    overlays: vec![overlay],
5325                }],
5326            },
5327            Some(read_version),
5328            None,
5329            None,
5330            Arc::new(Default::default()),
5331            false,
5332        )
5333        .await
5334        .unwrap();
5335
5336        table.checkout_latest().await.unwrap();
5337        let with_overlay = table.stats().await.unwrap().total_bytes;
5338        assert_eq!(with_overlay, with_index + overlay_bytes);
5339    }
5340
5341    /// `stats()` must stay manifest-only. Summing per-field `bytes_on_disk`
5342    /// instead opens every data file, so cost would grow with fragment count.
5343    #[tokio::test]
5344    pub async fn test_stats_does_not_read_data_files() {
5345        let tmp_dir = tempdir().unwrap();
5346        let uri = tmp_dir.path().to_str().unwrap();
5347
5348        let conn = ConnectBuilder::new(uri).execute().await.unwrap();
5349
5350        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
5351        let batch = RecordBatch::try_new(
5352            schema.clone(),
5353            vec![Arc::new(Int32Array::from_iter_values(0..10))],
5354        )
5355        .unwrap();
5356
5357        conn.create_table("test_stats_io", batch.clone())
5358            .execute()
5359            .await
5360            .unwrap();
5361        let table = conn.open_table("test_stats_io").execute().await.unwrap();
5362        const NUM_APPENDS: usize = 20;
5363        for _ in 0..NUM_APPENDS {
5364            table.add(batch.clone()).execute().await.unwrap();
5365        }
5366
5367        // Reopen through a tracking store so the counters cover `stats()` alone and
5368        // not the writes above.
5369        let (wrapper, io_stats) = IoTrackingStore::new_wrapper();
5370        let table = conn
5371            .open_table("test_stats_io")
5372            .lance_read_params(ReadParams {
5373                store_options: Some(ObjectStoreParams {
5374                    object_store_wrapper: Some(wrapper),
5375                    ..Default::default()
5376                }),
5377                ..Default::default()
5378            })
5379            .execute()
5380            .await
5381            .unwrap();
5382        io_stats.lock().unwrap().read_iops = 0;
5383
5384        let stats = table.stats().await.unwrap();
5385        let read_iops = io_stats.lock().unwrap().read_iops;
5386
5387        assert_eq!(stats.fragment_stats.num_fragments, NUM_APPENDS + 1);
5388        assert!(stats.total_bytes > 0);
5389        // Reading the fragments' data files would take at least one IOP each.
5390        assert!(
5391            read_iops < stats.fragment_stats.num_fragments as u64,
5392            "stats() issued {} read IOPs across {} fragments",
5393            read_iops,
5394            stats.fragment_stats.num_fragments
5395        );
5396    }
5397}