Skip to main content

lance/
dataset.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Lance Dataset
5//!
6
7use arrow_array::{RecordBatch, RecordBatchReader};
8use arrow_schema::DataType;
9use byteorder::{ByteOrder, LittleEndian};
10use chrono::{Duration, prelude::*};
11use futures::future::BoxFuture;
12use futures::stream::{self, BoxStream, StreamExt, TryStreamExt};
13use futures::{FutureExt, Stream};
14use lance_core::deepsize::DeepSizeOf;
15
16use crate::dataset::metadata::UpdateFieldMetadataBuilder;
17use crate::dataset::transaction::translate_schema_metadata_updates;
18use crate::index::DatasetIndexExt;
19use crate::session::caches::{DSMetadataCache, ManifestKey, TransactionKey};
20use crate::session::index_caches::DSIndexCache;
21use itertools::Itertools;
22use lance_core::ROW_ADDR;
23use lance_core::datatypes::{OnMissing, OnTypeMismatch, Projectable, Projection};
24use lance_core::traits::DatasetTakeRows;
25use lance_core::utils::address::RowAddress;
26use lance_core::utils::tracing::{
27    DATASET_DELETING_EVENT, DATASET_DROPPING_COLUMN_EVENT, TRACE_DATASET_EVENTS,
28};
29use lance_datafusion::projection::ProjectionPlan;
30use lance_file::reader::{FileReader, FileReaderOptions};
31use lance_file::version::{ConcreteFileVersion, LanceFileVersion};
32use lance_index::{IndexType, progress::IndexBuildProgress};
33use lance_io::object_store::{
34    ChainedWrappingObjectStore, LanceNamespaceStorageOptionsProvider, ObjectStore,
35    ObjectStoreParams, StorageOptions, StorageOptionsAccessor, StorageOptionsProvider,
36    WrappingObjectStore,
37};
38use lance_io::scheduler::{ScanScheduler, SchedulerConfig};
39use lance_io::traits::{WriteExt, Writer};
40use lance_io::utils::{
41    CachedFileSize, read_last_block, read_message, read_metadata_offset, read_struct,
42};
43use lance_namespace::LanceNamespace;
44use lance_table::format::{
45    DataFile, DataStorageFormat, DeletionFile, Fragment, IndexMetadata, MAGIC, Manifest, RowIdMeta,
46    pb, populate_manifest_schema_dictionaries,
47};
48use lance_table::io::commit::{
49    CommitConfig, CommitError, CommitHandler, CommitLock, ManifestLocation, ManifestNamingScheme,
50    VERSIONS_DIR, external_manifest::ExternalManifestCommitHandler, migrate_scheme_to_v2,
51    write_manifest_file_to_path,
52};
53
54use crate::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
55use lance_table::io::manifest::{read_manifest, read_manifest_indexes};
56use object_store::ObjectStoreExt;
57use object_store::path::Path;
58use prost::Message;
59use roaring::RoaringBitmap;
60use rowids::get_row_id_index;
61use serde::{Deserialize, Serialize};
62use std::borrow::Cow;
63use std::collections::{BTreeMap, HashMap, HashSet};
64use std::fmt::Debug;
65use std::num::NonZero;
66use std::ops::Range;
67use std::pin::Pin;
68use std::sync::Arc;
69use tracing::{info, instrument};
70
71pub(crate) mod blob;
72pub(crate) mod branch_location;
73pub mod builder;
74pub mod cleanup;
75pub mod delta;
76pub mod files;
77pub mod fragment;
78mod hash_joiner;
79pub mod index;
80pub mod mem_wal;
81mod metadata;
82pub mod optimize;
83pub(crate) mod overlay;
84pub mod progress;
85pub mod refs;
86pub mod rowids;
87pub mod scanner;
88mod schema_evolution;
89pub mod sql;
90pub mod statistics;
91mod take;
92pub mod transaction;
93pub mod udtf;
94pub mod updater;
95mod utils;
96pub mod write;
97
98pub(crate) use take::row_offsets_to_row_addresses;
99
100use self::builder::DatasetBuilder;
101use self::cleanup::RemovalStats;
102use self::fragment::FileFragment;
103use self::refs::Refs;
104use self::scanner::{DatasetRecordBatchStream, Scanner};
105use self::statistics::DatasetStatistics;
106use self::transaction::{Operation, Transaction, TransactionBuilder, UpdateMapEntry};
107use self::write::{cleanup_data_fragments, write_fragments_internal};
108use crate::dataset::branch_location::BranchLocation;
109use crate::dataset::cleanup::{CleanupOperation, CleanupPolicy, CleanupPolicyBuilder};
110use crate::dataset::refs::{BranchContents, BranchIdentifier, Branches, Tags};
111use crate::dataset::sql::SqlQueryBuilder;
112use crate::datatypes::Schema;
113use crate::index::retain_supported_indices;
114use crate::io::commit::{
115    commit_detached_transaction, commit_new_dataset, commit_transaction,
116    detect_overlapping_fragments,
117};
118use crate::session::Session;
119use crate::utils::temporal::{SystemTime, timestamp_to_nanos, utc_now};
120use crate::{Error, Result};
121pub use blob::{
122    BlobFile, BlobRangeRequest, BlobReadRange, ReadBlob, ReadBlobRange, ReadBlobRangesBuilder,
123    ReadBlobRangesStream, ReadBlobsBuilder, ReadBlobsStream,
124};
125use hash_joiner::HashJoiner;
126pub use lance_core::ROW_ID;
127use lance_core::box_error;
128use lance_index::scalar::lance_format::LanceIndexStore;
129use lance_namespace::models::{DeclareTableRequest, DescribeTableRequest};
130use lance_table::feature_flags::{apply_feature_flags, can_read_dataset};
131use lance_table::io::deletion::{DELETIONS_DIR, relative_deletion_file_path};
132pub use schema_evolution::{
133    BatchInfo, BatchUDF, ColumnAlteration, NewColumnTransform, UDFCheckpointStore,
134};
135pub use take::TakeBuilder;
136use uuid::Uuid;
137pub use write::merge_insert::{
138    MergeInsertBuilder, MergeInsertJob, MergeStats, UncommittedMergeInsert, WhenMatched,
139    WhenNotMatched, WhenNotMatchedBySource,
140};
141
142use crate::dataset::index::LanceIndexStoreExt;
143pub use write::update::{UpdateBuilder, UpdateJob};
144#[allow(deprecated)]
145pub use write::{
146    AutoCleanupParams, CommitBuilder, DEFAULT_COMMIT_TIMEOUT, DeleteBuilder, DeleteResult,
147    ExternalBlobMode, InsertBuilder, UncommittedDelete, WriteDestination, WriteMode, WriteParams,
148    WriteProgressFn, WriteStats, write_fragments,
149};
150
151pub(crate) const INDICES_DIR: &str = "_indices";
152pub(crate) const DATA_DIR: &str = "data";
153pub(crate) const TRANSACTIONS_DIR: &str = "_transactions";
154
155// We default to 6GB for the index cache, since indices are often large but
156// worth caching.
157pub const DEFAULT_INDEX_CACHE_SIZE: usize = 6 * 1024 * 1024 * 1024;
158// Default to 1 GiB for the metadata cache. Column metadata can be like 40MB,
159// so this should be enough for a few hundred columns. Other metadata is much
160// smaller.
161pub const DEFAULT_METADATA_CACHE_SIZE: usize = 1024 * 1024 * 1024;
162
163/// Lance Dataset
164#[derive(Clone)]
165pub struct Dataset {
166    /// The primary dataset object store. Use [`Self::object_store`] when
167    /// resolving files that may carry a base id.
168    pub(crate) object_store: Arc<ObjectStore>,
169    pub(crate) commit_handler: Arc<dyn CommitHandler>,
170    /// Uri of the dataset.
171    ///
172    /// On cloud storage, we can not use [Dataset::base] to build the full uri because the
173    /// `bucket` is swallowed in the inner [ObjectStore].
174    uri: String,
175    pub(crate) base: Path,
176    pub manifest: Arc<Manifest>,
177    // Path for the manifest that is loaded. Used to get additional information,
178    // such as the index metadata.
179    pub(crate) manifest_location: ManifestLocation,
180    pub(crate) session: Arc<Session>,
181    pub refs: Refs,
182
183    // Bitmap of fragment ids in this dataset.
184    pub(crate) fragment_bitmap: Arc<RoaringBitmap>,
185
186    // These are references to session caches, but with the dataset URI as a prefix.
187    pub(crate) index_cache: Arc<DSIndexCache>,
188    pub(crate) metadata_cache: Arc<DSMetadataCache>,
189
190    /// File reader options to use when reading data files.
191    pub(crate) file_reader_options: Option<FileReaderOptions>,
192
193    /// Object store parameters used when opening this dataset.
194    /// These are used when creating object stores for additional base paths.
195    pub(crate) store_params: Option<Box<ObjectStoreParams>>,
196    /// Optional runtime-only object store parameters keyed by base path URI.
197    pub(crate) base_store_params: Option<Arc<HashMap<String, ObjectStoreParams>>>,
198}
199
200impl std::fmt::Debug for Dataset {
201    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202        f.debug_struct("Dataset")
203            .field("uri", &self.uri)
204            .field("base", &self.base)
205            .field("version", &self.manifest.version)
206            .field("cache_num_items", &self.session.approx_num_items())
207            .field("base_store_params", &self.base_store_params.is_some())
208            .finish()
209    }
210}
211
212/// Dataset Version
213#[derive(Deserialize, Serialize, Debug)]
214pub struct Version {
215    /// version number
216    pub version: u64,
217
218    /// Timestamp of dataset creation in UTC.
219    pub timestamp: DateTime<Utc>,
220
221    /// Key-value pairs of metadata.
222    pub metadata: BTreeMap<String, String>,
223}
224
225/// Convert Manifest to Data Version.
226impl From<&Manifest> for Version {
227    fn from(m: &Manifest) -> Self {
228        Self {
229            version: m.version,
230            timestamp: m.timestamp(),
231            metadata: m.summary().into(),
232        }
233    }
234}
235
236/// The transaction that produced a version of the dataset, along with the
237/// version's commit timestamp.
238///
239/// Returned by [`Dataset::read_version_transaction`], which reads this
240/// information directly from storage without checking out the version.
241#[derive(Debug, Clone)]
242pub struct VersionTransaction {
243    /// Version number.
244    pub version: u64,
245
246    /// Timestamp the version was committed, in UTC.
247    pub timestamp: DateTime<Utc>,
248
249    /// The transaction that produced this version, if one was recorded.
250    pub transaction: Option<Transaction>,
251}
252
253/// Customize read behavior of a dataset.
254#[derive(Clone, Debug)]
255pub struct ReadParams {
256    /// Size of the index cache in bytes. This cache stores index data in memory
257    /// for faster lookups. The default is 6 GiB.
258    pub index_cache_size_bytes: usize,
259
260    /// Size of the metadata cache in bytes. This cache stores metadata in memory
261    /// for faster open table and scans. The default is 1 GiB.
262    pub metadata_cache_size_bytes: usize,
263
264    /// If present, dataset will use this shared [`Session`] instead creating a new one.
265    ///
266    /// This is useful for sharing the same session across multiple datasets.
267    pub session: Option<Arc<Session>>,
268
269    pub store_options: Option<ObjectStoreParams>,
270
271    /// If present, dataset will use this to resolve the latest version
272    ///
273    /// Lance needs to be able to make atomic updates to the manifest.  This involves
274    /// coordination between readers and writers and we can usually rely on the filesystem
275    /// to do this coordination for us.
276    ///
277    /// Some file systems (e.g. S3) do not support atomic operations.  In this case, for
278    /// safety, we recommend an external commit mechanism (such as dynamodb) and, on the
279    /// read path, we need to reach out to that external mechanism to figure out the latest
280    /// version of the dataset.
281    ///
282    /// If this is not set then a default behavior is chosen that is appropriate for the
283    /// filesystem.
284    ///
285    /// If a custom object store is provided (via store_params.object_store) then this
286    /// must also be provided.
287    pub commit_handler: Option<Arc<dyn CommitHandler>>,
288
289    /// File reader options to use when reading data files.
290    ///
291    /// This allows control over features like caching repetition indices and validation.
292    /// Options set here act as dataset-level defaults and can be overridden on a
293    /// per-scan basis via [`Scanner::batch_size_bytes`](crate::dataset::scanner::Scanner::batch_size_bytes) or
294    /// [`Scanner::with_file_reader_options`](crate::dataset::scanner::Scanner::with_file_reader_options).
295    pub file_reader_options: Option<FileReaderOptions>,
296}
297
298impl ReadParams {
299    /// Set the cache size for indices. Set to zero, to disable the cache.
300    #[deprecated(
301        since = "0.30.0",
302        note = "Use `index_cache_size_bytes` instead, which accepts a size in bytes."
303    )]
304    pub fn index_cache_size(&mut self, cache_size: usize) -> &mut Self {
305        let assumed_entry_size = 20 * 1024 * 1024; // 20 MiB per entry
306        self.index_cache_size_bytes = cache_size * assumed_entry_size;
307        self
308    }
309
310    pub fn index_cache_size_bytes(&mut self, cache_size: usize) -> &mut Self {
311        self.index_cache_size_bytes = cache_size;
312        self
313    }
314
315    /// Set the cache size for the file metadata. Set to zero to disable this cache.
316    #[deprecated(
317        since = "0.30.0",
318        note = "Use `metadata_cache_size_bytes` instead, which accepts a size in bytes."
319    )]
320    pub fn metadata_cache_size(&mut self, cache_size: usize) -> &mut Self {
321        let assumed_entry_size = 10 * 1024 * 1024; // 10 MiB per entry
322        self.metadata_cache_size_bytes = cache_size * assumed_entry_size;
323        self
324    }
325
326    /// Set the cache size for the file metadata in bytes.
327    pub fn metadata_cache_size_bytes(&mut self, cache_size: usize) -> &mut Self {
328        self.metadata_cache_size_bytes = cache_size;
329        self
330    }
331
332    /// Set a shared session for the datasets.
333    pub fn session(&mut self, session: Arc<Session>) -> &mut Self {
334        self.session = Some(session);
335        self
336    }
337
338    /// Use the explicit locking to resolve the latest version
339    pub fn set_commit_lock<T: CommitLock + Send + Sync + 'static>(&mut self, lock: Arc<T>) {
340        self.commit_handler = Some(Arc::new(lock));
341    }
342
343    /// Set the file reader options.
344    pub fn file_reader_options(&mut self, options: FileReaderOptions) -> &mut Self {
345        self.file_reader_options = Some(options);
346        self
347    }
348}
349
350impl Default for ReadParams {
351    fn default() -> Self {
352        Self {
353            index_cache_size_bytes: DEFAULT_INDEX_CACHE_SIZE,
354            metadata_cache_size_bytes: DEFAULT_METADATA_CACHE_SIZE,
355            session: None,
356            store_options: None,
357            commit_handler: None,
358            file_reader_options: None,
359        }
360    }
361}
362
363#[derive(Debug, Clone)]
364pub enum ProjectionRequest {
365    Schema(Arc<Schema>),
366    Sql(Vec<(String, String)>),
367}
368
369impl ProjectionRequest {
370    pub fn from_columns(
371        columns: impl IntoIterator<Item = impl AsRef<str>>,
372        dataset_schema: &Schema,
373    ) -> Self {
374        let columns = columns
375            .into_iter()
376            .map(|s| s.as_ref().to_string())
377            .collect::<Vec<_>>();
378
379        let schema = dataset_schema
380            .project_preserve_system_columns(&columns)
381            .unwrap();
382        Self::Schema(Arc::new(schema))
383    }
384
385    pub fn from_schema(schema: Schema) -> Self {
386        Self::Schema(Arc::new(schema))
387    }
388
389    /// Provide a list of projection with SQL transform.
390    ///
391    /// # Parameters
392    /// - `columns`: A list of tuples where the first element is resulted column name and the second
393    ///   element is the SQL expression.
394    pub fn from_sql(
395        columns: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
396    ) -> Self {
397        Self::Sql(
398            columns
399                .into_iter()
400                .map(|(a, b)| (a.into(), b.into()))
401                .collect(),
402        )
403    }
404
405    pub fn into_projection_plan(self, dataset: Arc<Dataset>) -> Result<ProjectionPlan> {
406        match self {
407            Self::Schema(schema) => {
408                // The schema might contain system columns (_rowid, _rowaddr) which are not
409                // in the dataset schema. We handle these specially in ProjectionPlan::from_schema.
410                let system_columns_present = schema
411                    .fields
412                    .iter()
413                    .any(|f| lance_core::is_system_column(&f.name));
414
415                if system_columns_present {
416                    // If system columns are present, we can't use project_by_schema directly
417                    // Just pass the schema to ProjectionPlan::from_schema which handles it
418                    ProjectionPlan::from_schema(dataset, schema.as_ref())
419                } else {
420                    // No system columns, use normal path with validation
421                    let projection = dataset.schema().project_by_schema(
422                        schema.as_ref(),
423                        OnMissing::Error,
424                        OnTypeMismatch::Error,
425                    )?;
426                    ProjectionPlan::from_schema(dataset, &projection)
427                }
428            }
429            Self::Sql(columns) => ProjectionPlan::from_expressions(dataset, &columns),
430        }
431    }
432}
433
434impl From<Arc<Schema>> for ProjectionRequest {
435    fn from(schema: Arc<Schema>) -> Self {
436        Self::Schema(schema)
437    }
438}
439
440impl From<Schema> for ProjectionRequest {
441    fn from(schema: Schema) -> Self {
442        Self::from(Arc::new(schema))
443    }
444}
445
446impl Dataset {
447    /// Open an existing dataset.
448    ///
449    /// See also [DatasetBuilder].
450    #[instrument]
451    pub async fn open(uri: &str) -> Result<Self> {
452        DatasetBuilder::from_uri(uri).load().await
453    }
454
455    /// Check out a dataset version with a ref
456    pub async fn checkout_version(&self, version: impl Into<refs::Ref>) -> Result<Self> {
457        let reference: refs::Ref = version.into();
458        match reference {
459            refs::Ref::Version(branch, version_number) => {
460                self.checkout_by_ref(version_number, branch.as_deref())
461                    .await
462            }
463            refs::Ref::VersionNumber(version_number) => {
464                self.checkout_by_ref(Some(version_number), self.manifest.branch.as_deref())
465                    .await
466            }
467            refs::Ref::Tag(tag_name) => {
468                let tag_contents = self.tags().get(tag_name.as_str()).await?;
469                self.checkout_by_ref(Some(tag_contents.version), tag_contents.branch.as_deref())
470                    .await
471            }
472        }
473    }
474
475    pub fn tags(&self) -> Tags<'_> {
476        self.refs.tags()
477    }
478
479    /// A handle for cheap, index-derived statistics about this dataset (e.g. a
480    /// column's global value range) that never scan data.
481    pub fn statistics(&self) -> DatasetStatistics<'_> {
482        DatasetStatistics::new(self)
483    }
484
485    pub fn branches(&self) -> Branches<'_> {
486        self.refs.branches()
487    }
488
489    /// Check out the latest version of the dataset
490    pub async fn checkout_latest(&mut self) -> Result<()> {
491        let (manifest, manifest_location) = self.latest_manifest().await?;
492        self.manifest = manifest;
493        self.manifest_location = manifest_location;
494        self.fragment_bitmap = Arc::new(
495            self.manifest
496                .fragments
497                .iter()
498                .map(|f| f.id as u32)
499                .collect(),
500        );
501        Ok(())
502    }
503
504    /// Check out the latest version of the branch
505    pub async fn checkout_branch(&self, branch: &str) -> Result<Self> {
506        self.checkout_by_ref(None, Some(branch)).await
507    }
508
509    /// This is a two-phase operation:
510    /// - Create the branch dataset by shallow cloning.
511    /// - Create the branch metadata (a.k.a. `BranchContents`).
512    ///
513    /// These two phases are not atomic. We consider `BranchContents` as the source of truth
514    /// for the branch.
515    ///
516    /// The cleanup procedure should:
517    /// - Clean up zombie branch datasets that have no related `BranchContents`.
518    /// - Delete broken `BranchContents` entries that have no related branch dataset.
519    ///
520    /// If `create_branch` stops at phase 1, it may leave a zombie branch dataset,
521    /// which can be cleaned up later. Such a zombie dataset may cause a branch creation
522    /// failure if we use the same name to `create_branch`. In that case, you need to call
523    /// `force_delete_branch` to interactively clean up the zombie dataset.
524    pub async fn create_branch(
525        &mut self,
526        branch: &str,
527        version: impl Into<refs::Ref>,
528        store_params: Option<ObjectStoreParams>,
529    ) -> Result<Self> {
530        let (source_branch, version_number) = self.resolve_reference(version.into()).await?;
531        let branch_location = self.branch_location().find_branch(Some(branch))?;
532        let source_location = self
533            .branch_location()
534            .find_branch(source_branch.as_deref())?;
535        let clone_op = Operation::Clone {
536            is_shallow: true,
537            ref_name: source_branch.clone(),
538            ref_version: version_number,
539            ref_path: source_location.uri,
540            branch_name: Some(branch.to_string()),
541        };
542        let transaction = Transaction::new(version_number, clone_op, None);
543
544        let builder = CommitBuilder::new(WriteDestination::Uri(branch_location.uri.as_str()))
545            // Fall back to the dataset's own store params
546            .with_store_params(
547                store_params.unwrap_or(self.store_params.as_deref().cloned().unwrap_or_default()),
548            )
549            .with_object_store(Arc::new(self.object_store.as_ref().clone()))
550            .with_commit_handler(self.commit_handler.clone())
551            .with_storage_format(self.manifest.data_storage_format.lance_file_version()?);
552        let dataset = builder.execute(transaction).await?;
553
554        // Create BranchContents after shallow_clone
555        self.branches()
556            .create(branch, version_number, source_branch.as_deref())
557            .await?;
558        Ok(dataset)
559    }
560
561    pub async fn delete_branch(&mut self, branch: &str) -> Result<()> {
562        self.branches().delete(branch, false).await
563    }
564
565    /// Delete the branch even if the BranchContents is not found.
566    /// This could be useful when we have zombie branches and want to clean them up immediately.
567    pub async fn force_delete_branch(&mut self, branch: &str) -> Result<()> {
568        self.branches().delete(branch, true).await
569    }
570
571    pub async fn list_branches(&self) -> Result<HashMap<String, BranchContents>> {
572        self.branches().list().await
573    }
574
575    fn already_checked_out(&self, location: &ManifestLocation, branch_name: Option<&str>) -> bool {
576        // We check the e_tag here just in case it has been overwritten. This can
577        // happen if the table has been dropped then re-created recently.
578        self.manifest.branch.as_deref() == branch_name
579            && self.manifest.version == location.version
580            && self.manifest_location.naming_scheme == location.naming_scheme
581            && location.e_tag.as_ref().is_some_and(|e_tag| {
582                self.manifest_location
583                    .e_tag
584                    .as_ref()
585                    .is_some_and(|current_e_tag| e_tag == current_e_tag)
586            })
587    }
588
589    async fn checkout_by_ref(
590        &self,
591        version_number: Option<u64>,
592        branch: Option<&str>,
593    ) -> Result<Self> {
594        // Reject malformed names at the boundary (mirroring the branch CRUD
595        // paths) so they fail as InvalidRef instead of tripping the wrong-chain
596        // check below
597        if let Some(branch_name) = branch
598            && !Branches::is_main_branch(branch)
599        {
600            refs::check_valid_branch(branch_name)?;
601        }
602
603        let new_location = self.branch_location().find_branch(branch)?;
604
605        let manifest_location = if let Some(version_number) = version_number {
606            self.commit_handler
607                .resolve_version_location(
608                    &new_location.path,
609                    version_number,
610                    &self.object_store.inner,
611                )
612                .await?
613        } else {
614            self.commit_handler
615                .resolve_latest_location(&new_location.path, &self.object_store)
616                .await?
617        };
618
619        if self.already_checked_out(&manifest_location, branch) {
620            return Ok(self.clone());
621        }
622
623        let manifest = Self::get_manifest(
624            self.object_store.as_ref(),
625            &manifest_location,
626            &new_location.uri,
627            self.session.as_ref(),
628        )
629        .await?;
630
631        // The resolved manifest must belong to the requested branch. A mismatch
632        // means the commit handler resolved against a different chain (for
633        // example an external manifest store that ignores branch-qualified
634        // paths); error loudly rather than hand back another branch's data.
635        let requested_branch = branch.and_then(refs::standardize_branch);
636        if manifest.branch.as_deref() != requested_branch.as_deref() {
637            return Err(Error::internal(format!(
638                "checkout of branch '{}' at version {} resolved a manifest belonging to branch '{}'",
639                refs::normalize_branch(branch),
640                manifest.version,
641                refs::normalize_branch(manifest.branch.as_deref()),
642            )));
643        }
644
645        Self::checkout_manifest(
646            self.object_store.clone(),
647            new_location.path,
648            new_location.uri,
649            manifest,
650            manifest_location,
651            self.session.clone(),
652            self.commit_handler.clone(),
653            self.file_reader_options.clone(),
654            self.store_params.as_deref().cloned(),
655            self.base_store_params.clone(),
656        )
657    }
658
659    pub(crate) async fn load_manifest(
660        object_store: &ObjectStore,
661        manifest_location: &ManifestLocation,
662        uri: &str,
663        session: &Session,
664    ) -> Result<Manifest> {
665        let object_reader = if let Some(size) = manifest_location.size {
666            object_store
667                .open_with_size(&manifest_location.path, size as usize)
668                .await
669        } else {
670            object_store.open(&manifest_location.path).await
671        };
672        let object_reader = object_reader.map_err(|e| match &e {
673            Error::NotFound { uri, .. } => Error::dataset_not_found(uri.clone(), box_error(e)),
674            _ => e,
675        })?;
676
677        let last_block =
678            read_last_block(object_reader.as_ref())
679                .await
680                .map_err(|err| match err {
681                    object_store::Error::NotFound { path, source } => {
682                        Error::dataset_not_found(path, source)
683                    }
684                    _ => Error::io_source(err.into()),
685                })?;
686
687        // A stale cached size yields a bogus footer offset. Detect it (the block
688        // lacks the trailing magic) and retry with the true size, like
689        // read_manifest.
690        if manifest_location.size.is_some() && !last_block.ends_with(MAGIC) {
691            let manifest_location = ManifestLocation {
692                size: None,
693                ..manifest_location.clone()
694            };
695            return Box::pin(Self::load_manifest(
696                object_store,
697                &manifest_location,
698                uri,
699                session,
700            ))
701            .await;
702        }
703
704        let offset = read_metadata_offset(&last_block)?;
705
706        // If manifest is in the last block, we can decode directly from memory.
707        let manifest_size = object_reader.size().await?;
708        let mut manifest = if manifest_size - offset <= last_block.len() {
709            let manifest_len = manifest_size - offset;
710            let offset_in_block = last_block.len() - manifest_len;
711            let message_len =
712                LittleEndian::read_u32(&last_block[offset_in_block..offset_in_block + 4]) as usize;
713            let message_data = &last_block[offset_in_block + 4..offset_in_block + 4 + message_len];
714            Manifest::try_from(lance_table::format::pb::Manifest::decode(message_data)?)
715        } else {
716            read_struct(object_reader.as_ref(), offset).await
717        }?;
718
719        if !can_read_dataset(manifest.reader_feature_flags) {
720            let message = format!(
721                "This dataset cannot be read by this version of Lance. \
722                 Please upgrade Lance to read this dataset.\n Flags: {}",
723                manifest.reader_feature_flags
724            );
725            return Err(Error::not_supported_source(message.into()));
726        }
727
728        // If indices were also in the last block, we can take the opportunity to
729        // decode them now and cache them.
730        if let Some(index_offset) = manifest.index_section
731            && manifest_size - index_offset <= last_block.len()
732        {
733            let offset_in_block = last_block.len() - (manifest_size - index_offset);
734            let message_len =
735                LittleEndian::read_u32(&last_block[offset_in_block..offset_in_block + 4]) as usize;
736            let message_data = &last_block[offset_in_block + 4..offset_in_block + 4 + message_len];
737            let section = lance_table::format::pb::IndexSection::decode(message_data)?;
738            let mut indices: Vec<IndexMetadata> = section
739                .indices
740                .into_iter()
741                .map(IndexMetadata::try_from)
742                .collect::<Result<Vec<_>>>()?;
743            retain_supported_indices(&mut indices);
744            let ds_index_cache = session.index_cache.for_dataset(uri);
745            let metadata_key = crate::session::index_caches::IndexMetadataKey {
746                version: manifest_location.version,
747                store_identity: &object_store.store_prefix,
748            };
749            ds_index_cache
750                .insert_with_key(&metadata_key, Arc::new(indices))
751                .await;
752        }
753
754        // If transaction is also in the last block, we can take the opportunity to
755        // decode them now and cache them.
756        if let Some(transaction_offset) = manifest.transaction_section
757            && manifest_size - transaction_offset <= last_block.len()
758        {
759            let offset_in_block = last_block.len() - (manifest_size - transaction_offset);
760            let message_len =
761                LittleEndian::read_u32(&last_block[offset_in_block..offset_in_block + 4]) as usize;
762            let message_data = &last_block[offset_in_block + 4..offset_in_block + 4 + message_len];
763            let transaction: Transaction =
764                lance_table::format::pb::Transaction::decode(message_data)?.try_into()?;
765
766            let metadata_cache = session.metadata_cache.for_dataset(uri);
767            let metadata_key = TransactionKey {
768                version: manifest_location.version,
769            };
770            metadata_cache
771                .insert_with_key(&metadata_key, Arc::new(transaction))
772                .await;
773        }
774
775        populate_manifest_schema_dictionaries(&mut manifest, object_reader.as_ref()).await?;
776
777        Ok(manifest)
778    }
779
780    /// Fetch the manifest for `manifest_location` from the session metadata
781    /// cache, loading and caching it on a miss.
782    pub(crate) async fn get_manifest(
783        object_store: &ObjectStore,
784        manifest_location: &ManifestLocation,
785        uri: &str,
786        session: &Session,
787    ) -> Result<Arc<Manifest>> {
788        if manifest_location.size.is_none() {
789            return Ok(Arc::new(
790                Self::load_manifest(object_store, manifest_location, uri, session).await?,
791            ));
792        }
793        let metadata_cache = session.metadata_cache.for_dataset(uri);
794        let manifest_key = ManifestKey {
795            version: manifest_location.version,
796            e_tag: manifest_location.e_tag.as_deref(),
797        };
798        if let Some(cached) = metadata_cache.get_with_key(&manifest_key).await {
799            return Ok(cached);
800        }
801        let loaded =
802            Arc::new(Self::load_manifest(object_store, manifest_location, uri, session).await?);
803        metadata_cache
804            .insert_with_key(&manifest_key, loaded.clone())
805            .await;
806        Ok(loaded)
807    }
808
809    #[allow(clippy::too_many_arguments)]
810    fn checkout_manifest(
811        object_store: Arc<ObjectStore>,
812        base_path: Path,
813        uri: String,
814        manifest: Arc<Manifest>,
815        manifest_location: ManifestLocation,
816        session: Arc<Session>,
817        commit_handler: Arc<dyn CommitHandler>,
818        file_reader_options: Option<FileReaderOptions>,
819        store_params: Option<ObjectStoreParams>,
820        base_store_params: Option<Arc<HashMap<String, ObjectStoreParams>>>,
821    ) -> Result<Self> {
822        let refs = Refs::new(
823            object_store.clone(),
824            commit_handler.clone(),
825            BranchLocation {
826                path: base_path.clone(),
827                uri: uri.clone(),
828                branch: manifest.branch.clone(),
829            },
830        );
831        let metadata_cache = Arc::new(session.metadata_cache.for_dataset(&uri));
832        let index_cache = Arc::new(session.index_cache.for_dataset(&uri));
833        let fragment_bitmap = Arc::new(manifest.fragments.iter().map(|f| f.id as u32).collect());
834        write::log_unregistered_base_scoped_options(
835            store_params.as_ref(),
836            &manifest.base_paths,
837            log::Level::Debug,
838        );
839        Ok(Self {
840            object_store,
841            base: base_path,
842            uri,
843            manifest,
844            manifest_location,
845            commit_handler,
846            session,
847            refs,
848            fragment_bitmap,
849            metadata_cache,
850            index_cache,
851            file_reader_options,
852            store_params: store_params.map(Box::new),
853            base_store_params,
854        })
855    }
856
857    /// Write to or Create a [Dataset] with a stream of [RecordBatch]s.
858    ///
859    /// `dest` can be a `&str`, `object_store::path::Path` or `Arc<Dataset>`.
860    ///
861    /// Returns the newly created [`Dataset`].
862    /// Or Returns [Error] if the dataset already exists.
863    ///
864    pub async fn write(
865        batches: impl RecordBatchReader + Send + 'static,
866        dest: impl Into<WriteDestination<'_>>,
867        params: Option<WriteParams>,
868    ) -> Result<Self> {
869        let mut builder = InsertBuilder::new(dest);
870        if let Some(params) = &params {
871            builder = builder.with_params(params);
872        }
873        Box::pin(builder.execute_stream(Box::new(batches) as Box<dyn RecordBatchReader + Send>))
874            .await
875    }
876
877    /// Write into a namespace client-managed table with automatic credential vending.
878    ///
879    /// For CREATE mode, calls declare_table() to initialize the table.
880    /// For other modes, calls describe_table() and opens dataset with namespace client credentials.
881    ///
882    /// # Arguments
883    ///
884    /// * `batches` - The record batches to write
885    /// * `namespace_client` - The namespace client to use for table management
886    /// * `table_id` - The table identifier
887    /// * `params` - Write parameters
888    pub async fn write_into_namespace(
889        batches: impl RecordBatchReader + Send + 'static,
890        namespace_client: Arc<dyn LanceNamespace>,
891        table_id: Vec<String>,
892        params: Option<WriteParams>,
893    ) -> Result<Self> {
894        Self::write_into_namespace_impl(batches, namespace_client, table_id, None, params).await
895    }
896
897    /// Write into a branch of a namespace client-managed table.
898    ///
899    /// Behaves like [`write_into_namespace`](Self::write_into_namespace), but APPEND and
900    /// OVERWRITE open and commit against `branch` instead of main. CREATE is rejected,
901    /// since a branch forks from an existing version.
902    pub async fn write_into_namespace_on_branch(
903        batches: impl RecordBatchReader + Send + 'static,
904        namespace_client: Arc<dyn LanceNamespace>,
905        table_id: Vec<String>,
906        branch: &str,
907        params: Option<WriteParams>,
908    ) -> Result<Self> {
909        Self::write_into_namespace_impl(
910            batches,
911            namespace_client,
912            table_id,
913            Some(branch.to_string()),
914            params,
915        )
916        .await
917    }
918
919    async fn write_into_namespace_impl(
920        batches: impl RecordBatchReader + Send + 'static,
921        namespace_client: Arc<dyn LanceNamespace>,
922        table_id: Vec<String>,
923        branch: Option<String>,
924        mut params: Option<WriteParams>,
925    ) -> Result<Self> {
926        let mut write_params = params.take().unwrap_or_default();
927
928        match write_params.mode {
929            WriteMode::Create => {
930                if branch.is_some() {
931                    return Err(Error::not_supported_source(
932                        "cannot create a table on a branch; create on main first, then branch it"
933                            .into(),
934                    ));
935                }
936                let declare_request = DeclareTableRequest {
937                    id: Some(table_id.clone()),
938                    ..Default::default()
939                };
940                let response = namespace_client
941                    .declare_table(declare_request)
942                    .await
943                    .map_err(|e| Error::namespace_source(Box::new(e)))?;
944
945                let uri = response.location.ok_or_else(|| {
946                    Error::namespace_source(Box::new(std::io::Error::other(
947                        "Table location not found in declare_table response",
948                    )))
949                })?;
950
951                // Set up commit handler when managed_versioning is enabled
952                if response.managed_versioning == Some(true) {
953                    // The store derives the branch a request targets from the
954                    // base path it is handed, resolved against the table root.
955                    let external_store = LanceNamespaceExternalManifestStore::for_table_uri(
956                        namespace_client.clone(),
957                        table_id.clone(),
958                        &uri,
959                    )?;
960                    let commit_handler: Arc<dyn CommitHandler> =
961                        Arc::new(ExternalManifestCommitHandler {
962                            external_manifest_store: Arc::new(external_store),
963                        });
964                    write_params.commit_handler = Some(commit_handler);
965                }
966
967                // Set initial credentials and provider from namespace_client
968                if let Some(namespace_storage_options) = response.storage_options {
969                    let provider: Arc<dyn StorageOptionsProvider> = Arc::new(
970                        LanceNamespaceStorageOptionsProvider::new(namespace_client, table_id),
971                    );
972
973                    // Merge namespace client storage options with any existing options
974                    let mut merged_options = write_params
975                        .store_params
976                        .as_ref()
977                        .and_then(|p| p.storage_options().cloned())
978                        .unwrap_or_default();
979                    merged_options.extend(namespace_storage_options);
980
981                    let accessor = Arc::new(StorageOptionsAccessor::with_initial_and_provider(
982                        merged_options,
983                        provider,
984                    ));
985
986                    let existing_params = write_params.store_params.take().unwrap_or_default();
987                    write_params.store_params = Some(ObjectStoreParams {
988                        storage_options_accessor: Some(accessor),
989                        ..existing_params
990                    });
991                }
992
993                Self::write(batches, uri.as_str(), Some(write_params)).await
994            }
995            WriteMode::Append | WriteMode::Overwrite => {
996                let request = DescribeTableRequest {
997                    id: Some(table_id.clone()),
998                    ..Default::default()
999                };
1000                let response = namespace_client
1001                    .describe_table(request)
1002                    .await
1003                    .map_err(|e| Error::namespace_source(Box::new(e)))?;
1004
1005                let uri = response.location.ok_or_else(|| {
1006                    Error::namespace_source(Box::new(std::io::Error::other(
1007                        "Table location not found in describe_table response",
1008                    )))
1009                })?;
1010
1011                // Set up commit handler when managed_versioning is enabled.
1012                // It must ride on the dataset opened below: InsertBuilder
1013                // commits through the destination dataset's handler and does
1014                // not consult write params for Dataset destinations.
1015                let commit_handler: Option<Arc<dyn CommitHandler>> =
1016                    if response.managed_versioning == Some(true) {
1017                        // The store derives the branch a request targets from the
1018                        // base path it is handed, resolved against the table root.
1019                        let external_store = LanceNamespaceExternalManifestStore::for_table_uri(
1020                            namespace_client.clone(),
1021                            table_id.clone(),
1022                            uri.as_str(),
1023                        )?;
1024                        Some(Arc::new(ExternalManifestCommitHandler {
1025                            external_manifest_store: Arc::new(external_store),
1026                        }))
1027                    } else {
1028                        None
1029                    };
1030
1031                // Set initial credentials and provider from namespace_client
1032                if let Some(namespace_storage_options) = response.storage_options {
1033                    let provider: Arc<dyn StorageOptionsProvider> =
1034                        Arc::new(LanceNamespaceStorageOptionsProvider::new(
1035                            namespace_client.clone(),
1036                            table_id.clone(),
1037                        ));
1038
1039                    // Merge namespace client storage options with any existing options
1040                    let mut merged_options = write_params
1041                        .store_params
1042                        .as_ref()
1043                        .and_then(|p| p.storage_options().cloned())
1044                        .unwrap_or_default();
1045                    merged_options.extend(namespace_storage_options);
1046
1047                    let accessor = Arc::new(StorageOptionsAccessor::with_initial_and_provider(
1048                        merged_options,
1049                        provider,
1050                    ));
1051
1052                    let existing_params = write_params.store_params.take().unwrap_or_default();
1053                    write_params.store_params = Some(ObjectStoreParams {
1054                        storage_options_accessor: Some(accessor),
1055                        ..existing_params
1056                    });
1057                }
1058
1059                // For APPEND/OVERWRITE modes, we must open the existing dataset first
1060                // and pass it to InsertBuilder. If we pass just the URI, InsertBuilder
1061                // assumes no dataset exists and converts the mode to CREATE.
1062                let mut builder = DatasetBuilder::from_uri(uri.as_str());
1063                if let Some(ref store_params) = write_params.store_params
1064                    && let Some(accessor) = &store_params.storage_options_accessor
1065                {
1066                    builder = builder.with_storage_options_accessor(accessor.clone());
1067                }
1068                if let Some(commit_handler) = commit_handler {
1069                    builder = builder.with_commit_handler(commit_handler);
1070                }
1071                if let Some(branch) = &branch {
1072                    builder = builder.with_branch(branch, None);
1073                }
1074                let dataset = Arc::new(builder.load().await?);
1075
1076                Self::write(batches, dataset, Some(write_params)).await
1077            }
1078        }
1079    }
1080
1081    /// Append to existing [Dataset] with a stream of [RecordBatch]s
1082    ///
1083    /// Returns void result or Returns [Error]
1084    pub async fn append(
1085        &mut self,
1086        batches: impl RecordBatchReader + Send + 'static,
1087        params: Option<WriteParams>,
1088    ) -> Result<()> {
1089        let write_params = WriteParams {
1090            mode: WriteMode::Append,
1091            ..params.unwrap_or_default()
1092        };
1093
1094        let new_dataset = InsertBuilder::new(WriteDestination::Dataset(Arc::new(self.clone())))
1095            .with_params(&write_params)
1096            .execute_stream(Box::new(batches) as Box<dyn RecordBatchReader + Send>)
1097            .await?;
1098
1099        *self = new_dataset;
1100
1101        Ok(())
1102    }
1103
1104    /// Get the fully qualified URI of this dataset.
1105    pub fn uri(&self) -> &str {
1106        &self.uri
1107    }
1108
1109    pub fn branch_location(&self) -> BranchLocation {
1110        BranchLocation {
1111            path: self.base.clone(),
1112            uri: self.uri.clone(),
1113            branch: self.manifest.branch.clone(),
1114        }
1115    }
1116
1117    pub async fn branch_identifier(&self) -> Result<BranchIdentifier> {
1118        self.refs
1119            .branches()
1120            .get_identifier(self.manifest.branch.as_deref())
1121            .await
1122    }
1123
1124    /// Get the full manifest of the dataset version.
1125    pub fn manifest(&self) -> &Manifest {
1126        &self.manifest
1127    }
1128
1129    pub fn manifest_location(&self) -> &ManifestLocation {
1130        &self.manifest_location
1131    }
1132
1133    /// Create a [`delta::DatasetDeltaBuilder`] to explore changes between dataset versions.
1134    ///
1135    /// # Example
1136    ///
1137    /// ```
1138    /// # use lance::{Dataset, Result};
1139    /// # async fn example(dataset: &Dataset) -> Result<()> {
1140    /// let delta = dataset.delta()
1141    ///     .compared_against_version(5)
1142    ///     .build()?;
1143    /// let inserted = delta.get_inserted_rows().await?;
1144    /// # Ok(())
1145    /// # }
1146    /// ```
1147    pub fn delta(&self) -> delta::DatasetDeltaBuilder {
1148        delta::DatasetDeltaBuilder::new(self.clone())
1149    }
1150
1151    // TODO: Cache this
1152    pub(crate) fn is_legacy_storage(&self) -> bool {
1153        self.manifest
1154            .data_storage_format
1155            .lance_file_version()
1156            .unwrap()
1157            == LanceFileVersion::Legacy
1158    }
1159
1160    pub async fn latest_manifest(&self) -> Result<(Arc<Manifest>, ManifestLocation)> {
1161        let location = self
1162            .commit_handler
1163            .resolve_latest_location(&self.base, &self.object_store)
1164            .await?;
1165
1166        // Check if manifest is in cache before reading from storage
1167        let manifest_key = ManifestKey {
1168            version: location.version,
1169            e_tag: location.e_tag.as_deref(),
1170        };
1171        let cached_manifest = self.metadata_cache.get_with_key(&manifest_key).await;
1172        if let Some(cached_manifest) = cached_manifest {
1173            return Ok((cached_manifest, location));
1174        }
1175
1176        if self.already_checked_out(&location, self.manifest.branch.as_deref()) {
1177            return Ok((self.manifest.clone(), self.manifest_location.clone()));
1178        }
1179        let mut manifest = read_manifest(&self.object_store, &location.path, location.size).await?;
1180        if manifest.schema.has_dictionary_types() {
1181            let reader = if let Some(size) = location.size {
1182                self.object_store
1183                    .open_with_size(&location.path, size as usize)
1184                    .await?
1185            } else {
1186                self.object_store.open(&location.path).await?
1187            };
1188            populate_manifest_schema_dictionaries(&mut manifest, reader.as_ref()).await?;
1189        }
1190        let manifest_arc = Arc::new(manifest);
1191        self.metadata_cache
1192            .insert_with_key(&manifest_key, manifest_arc.clone())
1193            .await;
1194        Ok((manifest_arc, location))
1195    }
1196
1197    /// Read the transaction file for this version of the dataset.
1198    ///
1199    /// If there was no transaction file written for this version of the dataset
1200    /// then this will return None.
1201    pub async fn read_transaction(&self) -> Result<Option<Transaction>> {
1202        let transaction_key = TransactionKey {
1203            version: self.manifest.version,
1204        };
1205        if let Some(transaction) = self.metadata_cache.get_with_key(&transaction_key).await {
1206            return Ok(Some((*transaction).clone()));
1207        }
1208
1209        let transaction = self
1210            .read_transaction_from_storage(&self.manifest, &self.manifest_location)
1211            .await?;
1212
1213        if let Some(tx) = transaction.as_ref() {
1214            self.metadata_cache
1215                .insert_with_key(&transaction_key, Arc::new(tx.clone()))
1216                .await;
1217        }
1218        Ok(transaction)
1219    }
1220
1221    /// Read the transaction recorded by `manifest` directly from storage,
1222    /// without consulting or populating any session cache.
1223    async fn read_transaction_from_storage(
1224        &self,
1225        manifest: &Manifest,
1226        manifest_location: &ManifestLocation,
1227    ) -> Result<Option<Transaction>> {
1228        // Prefer inline transaction from manifest when available
1229        if let Some(pos) = manifest.transaction_section {
1230            let reader = match manifest_location.size {
1231                Some(size) => {
1232                    self.object_store
1233                        .open_with_size(&manifest_location.path, size as usize)
1234                        .await?
1235                }
1236                None => self.object_store.open(&manifest_location.path).await?,
1237            };
1238
1239            // A concurrent overwrite can leave the listed size too small; retry
1240            // once with the true size.
1241            let tx: pb::Transaction = match read_message(reader.as_ref(), pos).await {
1242                Err(e)
1243                    if manifest_location.size.is_some()
1244                        && e.to_string().contains("file size is too small") =>
1245                {
1246                    let reader = self.object_store.open(&manifest_location.path).await?;
1247                    read_message(reader.as_ref(), pos).await?
1248                }
1249                other => other?,
1250            };
1251            Transaction::try_from(tx).map(Some)
1252        } else if let Some(path) = &manifest.transaction_file {
1253            // Fallback: read external transaction file if present
1254            let path = self.transactions_dir().join(path.as_str());
1255            let data = self.object_store.inner.get(&path).await?.bytes().await?;
1256            let transaction = lance_table::format::pb::Transaction::decode(data)?;
1257            Transaction::try_from(transaction).map(Some)
1258        } else {
1259            Ok(None)
1260        }
1261    }
1262
1263    /// Read the transaction (if any) and commit timestamp of a version of the
1264    /// dataset. `version` is a version number on this dataset's current branch.
1265    ///
1266    /// Reads the version's manifest transiently: no historical `Dataset` is
1267    /// constructed, no `IndexSection` is decoded, and no session cache is read
1268    /// or written, so scanning many historical versions does not fill the
1269    /// shared caches.
1270    ///
1271    /// Returns an error if the version does not exist (for example, if it has
1272    /// been cleaned up).
1273    ///
1274    /// # Example
1275    ///
1276    /// ```
1277    /// # use lance::{Dataset, Result};
1278    /// # async fn example(dataset: &Dataset) -> Result<()> {
1279    /// let record = dataset.read_version_transaction(5).await?;
1280    /// let committed_at = record.timestamp;
1281    /// let operation = record.transaction.as_ref().map(|t| t.operation.name());
1282    /// # Ok(())
1283    /// # }
1284    /// ```
1285    pub async fn read_version_transaction(&self, version: u64) -> Result<VersionTransaction> {
1286        // Resolve against this dataset's current branch.
1287        let manifest_location = self
1288            .commit_handler
1289            .resolve_version_location(&self.base, version, &self.object_store.inner)
1290            .await?;
1291
1292        // Keep the DatasetNotFound variant callers expect for a missing version.
1293        let manifest = read_manifest(
1294            &self.object_store,
1295            &manifest_location.path,
1296            manifest_location.size,
1297        )
1298        .await
1299        .map_err(|e| match &e {
1300            Error::NotFound { uri, .. } => Error::dataset_not_found(uri.clone(), box_error(e)),
1301            _ => e,
1302        })?;
1303
1304        // The resolved manifest must belong to this dataset's branch. A
1305        // mismatch means the commit handler resolved against a different chain
1306        // (for example an external manifest store that ignores
1307        // branch-qualified paths); error loudly rather than hand back another
1308        // branch's transaction.
1309        if manifest.branch != self.manifest.branch {
1310            return Err(Error::internal(format!(
1311                "reading version {} on branch '{}' resolved a manifest belonging to branch '{}'",
1312                version,
1313                refs::normalize_branch(self.manifest.branch.as_deref()),
1314                refs::normalize_branch(manifest.branch.as_deref()),
1315            )));
1316        }
1317
1318        let transaction = self
1319            .read_transaction_from_storage(&manifest, &manifest_location)
1320            .await?;
1321
1322        Ok(VersionTransaction {
1323            version: manifest.version,
1324            timestamp: manifest.timestamp(),
1325            transaction,
1326        })
1327    }
1328
1329    /// Read the transaction file for this version of the dataset.
1330    ///
1331    /// If there was no transaction file written for this version of the dataset
1332    /// then this will return None.
1333    ///
1334    /// Does not populate the session caches; see
1335    /// [`Self::read_version_transaction`].
1336    ///
1337    /// # Example
1338    ///
1339    /// ```
1340    /// # use lance::{Dataset, Result};
1341    /// # async fn example(dataset: &Dataset) -> Result<()> {
1342    /// let transaction = dataset.read_transaction_by_version(5).await?;
1343    /// let operation = transaction.as_ref().map(|t| t.operation.name());
1344    /// # Ok(())
1345    /// # }
1346    /// ```
1347    pub async fn read_transaction_by_version(&self, version: u64) -> Result<Option<Transaction>> {
1348        Ok(self.read_version_transaction(version).await?.transaction)
1349    }
1350
1351    /// List transactions for the dataset, up to a maximum number.
1352    ///
1353    /// This method iterates through dataset versions, starting from the current version,
1354    /// and collects the transaction for each version. It stops when either `recent_transactions`
1355    /// is reached or there are no more versions.
1356    ///
1357    /// # Arguments
1358    ///
1359    /// * `recent_transactions` - Maximum number of transactions to return
1360    ///
1361    /// # Returns
1362    ///
1363    /// A vector of optional transactions. Each element corresponds to a version,
1364    /// and may be None if no transaction file exists for that version.
1365    pub async fn get_transactions(
1366        &self,
1367        recent_transactions: usize,
1368    ) -> Result<Vec<Option<Transaction>>> {
1369        let mut transactions = vec![];
1370        let mut dataset = self.clone();
1371
1372        loop {
1373            let transaction = dataset.read_transaction().await?;
1374            transactions.push(transaction);
1375
1376            if transactions.len() >= recent_transactions {
1377                break;
1378            } else {
1379                match dataset
1380                    .checkout_version(dataset.version().version - 1)
1381                    .await
1382                {
1383                    Ok(ds) => dataset = ds,
1384                    Err(Error::DatasetNotFound { .. }) => break,
1385                    Err(err) => return Err(err),
1386                }
1387            }
1388        }
1389
1390        Ok(transactions)
1391    }
1392
1393    /// Restore the currently checked out version of the dataset as the latest version.
1394    pub async fn restore(&mut self) -> Result<()> {
1395        let (latest_manifest, _) = self.latest_manifest().await?;
1396        let latest_version = latest_manifest.version;
1397
1398        let transaction = Transaction::new(
1399            latest_version,
1400            Operation::Restore {
1401                version: self.manifest.version,
1402            },
1403            None,
1404        );
1405
1406        self.apply_commit(transaction, &Default::default(), &Default::default())
1407            .await?;
1408
1409        Ok(())
1410    }
1411
1412    /// Removes old versions of the dataset from disk
1413    ///
1414    /// This function will remove all versions of the dataset that are older than the provided
1415    /// timestamp.  This function will not remove the current version of the dataset.
1416    ///
1417    /// Once a version is removed it can no longer be checked out or restored.  Any data unique
1418    /// to that version will be lost.
1419    ///
1420    /// # Arguments
1421    ///
1422    /// * `older_than` - Versions older than this will be deleted.
1423    /// * `delete_unverified` - If false (the default) then files will only be deleted if they
1424    ///                        are listed in at least one manifest.  Otherwise these files will
1425    ///                        be kept since they cannot be distinguished from an in-progress
1426    ///                        transaction.  Set to true to delete these files if you are sure
1427    ///                        there are no other in-progress dataset operations.
1428    ///
1429    /// # Returns
1430    ///
1431    /// * `RemovalStats` - Statistics about the removal operation
1432    #[instrument(level = "debug", skip(self))]
1433    pub fn cleanup_old_versions(
1434        &self,
1435        older_than: Duration,
1436        delete_unverified: Option<bool>,
1437        error_if_tagged_old_versions: Option<bool>,
1438    ) -> BoxFuture<'_, Result<RemovalStats>> {
1439        let mut builder = CleanupPolicyBuilder::default();
1440        builder = builder.before_timestamp(utc_now() - older_than);
1441        if let Some(v) = delete_unverified {
1442            builder = builder.delete_unverified(v);
1443        }
1444        if let Some(v) = error_if_tagged_old_versions {
1445            builder = builder.error_if_tagged_old_versions(v);
1446        }
1447
1448        self.cleanup_with_policy(builder.build())
1449    }
1450
1451    /// Removes old versions of the dataset from storage
1452    ///
1453    /// This function will remove all versions of the dataset that satisfies the given policy.
1454    /// This function will not remove the current version of the dataset.
1455    ///
1456    /// Once a version is removed it can no longer be checked out or restored.  Any data unique
1457    /// to that version will be lost.
1458    ///
1459    /// # Arguments
1460    ///
1461    /// * `policy` - `CleanupPolicy` determines the behaviour of cleanup.
1462    ///
1463    /// # Returns
1464    ///
1465    /// * `RemovalStats` - Statistics about the removal operation
1466    #[instrument(level = "debug", skip(self))]
1467    pub fn cleanup_with_policy(
1468        &self,
1469        policy: CleanupPolicy,
1470    ) -> BoxFuture<'_, Result<RemovalStats>> {
1471        async move { self.cleanup(policy).execute().await }.boxed()
1472    }
1473
1474    /// Creates a cleanup operation for this dataset.
1475    ///
1476    /// The returned operation can be explained without deleting files, or
1477    /// executed to re-evaluate the current dataset state and remove files.
1478    pub fn cleanup(&self, policy: CleanupPolicy) -> CleanupOperation<'_> {
1479        CleanupOperation::new(self, policy)
1480    }
1481
1482    #[allow(clippy::too_many_arguments)]
1483    async fn do_commit(
1484        base_uri: WriteDestination<'_>,
1485        operation: Operation,
1486        read_version: Option<u64>,
1487        store_params: Option<ObjectStoreParams>,
1488        commit_handler: Option<Arc<dyn CommitHandler>>,
1489        session: Arc<Session>,
1490        enable_v2_manifest_paths: bool,
1491        detached: bool,
1492    ) -> Result<Self> {
1493        let read_version = read_version.map_or_else(
1494            || match operation {
1495                Operation::Overwrite { .. } | Operation::Restore { .. } => Ok(0),
1496                _ => Err(Error::invalid_input(
1497                    "read_version must be specified for this operation",
1498                )),
1499            },
1500            Ok,
1501        )?;
1502
1503        let transaction = Transaction::new(read_version, operation, None);
1504
1505        let mut builder = CommitBuilder::new(base_uri)
1506            .enable_v2_manifest_paths(enable_v2_manifest_paths)
1507            .with_session(session)
1508            .with_detached(detached);
1509
1510        if let Some(store_params) = store_params {
1511            builder = builder.with_store_params(store_params);
1512        }
1513
1514        if let Some(commit_handler) = commit_handler {
1515            builder = builder.with_commit_handler(commit_handler);
1516        }
1517
1518        builder.execute(transaction).await
1519    }
1520
1521    /// Commit changes to the dataset
1522    ///
1523    /// This operation is not needed if you are using append/write/delete to manipulate the dataset.
1524    /// It is used to commit changes to the dataset that are made externally.  For example, a bulk
1525    /// import tool may import large amounts of new data and write the appropriate lance files
1526    /// directly instead of using the write function.
1527    ///
1528    /// This method can be used to commit this change to the dataset's manifest.  This method will
1529    /// not verify that the provided fragments exist and correct, that is the caller's responsibility.
1530    /// Some validation can be performed using the function
1531    /// [crate::dataset::transaction::validate_operation].
1532    ///
1533    /// If this commit is a change to an existing dataset then it will often need to be based on an
1534    /// existing version of the dataset.  For example, if this change is a `delete` operation then
1535    /// the caller will have read in the existing data (at some version) to determine which fragments
1536    /// need to be deleted.  The base version that the caller used should be supplied as the `read_version`
1537    /// parameter.  Some operations (e.g. Overwrite) do not depend on a previous version and `read_version`
1538    /// can be None.  An error will be returned if the `read_version` is needed for an operation and
1539    /// it is not specified.
1540    ///
1541    /// All operations except Overwrite will fail if the dataset does not already exist.
1542    ///
1543    /// # Arguments
1544    ///
1545    /// * `base_uri` - The base URI of the dataset
1546    /// * `operation` - A description of the change to commit
1547    /// * `read_version` - The version of the dataset that this change is based on
1548    /// * `store_params` Parameters controlling object store access to the manifest
1549    /// * `enable_v2_manifest_paths`: If set to true, and this is a new dataset, uses the new v2 manifest
1550    ///   paths. These allow constant-time lookups for the latest manifest on object storage.
1551    ///   This parameter has no effect on existing datasets. To migrate an existing
1552    ///   dataset, use the [`Self::migrate_manifest_paths_v2`] method. WARNING: turning
1553    ///   this on will make the dataset unreadable for older versions of Lance
1554    ///   (prior to 0.17.0). Default is False.
1555    pub async fn commit(
1556        dest: impl Into<WriteDestination<'_>>,
1557        operation: Operation,
1558        read_version: Option<u64>,
1559        store_params: Option<ObjectStoreParams>,
1560        commit_handler: Option<Arc<dyn CommitHandler>>,
1561        session: Arc<Session>,
1562        enable_v2_manifest_paths: bool,
1563    ) -> Result<Self> {
1564        Self::do_commit(
1565            dest.into(),
1566            operation,
1567            read_version,
1568            store_params,
1569            commit_handler,
1570            session,
1571            enable_v2_manifest_paths,
1572            /*detached=*/ false,
1573        )
1574        .await
1575    }
1576
1577    /// Commits changes exactly the same as [`Self::commit`] but the commit will
1578    /// not be associated with the dataset lineage.
1579    ///
1580    /// The commit will not show up in the dataset's history and will never be
1581    /// the latest version of the dataset.
1582    ///
1583    /// This can be used to stage changes or to handle "secondary" datasets whose
1584    /// lineage is tracked elsewhere.
1585    pub async fn commit_detached(
1586        dest: impl Into<WriteDestination<'_>>,
1587        operation: Operation,
1588        read_version: Option<u64>,
1589        store_params: Option<ObjectStoreParams>,
1590        commit_handler: Option<Arc<dyn CommitHandler>>,
1591        session: Arc<Session>,
1592        enable_v2_manifest_paths: bool,
1593    ) -> Result<Self> {
1594        Self::do_commit(
1595            dest.into(),
1596            operation,
1597            read_version,
1598            store_params,
1599            commit_handler,
1600            session,
1601            enable_v2_manifest_paths,
1602            /*detached=*/ true,
1603        )
1604        .await
1605    }
1606
1607    pub(crate) async fn apply_commit(
1608        &mut self,
1609        transaction: Transaction,
1610        write_config: &ManifestWriteConfig,
1611        commit_config: &CommitConfig,
1612    ) -> Result<()> {
1613        let (manifest, manifest_location) = commit_transaction(
1614            self,
1615            self.object_store.as_ref(),
1616            self.commit_handler.as_ref(),
1617            &transaction,
1618            write_config,
1619            commit_config,
1620            self.manifest_location.naming_scheme,
1621            None,
1622        )
1623        .await?;
1624
1625        self.manifest = Arc::new(manifest);
1626        self.manifest_location = manifest_location;
1627        self.fragment_bitmap = Arc::new(
1628            self.manifest
1629                .fragments
1630                .iter()
1631                .map(|f| f.id as u32)
1632                .collect(),
1633        );
1634
1635        Ok(())
1636    }
1637
1638    /// Create a Scanner to scan the dataset.
1639    pub fn scan(&self) -> Scanner {
1640        Scanner::new(Arc::new(self.clone()))
1641    }
1642
1643    /// Count the number of rows in the dataset.
1644    ///
1645    /// It offers a fast path of counting rows by just computing via metadata.
1646    #[instrument(skip_all)]
1647    pub async fn count_rows(&self, filter: Option<String>) -> Result<usize> {
1648        // TODO: consolidate the count_rows into Scanner plan.
1649        if let Some(filter) = filter {
1650            let mut scanner = self.scan();
1651            scanner.filter(&filter)?;
1652            Ok(scanner
1653                .project::<String>(&[])?
1654                .with_row_id() // TODO: fix scan plan to not require row_id for count_rows.
1655                .count_rows()
1656                .await? as usize)
1657        } else {
1658            self.count_all_rows().await
1659        }
1660    }
1661
1662    pub(crate) async fn count_all_rows(&self) -> Result<usize> {
1663        let cnts = stream::iter(self.get_fragments())
1664            .map(|f| async move { f.count_rows(None).await })
1665            .buffer_unordered(16)
1666            .try_collect::<Vec<_>>()
1667            .await?;
1668        Ok(cnts.iter().sum())
1669    }
1670
1671    /// Take rows by indices.
1672    #[instrument(skip_all, fields(num_rows=row_indices.len()))]
1673    pub async fn take(
1674        &self,
1675        row_indices: &[u64],
1676        projection: impl Into<ProjectionRequest>,
1677    ) -> Result<RecordBatch> {
1678        take::take(self, row_indices, projection.into()).await
1679    }
1680
1681    /// Take Rows by the internal ROW ids.
1682    ///
1683    /// In Lance format, each row has a unique `u64` id, which is used to identify the row globally.
1684    ///
1685    /// ```rust
1686    /// # use std::sync::Arc;
1687    /// # use tokio::runtime::Runtime;
1688    /// # use arrow_array::{RecordBatch, RecordBatchIterator, Int64Array};
1689    /// # use arrow_schema::{Schema, Field, DataType};
1690    /// # use lance::dataset::{WriteParams, Dataset, ProjectionRequest};
1691    /// #
1692    /// # let mut rt = Runtime::new().unwrap();
1693    /// # rt.block_on(async {
1694    /// # let test_dir = tempfile::tempdir().unwrap();
1695    /// # let uri = test_dir.path().to_str().unwrap().to_string();
1696    /// #
1697    /// # let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
1698    /// # let write_params = WriteParams::default();
1699    /// # let array = Arc::new(Int64Array::from_iter(0..128));
1700    /// # let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
1701    /// # let reader = RecordBatchIterator::new(
1702    /// #    vec![batch].into_iter().map(Ok), schema
1703    /// # );
1704    /// # let dataset = Dataset::write(reader, &uri, Some(write_params)).await.unwrap();
1705    /// #
1706    /// let schema = dataset.schema().clone();
1707    /// let row_ids = vec![0, 4, 7];
1708    /// let rows = dataset.take_rows(&row_ids, schema).await.unwrap();
1709    ///
1710    /// // We can have more fine-grained control over the projection, i.e., SQL projection.
1711    /// let projection = ProjectionRequest::from_sql([("identity", "id * 2")]);
1712    /// let rows = dataset.take_rows(&row_ids, projection).await.unwrap();
1713    /// # });
1714    /// ```
1715    pub async fn take_rows(
1716        &self,
1717        row_ids: &[u64],
1718        projection: impl Into<ProjectionRequest>,
1719    ) -> Result<RecordBatch> {
1720        Arc::new(self.clone())
1721            .take_builder(row_ids, projection)?
1722            .execute()
1723            .await
1724    }
1725
1726    pub fn take_builder(
1727        self: &Arc<Self>,
1728        row_ids: &[u64],
1729        projection: impl Into<ProjectionRequest>,
1730    ) -> Result<TakeBuilder> {
1731        TakeBuilder::try_new_from_ids(self.clone(), row_ids.to_vec(), projection.into())
1732    }
1733
1734    /// Take [BlobFile] by row IDs.
1735    ///
1736    /// The returned vector has one element per row ID. Null blob values are
1737    /// represented as `None`; valid empty blobs return a `BlobFile` with size
1738    /// zero.
1739    ///
1740    /// ```
1741    /// # use std::sync::Arc;
1742    /// # use lance::dataset::Dataset;
1743    /// # use lance::Result;
1744    /// # async fn example(dataset: Arc<Dataset>) -> Result<()> {
1745    /// let blobs = dataset.take_blobs(&[42], "images").await?;
1746    /// match &blobs[0] {
1747    ///     None => { /* The selected blob is null. */ }
1748    ///     Some(blob) if blob.size() == 0 => { /* The selected blob is valid but empty. */ }
1749    ///     Some(blob) => { let _size = blob.size(); }
1750    /// }
1751    /// # Ok(())
1752    /// # }
1753    /// ```
1754    pub async fn take_blobs(
1755        self: &Arc<Self>,
1756        row_ids: &[u64],
1757        column: impl AsRef<str>,
1758    ) -> Result<Vec<Option<BlobFile>>> {
1759        blob::take_blobs(self, row_ids, column.as_ref()).await
1760    }
1761
1762    /// Take [BlobFile] by row addresses.
1763    ///
1764    /// Row addresses are `u64` values encoding `(fragment_id << 32) | row_offset`.
1765    /// Use this method when you already have row addresses, for example from
1766    /// a scan with `with_row_address()`. For row IDs (stable identifiers), use
1767    /// [`Self::take_blobs`]. For row indices (offsets), use
1768    /// [`Self::take_blobs_by_indices`]. The result has the same null and empty
1769    /// blob representation as [`Self::take_blobs`].
1770    ///
1771    /// ```
1772    /// # use std::sync::Arc;
1773    /// # use lance::dataset::Dataset;
1774    /// # use lance::Result;
1775    /// # async fn example(dataset: Arc<Dataset>, row_address: u64) -> Result<()> {
1776    /// let blobs = dataset
1777    ///     .take_blobs_by_addresses(&[row_address], "images")
1778    ///     .await?;
1779    /// match &blobs[0] {
1780    ///     None => { /* The selected blob is null. */ }
1781    ///     Some(blob) if blob.size() == 0 => { /* The selected blob is valid but empty. */ }
1782    ///     Some(blob) => { let _size = blob.size(); }
1783    /// }
1784    /// # Ok(())
1785    /// # }
1786    /// ```
1787    pub async fn take_blobs_by_addresses(
1788        self: &Arc<Self>,
1789        row_addrs: &[u64],
1790        column: impl AsRef<str>,
1791    ) -> Result<Vec<Option<BlobFile>>> {
1792        blob::take_blobs_by_addresses(self, row_addrs, column.as_ref()).await
1793    }
1794
1795    /// Take [BlobFile] by row indices (offsets in the dataset).
1796    ///
1797    /// The result has the same null and empty blob representation as
1798    /// [`Self::take_blobs`].
1799    ///
1800    /// ```
1801    /// # use std::sync::Arc;
1802    /// # use lance::dataset::Dataset;
1803    /// # use lance::Result;
1804    /// # async fn example(dataset: Arc<Dataset>) -> Result<()> {
1805    /// let blobs = dataset.take_blobs_by_indices(&[0], "images").await?;
1806    /// match &blobs[0] {
1807    ///     None => { /* The selected blob is null. */ }
1808    ///     Some(blob) if blob.size() == 0 => { /* The selected blob is valid but empty. */ }
1809    ///     Some(blob) => { let _size = blob.size(); }
1810    /// }
1811    /// # Ok(())
1812    /// # }
1813    /// ```
1814    pub async fn take_blobs_by_indices(
1815        self: &Arc<Self>,
1816        row_indices: &[u64],
1817        column: impl AsRef<str>,
1818    ) -> Result<Vec<Option<BlobFile>>> {
1819        let fragments = self.get_fragments();
1820        let row_addrs = row_offsets_to_row_addresses(&fragments, row_indices).await?;
1821        blob::take_blobs_by_addresses(self, &row_addrs, column.as_ref()).await
1822    }
1823
1824    /// Create a planned blob reader for a blob column.
1825    ///
1826    /// This API complements [`Self::take_blobs`]. `take_blobs` returns
1827    /// [`BlobFile`] handles for caller-driven random access, while
1828    /// `read_blobs` builds a streaming read plan for sequential or batched blob
1829    /// retrieval. Every selected row produces one result: null blob values have
1830    /// `ReadBlob::data` set to `None`, while valid empty blobs contain an empty
1831    /// buffer.
1832    ///
1833    /// ```rust
1834    /// # use std::sync::Arc;
1835    /// # use futures::TryStreamExt;
1836    /// # use lance::dataset::Dataset;
1837    /// # use lance::Result;
1838    /// # async fn example(dataset: Arc<Dataset>) -> Result<()> {
1839    /// let blobs = dataset
1840    ///     .read_blobs("images")?
1841    ///     .with_row_indices(vec![0, 1, 2])
1842    ///     .execute()
1843    ///     .await?;
1844    /// # let _ = blobs;
1845    /// # Ok(())
1846    /// # }
1847    /// ```
1848    pub fn read_blobs(self: &Arc<Self>, column: impl AsRef<str>) -> Result<ReadBlobsBuilder> {
1849        let column = column.as_ref();
1850        let blob_field_id = blob::validate_blob_column(self, column)?;
1851        Ok(ReadBlobsBuilder::new(
1852            self.clone(),
1853            column.to_string(),
1854            blob_field_id,
1855        ))
1856    }
1857
1858    /// Create a planned reader for row-specific blob-local byte ranges.
1859    ///
1860    /// Each [`BlobRangeRequest`] contains both its row selector and byte range,
1861    /// so requests can be repeated or reordered without coordinating parallel
1862    /// selector and range lists. Every request produces one result. A null blob
1863    /// has `ReadBlobRange::data` set to `None`; an empty range on a non-null blob
1864    /// contains an empty buffer.
1865    ///
1866    /// ```rust
1867    /// # use std::sync::Arc;
1868    /// # use lance::dataset::{BlobRangeRequest, Dataset};
1869    /// # use lance::Result;
1870    /// # async fn example(dataset: Arc<Dataset>) -> Result<()> {
1871    /// let ranges = dataset
1872    ///     .read_blob_ranges("images")?
1873    ///     .with_row_indices([
1874    ///         BlobRangeRequest::new(7, 0, 1024),
1875    ///         BlobRangeRequest::new(7, 4096, 1024),
1876    ///     ])
1877    ///     .execute()
1878    ///     .await?;
1879    /// # let _ = ranges;
1880    /// # Ok(())
1881    /// # }
1882    /// ```
1883    pub fn read_blob_ranges(
1884        self: &Arc<Self>,
1885        column: impl AsRef<str>,
1886    ) -> Result<ReadBlobRangesBuilder> {
1887        Ok(ReadBlobRangesBuilder::new(self.read_blobs(column)?))
1888    }
1889
1890    /// Get a stream of batches based on iterator of ranges of row numbers.
1891    ///
1892    /// This is an experimental API. It may change at any time.
1893    pub fn take_scan(
1894        &self,
1895        row_ranges: Pin<Box<dyn Stream<Item = Result<Range<u64>>> + Send>>,
1896        projection: Arc<Schema>,
1897        batch_readahead: usize,
1898    ) -> DatasetRecordBatchStream {
1899        take::take_scan(self, row_ranges, projection, batch_readahead)
1900    }
1901
1902    /// Randomly sample `n` rows from the dataset.
1903    ///
1904    /// If `fragment_ids` is provided, sampling is limited to rows from those
1905    /// fragments in the current dataset version.
1906    ///
1907    /// The returned rows are in row-id order (not random order), which allows
1908    /// the underlying take operation to use an efficient sorted code path.
1909    pub async fn sample(
1910        &self,
1911        n: usize,
1912        projection: &Schema,
1913        fragment_ids: Option<&[u32]>,
1914    ) -> Result<RecordBatch> {
1915        use rand::seq::IteratorRandom;
1916
1917        match fragment_ids {
1918            None => {
1919                let num_rows = self.count_rows(None).await?;
1920                let mut ids = (0..num_rows as u64).choose_multiple(&mut rand::rng(), n);
1921                ids.sort_unstable();
1922                self.take(&ids, projection.clone()).await
1923            }
1924            Some(fragment_ids) => {
1925                if fragment_ids.is_empty() {
1926                    return Err(Error::invalid_input(
1927                        "Dataset::sample does not accept an empty fragment_ids list".to_string(),
1928                    ));
1929                }
1930
1931                let selected_fragments = self.get_fragments_from_ids(fragment_ids)?;
1932
1933                let num_rows = stream::iter(selected_fragments.iter().cloned())
1934                    .map(|fragment| async move { fragment.count_rows(None).await })
1935                    .buffer_unordered(16)
1936                    .try_fold(0_u64, |acc, rows| async move { Ok(acc + rows as u64) })
1937                    .await?;
1938
1939                let mut offsets = (0..num_rows).choose_multiple(&mut rand::rng(), n);
1940                offsets.sort_unstable();
1941
1942                let row_addrs = row_offsets_to_row_addresses(&selected_fragments, &offsets).await?;
1943                let dataset = Arc::new(self.clone());
1944                let projection = Arc::new(
1945                    ProjectionRequest::from(projection.clone())
1946                        .into_projection_plan(dataset.clone())?,
1947                );
1948                TakeBuilder::try_new_from_addresses(dataset, row_addrs, projection)?
1949                    .execute()
1950                    .await
1951            }
1952        }
1953    }
1954
1955    /// Delete rows based on a predicate.
1956    pub async fn delete(&mut self, predicate: &str) -> Result<write::delete::DeleteResult> {
1957        info!(target: TRACE_DATASET_EVENTS, event=DATASET_DELETING_EVENT, uri = &self.uri, predicate=predicate);
1958        write::delete::delete(self, predicate).await
1959    }
1960
1961    /// Truncate the dataset by deleting all rows.
1962    pub async fn truncate_table(&mut self) -> Result<()> {
1963        self.delete("true").await.map(|_| ())
1964    }
1965
1966    /// Add new base paths to the dataset.
1967    ///
1968    /// This method allows you to register additional storage locations (buckets)
1969    /// that can be used for future data writes. The base paths are added to the
1970    /// dataset's manifest and can be referenced by name in subsequent write operations.
1971    ///
1972    /// # Arguments
1973    ///
1974    /// * `new_bases` - A vector of `lance_table::format::BasePath` objects representing the new storage
1975    ///   locations to add. Each base path should have a unique name and path.
1976    ///
1977    /// # Returns
1978    ///
1979    /// Returns a new `Dataset` instance with the updated manifest containing the
1980    /// new base paths.
1981    pub async fn add_bases(
1982        self: &Arc<Self>,
1983        new_bases: Vec<lance_table::format::BasePath>,
1984        transaction_properties: Option<HashMap<String, String>>,
1985    ) -> Result<Self> {
1986        let operation = Operation::UpdateBases { new_bases };
1987
1988        let transaction = TransactionBuilder::new(self.manifest.version, operation)
1989            .transaction_properties(transaction_properties.map(Arc::new))
1990            .build();
1991
1992        let new_dataset = CommitBuilder::new(self.clone())
1993            .execute(transaction)
1994            .await?;
1995
1996        Ok(new_dataset)
1997    }
1998
1999    pub async fn count_deleted_rows(&self) -> Result<usize> {
2000        futures::stream::iter(self.get_fragments())
2001            .map(|f| async move { f.count_deletions().await })
2002            .buffer_unordered(self.object_store.io_parallelism())
2003            .try_fold(0, |acc, x| futures::future::ready(Ok(acc + x)))
2004            .await
2005    }
2006
2007    /// Clone this dataset with a different object store binding.
2008    ///
2009    /// The returned dataset shares metadata, session state, and caches with the
2010    /// original dataset, but all subsequent operations on the returned dataset
2011    /// use the supplied object store.
2012    pub fn with_object_store(
2013        &self,
2014        object_store: Arc<ObjectStore>,
2015        store_params: Option<ObjectStoreParams>,
2016    ) -> Self {
2017        let mut cloned = self.clone();
2018        cloned.object_store = object_store;
2019        if let Some(store_params) = store_params {
2020            cloned.store_params = Some(Box::new(store_params));
2021        }
2022        cloned
2023    }
2024
2025    /// Clone this dataset with extra object store wrappers applied to all read stores.
2026    ///
2027    /// The returned dataset uses the wrappers for the already-open primary object
2028    /// store and appends the same wrappers to the dataset-level and base-specific
2029    /// object store params used when additional base stores are opened later.
2030    pub fn with_object_store_wrappers(
2031        &self,
2032        wrappers: impl IntoIterator<Item = Arc<dyn WrappingObjectStore>>,
2033    ) -> Self {
2034        let wrappers = wrappers.into_iter().collect::<Vec<_>>();
2035        if wrappers.is_empty() {
2036            return self.clone();
2037        }
2038
2039        let mut cloned = self.clone();
2040        let mut object_store = self.object_store.as_ref().clone();
2041        for wrapper in &wrappers {
2042            object_store.inner =
2043                wrapper.wrap(&object_store.store_prefix, object_store.inner.clone());
2044        }
2045        cloned.object_store = Arc::new(object_store);
2046        cloned.refs = Refs::new(
2047            cloned.object_store.clone(),
2048            cloned.commit_handler.clone(),
2049            cloned.branch_location(),
2050        );
2051
2052        let store_params = self.store_params.as_deref().cloned().unwrap_or_default();
2053        cloned.store_params = Some(Box::new(Self::append_object_store_wrappers(
2054            store_params,
2055            &wrappers,
2056        )));
2057        cloned.base_store_params = self.base_store_params.as_ref().map(|base_store_params| {
2058            Arc::new(
2059                base_store_params
2060                    .iter()
2061                    .map(|(base_path, store_params)| {
2062                        (
2063                            base_path.clone(),
2064                            Self::append_object_store_wrappers(store_params.clone(), &wrappers),
2065                        )
2066                    })
2067                    .collect(),
2068            )
2069        });
2070        cloned
2071    }
2072
2073    fn append_object_store_wrappers(
2074        mut store_params: ObjectStoreParams,
2075        wrappers: &[Arc<dyn WrappingObjectStore>],
2076    ) -> ObjectStoreParams {
2077        let mut all_wrappers = Vec::with_capacity(
2078            store_params.object_store_wrapper.as_ref().map_or(0, |_| 1) + wrappers.len(),
2079        );
2080        if let Some(wrapper) = store_params.object_store_wrapper.take() {
2081            all_wrappers.push(wrapper);
2082        }
2083        all_wrappers.extend(wrappers.iter().cloned());
2084        store_params.object_store_wrapper = match all_wrappers.len() {
2085            0 => None,
2086            1 => all_wrappers.pop(),
2087            _ => Some(Arc::new(ChainedWrappingObjectStore::new(all_wrappers))),
2088        };
2089        store_params
2090    }
2091
2092    pub(crate) fn store_params_for_base(
2093        &self,
2094        base_path: Option<&lance_table::format::BasePath>,
2095    ) -> ObjectStoreParams {
2096        // Base-specific bindings are exact ObjectStoreParams keyed by
2097        // `BasePath.path` and are used as-is. Otherwise the dataset-level
2098        // default params are resolved for the base scope: `base_<id>.<key>`
2099        // storage options overlay the shared defaults for that base.
2100        if let Some(params) = base_path.and_then(|base_path| {
2101            self.base_store_params
2102                .as_ref()
2103                .and_then(|params| params.get(&base_path.path))
2104        }) {
2105            return params.clone();
2106        }
2107        let default_params = self.store_params.as_deref().cloned().unwrap_or_default();
2108        match default_params.scoped_to_base(base_path.map(|base_path| base_path.id)) {
2109            Cow::Owned(scoped_params) => scoped_params,
2110            Cow::Borrowed(_) => default_params,
2111        }
2112    }
2113
2114    /// Returns the initial storage options used when opening this dataset, if any.
2115    ///
2116    /// This returns the static initial options without triggering any refresh.
2117    /// For the latest refreshed options, use [`Self::latest_storage_options`].
2118    #[deprecated(since = "0.25.0", note = "Use initial_storage_options() instead")]
2119    pub fn storage_options(&self) -> Option<&HashMap<String, String>> {
2120        self.initial_storage_options()
2121    }
2122
2123    /// Returns the initial storage options without triggering any refresh.
2124    ///
2125    /// For the latest refreshed options, use [`Self::latest_storage_options`].
2126    pub fn initial_storage_options(&self) -> Option<&HashMap<String, String>> {
2127        self.store_params
2128            .as_ref()
2129            .and_then(|params| params.storage_options())
2130    }
2131
2132    /// Returns the storage options provider used when opening this dataset, if any.
2133    pub fn storage_options_provider(
2134        &self,
2135    ) -> Option<Arc<dyn lance_io::object_store::StorageOptionsProvider>> {
2136        self.store_params
2137            .as_ref()
2138            .and_then(|params| params.storage_options_accessor.as_ref())
2139            .and_then(|accessor| accessor.provider().cloned())
2140    }
2141
2142    /// Returns the unified storage options accessor for this dataset, if any.
2143    ///
2144    /// The accessor handles both static and dynamic storage options with automatic
2145    /// caching and refresh. Use [`StorageOptionsAccessor::get_storage_options`] to
2146    /// get the latest options.
2147    pub fn storage_options_accessor(&self) -> Option<Arc<StorageOptionsAccessor>> {
2148        self.store_params
2149            .as_ref()
2150            .and_then(|params| params.get_accessor())
2151    }
2152
2153    /// Returns the latest (possibly refreshed) storage options.
2154    ///
2155    /// If a dynamic storage options provider is configured, this will return
2156    /// the cached options if still valid, or fetch fresh options if expired.
2157    ///
2158    /// For the initial static options without refresh, use [`Self::storage_options`].
2159    ///
2160    /// # Returns
2161    ///
2162    /// - `Ok(Some(options))` - Storage options are available (static or refreshed)
2163    /// - `Ok(None)` - No storage options were configured for this dataset
2164    /// - `Err(...)` - Error occurred while fetching/refreshing options from provider
2165    pub async fn latest_storage_options(&self) -> Result<Option<StorageOptions>> {
2166        // First check if we have an accessor (handles both static and dynamic options)
2167        if let Some(accessor) = self.storage_options_accessor() {
2168            let options = accessor.get_storage_options().await?;
2169            return Ok(Some(options));
2170        }
2171
2172        // Fallback to initial storage options if no accessor
2173        Ok(self.initial_storage_options().cloned().map(StorageOptions))
2174    }
2175
2176    pub fn data_dir(&self) -> Path {
2177        self.base.clone().join(DATA_DIR)
2178    }
2179
2180    pub fn indices_dir(&self) -> Path {
2181        self.base.clone().join(INDICES_DIR)
2182    }
2183
2184    pub fn transactions_dir(&self) -> Path {
2185        self.base.clone().join(TRANSACTIONS_DIR)
2186    }
2187
2188    pub fn deletions_dir(&self) -> Path {
2189        self.base.clone().join(DELETIONS_DIR)
2190    }
2191
2192    pub fn versions_dir(&self) -> Path {
2193        self.base.clone().join(VERSIONS_DIR)
2194    }
2195
2196    pub(crate) fn data_file_dir(&self, data_file: &DataFile) -> Result<Path> {
2197        self.data_file_dir_for_base(data_file.base_id)
2198    }
2199
2200    /// Create a [`DataFile`] by reading metadata from an existing lance file.
2201    ///
2202    /// This reads the file's schema and version information, matches columns to
2203    /// the dataset's schema to determine field IDs, and calculates column indices.
2204    /// This is useful for constructing `DataFile` metadata needed for operations
2205    /// like [`Operation::DataReplacement`].
2206    ///
2207    /// # Arguments
2208    ///
2209    /// * `path` - The path to the data file, relative to the dataset's data directory.
2210    /// * `base_id` - The base path ID if the file is outside the dataset directory.
2211    pub async fn create_data_file(&self, path: &str, base_id: Option<u32>) -> Result<DataFile> {
2212        let data_dir = self.data_file_dir_for_base(base_id)?;
2213        let filepath = data_dir.clone().join(path);
2214
2215        let object_store = self.object_store(base_id).await?;
2216
2217        // Get file size
2218        let file_size = object_store.size(&filepath).await?;
2219
2220        // Read file metadata
2221        let scheduler = ScanScheduler::new(
2222            object_store.clone(),
2223            SchedulerConfig::new(2 * 1024 * 1024 * 1024),
2224        );
2225        let file = scheduler
2226            .open_file(&filepath, &CachedFileSize::new(file_size))
2227            .await?;
2228        let file_metadata = FileReader::read_all_metadata(&file).await?;
2229
2230        let lance_file_format = ConcreteFileVersion::from_footer_numbers(
2231            file_metadata.major_version,
2232            file_metadata.minor_version,
2233        )?;
2234        let file_version: LanceFileVersion = lance_file_format.into();
2235
2236        let is_structural = file_version >= LanceFileVersion::V2_1;
2237        let physical_columns = file_metadata.column_metadatas.len();
2238        let has_footer_orphans = file_metadata.file_schema.fields.len() > physical_columns;
2239        let dataset_schema = self.schema();
2240        let mut represented_columns = 0usize;
2241        let mut column_names = Vec::new();
2242        let mut consumed_top_level_fields = 0usize;
2243
2244        fn physical_column_count(
2245            field: &lance_core::datatypes::Field,
2246            is_structural: bool,
2247        ) -> usize {
2248            if !is_structural {
2249                return 1 + field
2250                    .children
2251                    .iter()
2252                    .map(|child| physical_column_count(child, is_structural))
2253                    .sum::<usize>();
2254            }
2255
2256            if field.children.is_empty() || field.is_blob() || field.is_packed_struct() {
2257                1
2258            } else {
2259                field
2260                    .children
2261                    .iter()
2262                    .map(|child| physical_column_count(child, is_structural))
2263                    .sum()
2264            }
2265        }
2266
2267        fn field_contains_blob(field: &lance_core::datatypes::Field) -> bool {
2268            field.is_blob() || field.children.iter().any(field_contains_blob)
2269        }
2270
2271        fn field_names_match(
2272            fields: &[lance_core::datatypes::Field],
2273            start: usize,
2274            names: &[&str],
2275        ) -> bool {
2276            fields
2277                .get(start..start + names.len())
2278                .is_some_and(|candidate| {
2279                    candidate
2280                        .iter()
2281                        .zip(names)
2282                        .all(|(field, name)| field.name == *name)
2283                })
2284        }
2285
2286        fn blob_descriptor_orphan_len(
2287            fields: &[lance_core::datatypes::Field],
2288            start: usize,
2289        ) -> usize {
2290            const BLOB_V2_DESCRIPTOR_FIELDS: &[&str] =
2291                &["kind", "position", "size", "blob_id", "blob_uri"];
2292            const BLOB_V1_DESCRIPTOR_FIELDS: &[&str] = &["position", "size"];
2293
2294            if field_names_match(fields, start, BLOB_V2_DESCRIPTOR_FIELDS) {
2295                BLOB_V2_DESCRIPTOR_FIELDS.len()
2296            } else if field_names_match(fields, start, BLOB_V1_DESCRIPTOR_FIELDS) {
2297                BLOB_V1_DESCRIPTOR_FIELDS.len()
2298            } else {
2299                0
2300            }
2301        }
2302
2303        fn collect_columns(
2304            field: &lance_core::datatypes::Field,
2305            is_structural: bool,
2306            fields: &mut Vec<i32>,
2307            column_indices: &mut Vec<i32>,
2308            curr_column_idx: &mut i32,
2309        ) {
2310            let contributes = !is_structural
2311                || field.children.is_empty()
2312                || field.is_blob()
2313                || field.is_packed_struct();
2314            let recurse = !is_structural || (!field.is_blob() && !field.is_packed_struct());
2315
2316            if contributes {
2317                fields.push(field.id);
2318                column_indices.push(*curr_column_idx);
2319                *curr_column_idx += 1;
2320            }
2321
2322            if recurse {
2323                for child in &field.children {
2324                    collect_columns(
2325                        child,
2326                        is_structural,
2327                        fields,
2328                        column_indices,
2329                        curr_column_idx,
2330                    );
2331                }
2332            }
2333        }
2334
2335        fn validate_file_field_matches_dataset(
2336            dataset_field: &lance_core::datatypes::Field,
2337            file_field: &lance_core::datatypes::Field,
2338            path: &str,
2339        ) -> Result<()> {
2340            if dataset_field.name != file_field.name {
2341                return Err(Error::invalid_input(format!(
2342                    "Schema mismatch: expected field '{}' but file has '{}'",
2343                    path, file_field.name
2344                )));
2345            }
2346
2347            if dataset_field.is_blob() && file_field.is_blob() {
2348                return Ok(());
2349            }
2350
2351            if dataset_field.children.len() != file_field.children.len() {
2352                return Err(Error::invalid_input(format!(
2353                    "Schema mismatch: field '{}' has {} children in dataset schema but {} children in file schema",
2354                    path,
2355                    dataset_field.children.len(),
2356                    file_field.children.len()
2357                )));
2358            }
2359
2360            for (dataset_child, file_child) in
2361                dataset_field.children.iter().zip(&file_field.children)
2362            {
2363                let child_path = format!("{}.{}", path, dataset_child.name);
2364                validate_file_field_matches_dataset(dataset_child, file_child, &child_path)?;
2365            }
2366
2367            Ok(())
2368        }
2369
2370        let file_schema_fields = &file_metadata.file_schema.fields;
2371        let mut idx = 0usize;
2372        while represented_columns < physical_columns {
2373            let Some(field) = file_schema_fields.get(idx) else {
2374                return Err(Error::invalid_input(format!(
2375                    "Schema mismatch: file schema ended after representing {} physical columns but file has {} columns",
2376                    represented_columns, physical_columns
2377                )));
2378            };
2379
2380            let Some(dataset_field) = dataset_schema.field(&field.name) else {
2381                return Err(Error::invalid_input(format!(
2382                    "Schema mismatch: file has extra field '{}'",
2383                    field.name
2384                )));
2385            };
2386            validate_file_field_matches_dataset(dataset_field, field, &field.name)?;
2387
2388            represented_columns += physical_column_count(field, is_structural);
2389            column_names.push(field.name.as_str());
2390            consumed_top_level_fields = idx + 1;
2391            idx += 1;
2392
2393            if has_footer_orphans && field_contains_blob(field) {
2394                loop {
2395                    let skipped = blob_descriptor_orphan_len(file_schema_fields, idx);
2396                    if skipped == 0 {
2397                        break;
2398                    }
2399                    consumed_top_level_fields = idx + skipped;
2400                    idx += skipped;
2401                }
2402            }
2403        }
2404
2405        if represented_columns != physical_columns {
2406            return Err(Error::invalid_input(format!(
2407                "Schema mismatch: file schema represents {} physical columns but file has {} columns",
2408                represented_columns, physical_columns
2409            )));
2410        }
2411
2412        if let Some(field) = file_schema_fields.get(consumed_top_level_fields) {
2413            return Err(Error::invalid_input(format!(
2414                "Schema mismatch: file has extra field '{}'",
2415                field.name
2416            )));
2417        }
2418
2419        let projected_ds_schema = self.schema().project(&column_names)?;
2420
2421        let mut fields = Vec::new();
2422        let mut column_indices = Vec::new();
2423        let mut curr_column_idx: i32 = 0;
2424        for field in &projected_ds_schema.fields {
2425            collect_columns(
2426                field,
2427                is_structural,
2428                &mut fields,
2429                &mut column_indices,
2430                &mut curr_column_idx,
2431            );
2432        }
2433
2434        if curr_column_idx as usize != physical_columns {
2435            return Err(Error::invalid_input(format!(
2436                "Schema mismatch: dataset projection maps to {} physical columns but file has {} columns",
2437                curr_column_idx, physical_columns
2438            )));
2439        }
2440
2441        if fields.is_empty() && physical_columns > 0 {
2442            return Err(Error::invalid_input(
2443                "Schema mismatch: file has columns but none matched the dataset schema",
2444            ));
2445        }
2446
2447        let file_size_nz = NonZero::new(file_size);
2448        Ok(DataFile::new(
2449            path,
2450            fields,
2451            column_indices,
2452            lance_file_format,
2453            file_size_nz,
2454            base_id,
2455        ))
2456    }
2457
2458    /// Resolve the data directory for a given base_id.
2459    ///
2460    /// If `base_id` is `None`, returns the default data directory.
2461    pub(crate) fn data_file_dir_for_base(&self, base_id: Option<u32>) -> Result<Path> {
2462        match base_id {
2463            Some(base_id) => {
2464                let base_path = self.manifest.base_paths.get(&base_id).ok_or_else(|| {
2465                    Error::invalid_input(format!("base_path id {} not found", base_id))
2466                })?;
2467                let path = base_path.extract_path(self.session.store_registry())?;
2468                if base_path.is_dataset_root {
2469                    Ok(path.join(DATA_DIR))
2470                } else {
2471                    Ok(path)
2472                }
2473            }
2474            None => Ok(self.base.clone().join(DATA_DIR)),
2475        }
2476    }
2477
2478    async fn base_object_store(&self, base_id: u32) -> Result<Arc<ObjectStore>> {
2479        let base_path = self.manifest.base_paths.get(&base_id).ok_or_else(|| {
2480            Error::invalid_input(format!("Dataset base path with ID {} not found", base_id))
2481        })?;
2482        let store_params = self.store_params_for_base(Some(base_path));
2483
2484        let (store, _) = ObjectStore::from_uri_and_params(
2485            self.session.store_registry(),
2486            &base_path.path,
2487            &store_params,
2488        )
2489        .await?;
2490
2491        Ok(store)
2492    }
2493
2494    /// Resolve the object store for the primary dataset or an additional base.
2495    ///
2496    /// Pass `None` to get the primary dataset object store. Pass `Some(base_id)`
2497    /// when resolving a file whose metadata references an additional base.
2498    pub async fn object_store(&self, base_id: Option<u32>) -> Result<Arc<ObjectStore>> {
2499        match base_id {
2500            Some(base_id) => self.base_object_store(base_id).await,
2501            None => Ok(self.object_store.clone()),
2502        }
2503    }
2504
2505    /// The `ObjectStoreParams` this dataset was opened with, or `None` when
2506    /// opened without explicit params. Lets a caller re-open a derived path
2507    /// (e.g. a MemWAL SSTable) with the same store this dataset used.
2508    pub fn store_params(&self) -> Option<&ObjectStoreParams> {
2509        self.store_params.as_deref()
2510    }
2511
2512    pub(crate) async fn object_store_for_data_file(
2513        &self,
2514        data_file: &DataFile,
2515    ) -> Result<Arc<ObjectStore>> {
2516        self.object_store(data_file.base_id).await
2517    }
2518
2519    pub(crate) async fn object_store_for_deletion(
2520        &self,
2521        deletion_file: &DeletionFile,
2522    ) -> Result<Arc<ObjectStore>> {
2523        self.object_store(deletion_file.base_id).await
2524    }
2525
2526    pub(crate) async fn object_store_for_index(
2527        &self,
2528        index: &IndexMetadata,
2529    ) -> Result<Arc<ObjectStore>> {
2530        self.object_store(index.base_id).await
2531    }
2532
2533    pub(crate) fn dataset_dir_for_deletion(&self, deletion_file: &DeletionFile) -> Result<Path> {
2534        match deletion_file.base_id.as_ref() {
2535            Some(base_id) => {
2536                let base_paths = &self.manifest.base_paths;
2537                let base_path = base_paths.get(base_id).ok_or_else(|| {
2538                    Error::invalid_input(format!(
2539                        "base_path id {} not found for deletion_file {:?}",
2540                        base_id, deletion_file
2541                    ))
2542                })?;
2543
2544                if !base_path.is_dataset_root {
2545                    return Err(Error::internal(format!(
2546                        "base_path id {} is not a dataset root for deletion_file {:?}",
2547                        base_id, deletion_file
2548                    )));
2549                }
2550                base_path.extract_path(self.session.store_registry())
2551            }
2552            None => Ok(self.base.clone()),
2553        }
2554    }
2555
2556    /// Get the indices directory for a specific index, considering its base_id
2557    pub(crate) fn indice_files_dir(&self, index: &IndexMetadata) -> Result<Path> {
2558        match index.base_id.as_ref() {
2559            Some(base_id) => {
2560                let base_paths = &self.manifest.base_paths;
2561                let base_path = base_paths.get(base_id).ok_or_else(|| {
2562                    Error::invalid_input(format!(
2563                        "base_path id {} not found for index {}",
2564                        base_id, index.uuid
2565                    ))
2566                })?;
2567                let path = base_path.extract_path(self.session.store_registry())?;
2568                if base_path.is_dataset_root {
2569                    Ok(path.join(INDICES_DIR))
2570                } else {
2571                    // For non-dataset-root base paths, we assume the path already points to the indices directory
2572                    Ok(path)
2573                }
2574            }
2575            None => Ok(self.base.clone().join(INDICES_DIR)),
2576        }
2577    }
2578
2579    pub fn session(&self) -> Arc<Session> {
2580        self.session.clone()
2581    }
2582
2583    /// Get the currently checked-out version id.
2584    ///
2585    /// This is a cheap accessor that reads the id directly from the loaded
2586    /// manifest without constructing the full [Version] summary.
2587    pub fn version_id(&self) -> u64 {
2588        self.manifest.version
2589    }
2590
2591    /// Get the currently checked-out version details.
2592    ///
2593    /// This constructs a full [Version], including summary metadata derived
2594    /// from the loaded manifest fragments.
2595    pub fn version(&self) -> Version {
2596        Version::from(self.manifest.as_ref())
2597    }
2598
2599    /// Get the number of entries currently in the index cache.
2600    pub async fn index_cache_entry_count(&self) -> usize {
2601        self.session.index_cache.size().await
2602    }
2603
2604    /// Get cache hit ratio.
2605    pub async fn index_cache_hit_rate(&self) -> f32 {
2606        let stats = self.session.index_cache_stats().await;
2607        stats.hit_ratio()
2608    }
2609
2610    pub fn cache_size_bytes(&self) -> u64 {
2611        self.session.deep_size_of() as u64
2612    }
2613
2614    /// Get all versions.
2615    pub async fn versions(&self) -> Result<Vec<Version>> {
2616        let mut versions: Vec<Version> = self
2617            .commit_handler
2618            .list_manifest_locations(&self.base, &self.object_store, false)
2619            .try_filter_map(|location| async move {
2620                match read_manifest(&self.object_store, &location.path, location.size).await {
2621                    Ok(manifest) => Ok(Some(Version::from(&manifest))),
2622                    Err(e) => Err(e),
2623                }
2624            })
2625            .try_collect()
2626            .await?;
2627
2628        // TODO: this API should support pagination
2629        versions.sort_by_key(|v| v.version);
2630
2631        Ok(versions)
2632    }
2633
2634    /// List all detached manifest locations.
2635    ///
2636    /// Detached manifests are versions that are not part of the main version history.
2637    /// They are created by `commit_detached` and can be used for staging changes.
2638    ///
2639    /// To read transaction properties from a detached manifest:
2640    /// ```ignore
2641    /// let detached = dataset.list_detached_manifests().await?;
2642    /// for location in detached {
2643    ///     let ds = dataset.checkout_version(location.version).await?;
2644    ///     let tx = ds.read_transaction().await?;
2645    ///     // Access tx.transaction_properties
2646    /// }
2647    /// ```
2648    pub async fn list_detached_manifests(&self) -> Result<Vec<ManifestLocation>> {
2649        self.commit_handler
2650            .list_detached_manifest_locations(&self.base, &self.object_store)
2651            .try_collect()
2652            .await
2653    }
2654
2655    /// Get the latest version of the dataset
2656    /// This is meant to be a fast path for checking if a dataset has changed. This is why
2657    /// we don't return the full version struct.
2658    pub async fn latest_version_id(&self) -> Result<u64> {
2659        Ok(self
2660            .commit_handler
2661            .resolve_latest_location(&self.base, &self.object_store)
2662            .await?
2663            .version)
2664    }
2665
2666    /// Return whether the dataset has a newer committed version.
2667    pub async fn is_stale(&self) -> Result<bool> {
2668        let latest_version = self.latest_version_id().await?;
2669        Ok(latest_version != self.manifest.version)
2670    }
2671
2672    /// Return whether the immediate attached successor manifest exists.
2673    ///
2674    /// This is a fast contiguous-history probe. It does not resolve the latest
2675    /// version and may return `false` if intermediate manifests have been
2676    /// removed. Callers that need a general freshness check should use
2677    /// [`Self::is_stale`].
2678    #[doc(hidden)]
2679    pub async fn has_successor_version(&self) -> Result<bool> {
2680        let Some(next_version) = self.manifest.version.checked_add(1) else {
2681            return Ok(false);
2682        };
2683        if lance_table::format::is_detached_version(next_version) {
2684            return Ok(false);
2685        }
2686
2687        let exists = self
2688            .commit_handler
2689            .version_exists(
2690                &self.base,
2691                next_version,
2692                self.object_store.inner.as_ref(),
2693                self.manifest_location.naming_scheme,
2694            )
2695            .await?;
2696        Ok(exists)
2697    }
2698
2699    pub fn count_fragments(&self) -> usize {
2700        self.manifest.fragments.len()
2701    }
2702
2703    /// Get the schema of the dataset
2704    pub fn schema(&self) -> &Schema {
2705        &self.manifest.schema
2706    }
2707
2708    /// Similar to [Self::schema], but only returns fields that are not marked as blob columns
2709    /// Creates a new empty projection into the dataset schema
2710    pub fn empty_projection(self: &Arc<Self>) -> Projection {
2711        Projection::empty(self.clone())
2712    }
2713
2714    /// Creates a projection that includes all columns in the dataset
2715    pub fn full_projection(self: &Arc<Self>) -> Projection {
2716        Projection::full(self.clone())
2717    }
2718
2719    /// Get fragments.
2720    pub fn get_fragments(&self) -> Vec<FileFragment> {
2721        let dataset = Arc::new(self.clone());
2722        self.manifest
2723            .fragments
2724            .iter()
2725            .map(|f| FileFragment::new(dataset.clone(), f.clone()))
2726            .collect()
2727    }
2728
2729    /// Iterate over manifest fragments without allocating [`FileFragment`] wrappers.
2730    pub fn iter_fragments(&self) -> impl Iterator<Item = &Fragment> {
2731        self.manifest.fragments.iter()
2732    }
2733
2734    pub fn get_fragment(&self, fragment_id: usize) -> Option<FileFragment> {
2735        let dataset = Arc::new(self.clone());
2736        let fragment = self
2737            .manifest
2738            .fragments
2739            .iter()
2740            .find(|f| f.id == fragment_id as u64)?;
2741        Some(FileFragment::new(dataset, fragment.clone()))
2742    }
2743
2744    pub fn fragments(&self) -> &Arc<Vec<Fragment>> {
2745        &self.manifest.fragments
2746    }
2747
2748    pub(crate) fn normalize_fragment_ids(fragment_ids: &[u32]) -> Vec<u32> {
2749        let mut ids = fragment_ids.to_vec();
2750        ids.sort_unstable();
2751        ids.dedup();
2752        ids
2753    }
2754
2755    pub(crate) fn get_fragments_from_ids(&self, fragment_ids: &[u32]) -> Result<Vec<FileFragment>> {
2756        let ordered_ids = Self::normalize_fragment_ids(fragment_ids);
2757        let fragments = self.get_frags_from_ordered_ids(&ordered_ids);
2758        if let Some(missing_id) = fragments
2759            .iter()
2760            .zip(ordered_ids.iter())
2761            .find_map(|(fragment, fragment_id)| fragment.is_none().then_some(*fragment_id))
2762        {
2763            return Err(Error::invalid_input(format!(
2764                "Unknown fragment id {missing_id} in fragment filter; not part of the current dataset version"
2765            )));
2766        }
2767
2768        Ok(fragments.into_iter().flatten().collect())
2769    }
2770
2771    pub(crate) fn get_existing_fragments_from_ids(
2772        &self,
2773        fragment_ids: &[u32],
2774    ) -> Vec<FileFragment> {
2775        let ordered_ids = Self::normalize_fragment_ids(fragment_ids);
2776        self.get_frags_from_ordered_ids(&ordered_ids)
2777            .into_iter()
2778            .flatten()
2779            .collect()
2780    }
2781
2782    pub(crate) fn get_fragment_metadata_from_ids(
2783        &self,
2784        fragment_ids: &[u32],
2785    ) -> Result<Vec<Fragment>> {
2786        Ok(self
2787            .get_fragments_from_ids(fragment_ids)?
2788            .into_iter()
2789            .map(|fragment| fragment.metadata().clone())
2790            .collect())
2791    }
2792
2793    pub(crate) fn get_existing_fragment_metadata_from_ids(
2794        &self,
2795        fragment_ids: &[u32],
2796    ) -> Vec<Fragment> {
2797        self.get_existing_fragments_from_ids(fragment_ids)
2798            .into_iter()
2799            .map(|fragment| fragment.metadata().clone())
2800            .collect()
2801    }
2802
2803    pub(crate) async fn count_rows_in_fragments(&self, fragment_ids: &[u32]) -> Result<usize> {
2804        let fragments = self.get_fragments_from_ids(fragment_ids)?;
2805        self.count_rows_in_resolved_fragments(fragments).await
2806    }
2807
2808    pub(crate) async fn count_rows_in_existing_fragments(
2809        &self,
2810        fragment_ids: &[u32],
2811    ) -> Result<usize> {
2812        let fragments = self.get_existing_fragments_from_ids(fragment_ids);
2813        self.count_rows_in_resolved_fragments(fragments).await
2814    }
2815
2816    async fn count_rows_in_resolved_fragments(
2817        &self,
2818        fragments: Vec<FileFragment>,
2819    ) -> Result<usize> {
2820        let counts = stream::iter(fragments)
2821            .map(|fragment| async move { fragment.count_rows(None).await })
2822            .buffer_unordered(16)
2823            .try_collect::<Vec<_>>()
2824            .await?;
2825        Ok(counts.iter().sum())
2826    }
2827
2828    /// Resolves fragments for the given ids without scanning the manifest.
2829    ///
2830    /// The ids do not need to be sorted or deduplicated. Each id is resolved
2831    /// independently via the fragment bitmap.
2832    pub fn get_frags_from_ordered_ids(&self, ordered_ids: &[u32]) -> Vec<Option<FileFragment>> {
2833        let dataset = Arc::new(self.clone());
2834        ordered_ids
2835            .iter()
2836            .map(|id| {
2837                if !self.fragment_bitmap.contains(*id) {
2838                    return None;
2839                }
2840                let fragment_index = self.fragment_bitmap.rank(*id) as usize - 1;
2841                let fragment = self.manifest.fragments.get(fragment_index)?;
2842                debug_assert_eq!(
2843                    fragment.id, *id as u64,
2844                    "fragment_bitmap rank({id}) resolved to fragment {}, but fragment_bitmap and manifest.fragments are expected to stay in sync",
2845                    fragment.id
2846                );
2847                Some(FileFragment::new(dataset.clone(), fragment.clone()))
2848            })
2849            .collect()
2850    }
2851
2852    // This method filters deleted items from `addr_or_ids` using `addrs` as a reference
2853    async fn filter_addr_or_ids(&self, addr_or_ids: &[u64], addrs: &[u64]) -> Result<Vec<u64>> {
2854        // The final zip pairs these positionally; misalignment must fail
2855        // loud rather than truncate.
2856        if addr_or_ids.len() != addrs.len() {
2857            return Err(Error::internal(format!(
2858                "filter_addr_or_ids: addr_or_ids has {} entries but addrs has {}",
2859                addr_or_ids.len(),
2860                addrs.len()
2861            )));
2862        }
2863        if addrs.is_empty() {
2864            return Ok(Vec::new());
2865        }
2866
2867        let mut perm = permutation::sort(addrs);
2868        // First we sort the addrs, then we transform from Vec<u64> to Vec<Option<u64>> and then
2869        // we un-sort and use the None values to filter `addr_or_ids`
2870        let sorted_addrs = perm.apply_slice(addrs);
2871
2872        // Only collect deletion vectors for the fragments referenced by the given addrs
2873        let referenced_frag_ids = sorted_addrs
2874            .iter()
2875            .map(|addr| RowAddress::from(*addr).fragment_id())
2876            .dedup()
2877            .collect::<Vec<_>>();
2878        let frags = self.get_frags_from_ordered_ids(&referenced_frag_ids);
2879        let dv_futs = frags
2880            .iter()
2881            .map(|frag| {
2882                if let Some(frag) = frag {
2883                    frag.get_deletion_vector().boxed()
2884                } else {
2885                    std::future::ready(Ok(None)).boxed()
2886                }
2887            })
2888            .collect::<Vec<_>>();
2889        let dvs = stream::iter(dv_futs)
2890            .buffered(self.object_store.io_parallelism())
2891            .try_collect::<Vec<_>>()
2892            .await?;
2893
2894        // Iterate through the sorted addresses and sorted fragments (and sorted deletion vectors)
2895        // and filter out addresses that have been deleted
2896        let mut filtered_sorted_addrs = Vec::with_capacity(sorted_addrs.len());
2897        let mut sorted_addr_iter = sorted_addrs.into_iter().map(RowAddress::from);
2898        let mut next_addr = sorted_addr_iter.next().unwrap();
2899        let mut exhausted = false;
2900
2901        for frag_dv in frags.iter().zip(dvs).zip(referenced_frag_ids.iter()) {
2902            let ((frag, dv), frag_id) = frag_dv;
2903            if frag.is_some() {
2904                // Frag exists
2905                if let Some(dv) = dv.as_ref() {
2906                    // Deletion vector exists, scan DV
2907                    for deleted in dv.to_sorted_iter() {
2908                        while next_addr.fragment_id() == *frag_id
2909                            && next_addr.row_offset() < deleted
2910                        {
2911                            filtered_sorted_addrs.push(Some(u64::from(next_addr)));
2912                            if let Some(next) = sorted_addr_iter.next() {
2913                                next_addr = next;
2914                            } else {
2915                                exhausted = true;
2916                                break;
2917                            }
2918                        }
2919                        if exhausted {
2920                            break;
2921                        }
2922                        if next_addr.fragment_id() != *frag_id {
2923                            break;
2924                        }
2925                        if next_addr.row_offset() == deleted {
2926                            filtered_sorted_addrs.push(None);
2927                            if let Some(next) = sorted_addr_iter.next() {
2928                                next_addr = next;
2929                            } else {
2930                                exhausted = true;
2931                                break;
2932                            }
2933                        }
2934                    }
2935                }
2936                if exhausted {
2937                    break;
2938                }
2939                // Either no deletion vector, or we've exhausted it, keep everything else
2940                // in this frag
2941                while next_addr.fragment_id() == *frag_id {
2942                    filtered_sorted_addrs.push(Some(u64::from(next_addr)));
2943                    if let Some(next) = sorted_addr_iter.next() {
2944                        next_addr = next;
2945                    } else {
2946                        break;
2947                    }
2948                }
2949            } else {
2950                // Frag doesn't exist (possibly deleted), delete all items
2951                while next_addr.fragment_id() == *frag_id {
2952                    filtered_sorted_addrs.push(None);
2953                    if let Some(next) = sorted_addr_iter.next() {
2954                        next_addr = next;
2955                    } else {
2956                        break;
2957                    }
2958                }
2959            }
2960        }
2961
2962        // filtered_sorted_ids is now a Vec with the same size as sorted_addrs, but with None
2963        // values where the corresponding address was deleted.  We now need to un-sort it and
2964        // filter out the deleted addresses.
2965        perm.apply_inv_slice_in_place(&mut filtered_sorted_addrs);
2966        Ok(addr_or_ids
2967            .iter()
2968            .zip(filtered_sorted_addrs)
2969            .filter_map(|(addr_or_id, maybe_addr)| maybe_addr.map(|_| *addr_or_id))
2970            .collect())
2971    }
2972
2973    pub(crate) async fn filter_deleted_ids(&self, ids: &[u64]) -> Result<Vec<u64>> {
2974        let (ids, addresses) = if let Some(row_id_index) = get_row_id_index(self).await? {
2975            // Ids absent from the deletion-aware index are deleted; drop
2976            // them from both lists to keep the zip aligned. ids.len() is an
2977            // upper bound on the output size, so allocate once up front.
2978            let mut live_ids = Vec::with_capacity(ids.len());
2979            let mut addresses = Vec::with_capacity(ids.len());
2980            for id in ids {
2981                if let Some(address) = row_id_index.get(*id) {
2982                    live_ids.push(*id);
2983                    addresses.push(u64::from(address));
2984                }
2985            }
2986            (Cow::Owned(live_ids), Cow::Owned(addresses))
2987        } else {
2988            (Cow::Borrowed(ids), Cow::Borrowed(ids))
2989        };
2990
2991        self.filter_addr_or_ids(&ids, &addresses).await
2992    }
2993
2994    /// Gets the number of files that are so small they don't even have a full
2995    /// group. These are considered too small because reading many of them is
2996    /// much less efficient than reading a single file because the separate files
2997    /// split up what would otherwise be single IO requests into multiple.
2998    pub async fn num_small_files(&self, max_rows_per_group: usize) -> usize {
2999        futures::stream::iter(self.get_fragments())
3000            .map(|f| async move { f.physical_rows().await })
3001            .buffered(self.object_store.io_parallelism())
3002            .try_filter(|row_count| futures::future::ready(*row_count < max_rows_per_group))
3003            .count()
3004            .await
3005    }
3006
3007    pub async fn validate(&self) -> Result<()> {
3008        // All fragments have unique ids
3009        let id_counts =
3010            self.manifest
3011                .fragments
3012                .iter()
3013                .map(|f| f.id)
3014                .fold(HashMap::new(), |mut acc, id| {
3015                    *acc.entry(id).or_insert(0) += 1;
3016                    acc
3017                });
3018        for (id, count) in id_counts {
3019            if count > 1 {
3020                return Err(Error::corrupt_file(
3021                    self.base.clone(),
3022                    format!(
3023                        "Duplicate fragment id {} found in dataset {:?}",
3024                        id, self.base
3025                    ),
3026                ));
3027            }
3028        }
3029
3030        // Fragments are sorted in increasing fragment id order
3031        self.manifest
3032            .fragments
3033            .iter()
3034            .map(|f| f.id)
3035            .try_fold(0, |prev, id| {
3036                if id < prev {
3037                    Err(Error::corrupt_file(self.base.clone(), format!(
3038                        "Fragment ids are not sorted in increasing fragment-id order. Found {} after {} in dataset {:?}",
3039                        id, prev, self.base
3040                    )))
3041                } else {
3042                    Ok(id)
3043                }
3044            })?;
3045
3046        // All fragments have equal lengths
3047        futures::stream::iter(self.get_fragments())
3048            .map(|f| async move { f.validate().await })
3049            .buffer_unordered(self.object_store.io_parallelism())
3050            .try_collect::<Vec<()>>()
3051            .await?;
3052
3053        // Validate indices
3054        let indices = self.load_indices().await?;
3055        self.validate_indices(&indices)?;
3056
3057        Ok(())
3058    }
3059
3060    fn validate_indices(&self, indices: &[IndexMetadata]) -> Result<()> {
3061        // Make sure there are no duplicate ids
3062        let mut index_ids = HashSet::new();
3063        for index in indices.iter() {
3064            if !index_ids.insert(&index.uuid) {
3065                return Err(Error::corrupt_file(
3066                    self.manifest_location.path.clone(),
3067                    format!(
3068                        "Duplicate index id {} found in dataset {:?}",
3069                        index.uuid, self.base
3070                    ),
3071                ));
3072            }
3073        }
3074
3075        // For each index name, make sure there is no overlap in fragment bitmaps
3076        if let Err(err) = detect_overlapping_fragments(indices) {
3077            let mut message = "Overlapping fragments detected in dataset.".to_string();
3078            for (index_name, overlapping_frags) in err.bad_indices {
3079                message.push_str(&format!(
3080                    "\nIndex {:?} has overlapping fragments: {:?}",
3081                    index_name, overlapping_frags
3082                ));
3083            }
3084            return Err(Error::corrupt_file(
3085                self.manifest_location.path.clone(),
3086                message,
3087            ));
3088        };
3089
3090        Ok(())
3091    }
3092
3093    /// Migrate the dataset to use the new manifest path scheme.
3094    ///
3095    /// This function will rename all V1 manifests to [ManifestNamingScheme::V2].
3096    /// These paths provide more efficient opening of datasets with many versions
3097    /// on object stores.
3098    ///
3099    /// This function is idempotent, and can be run multiple times without
3100    /// changing the state of the object store.
3101    ///
3102    /// However, it should not be run while other concurrent operations are happening.
3103    /// And it should also run until completion before resuming other operations.
3104    ///
3105    /// ```rust
3106    /// # use lance::dataset::Dataset;
3107    /// # use lance_table::io::commit::ManifestNamingScheme;
3108    /// # use lance_datagen::{array, RowCount, BatchCount};
3109    /// # use arrow_array::types::Int32Type;
3110    /// # use lance::dataset::write::WriteParams;
3111    /// # let data = lance_datagen::gen_batch()
3112    /// #  .col("key", array::step::<Int32Type>())
3113    /// #  .into_reader_rows(RowCount::from(10), BatchCount::from(1));
3114    /// # let fut = async {
3115    /// # let params = WriteParams {
3116    /// #     enable_v2_manifest_paths: false,
3117    /// #     ..Default::default()
3118    /// # };
3119    /// let mut dataset = Dataset::write(data, "memory://test", Some(params)).await.unwrap();
3120    /// assert_eq!(dataset.manifest_location().naming_scheme, ManifestNamingScheme::V1);
3121    ///
3122    /// dataset.migrate_manifest_paths_v2().await.unwrap();
3123    /// assert_eq!(dataset.manifest_location().naming_scheme, ManifestNamingScheme::V2);
3124    /// # };
3125    /// # tokio::runtime::Runtime::new().unwrap().block_on(fut);
3126    /// ```
3127    pub async fn migrate_manifest_paths_v2(&mut self) -> Result<()> {
3128        migrate_scheme_to_v2(self.object_store.as_ref(), &self.base).await?;
3129        // We need to re-open.
3130        let latest_version = self.latest_version_id().await?;
3131        *self = self.checkout_version(latest_version).await?;
3132        Ok(())
3133    }
3134
3135    /// Shallow clone the target version into a new dataset at target_path.
3136    /// 'target_path': the uri string to clone the dataset into.
3137    /// 'version': the version cloned from, could be a version number or tag.
3138    /// 'store_params': the object store params to use for the new dataset.
3139    pub async fn shallow_clone(
3140        &mut self,
3141        target_path: &str,
3142        version: impl Into<refs::Ref>,
3143        store_params: Option<ObjectStoreParams>,
3144    ) -> Result<Self> {
3145        let (ref_name, version_number) = self.resolve_reference(version.into()).await?;
3146        let source_location = self.branch_location().find_branch(ref_name.as_deref())?;
3147        let clone_op = Operation::Clone {
3148            is_shallow: true,
3149            ref_name,
3150            ref_version: version_number,
3151            ref_path: source_location.uri,
3152            branch_name: None,
3153        };
3154        let transaction = Transaction::new(version_number, clone_op, None);
3155
3156        let builder = CommitBuilder::new(WriteDestination::Uri(target_path))
3157            .with_store_params(
3158                store_params.unwrap_or(self.store_params.as_deref().cloned().unwrap_or_default()),
3159            )
3160            .with_object_store(Arc::new(self.object_store.as_ref().clone()))
3161            .with_commit_handler(self.commit_handler.clone())
3162            .with_storage_format(self.manifest.data_storage_format.lance_file_version()?);
3163        builder.execute(transaction).await
3164    }
3165
3166    /// Deep clone the target version into a new dataset at target_path.
3167    /// This copies all relevant dataset files (data files, deletion files, and
3168    /// index files) into the target dataset without loading data into memory.
3169    ///
3170    /// The source files are read through this dataset's own object store while the
3171    /// copies are written through the target object store built from `store_params`.
3172    /// This makes the clone work across accounts/stores (e.g. between two abfss
3173    /// accounts): when the source and target stores are the same the copy stays
3174    /// server-side, otherwise the data is streamed through this process.
3175    ///
3176    /// Parameters:
3177    /// - `target_path`: the URI string to clone the dataset into.
3178    /// - `version`: the version cloned from, could be a version number, branch head, or tag.
3179    /// - `store_params`: the object store params for the target dataset (e.g. the
3180    ///   credentials of the target account).
3181    ///
3182    /// Note: external `base_paths` referenced by the source manifest are read through
3183    /// this dataset's object store; per-base distinct source credentials are not yet
3184    /// supported (see <https://github.com/lance-format/lance/issues/6093>).
3185    pub async fn deep_clone(
3186        &mut self,
3187        target_path: &str,
3188        version: impl Into<refs::Ref>,
3189        store_params: Option<ObjectStoreParams>,
3190    ) -> Result<Self> {
3191        use futures::StreamExt;
3192
3193        // Resolve source dataset and its manifest using checkout_version
3194        let src_ds = self.checkout_version(version).await?;
3195        let src_paths = src_ds.collect_paths().await?;
3196
3197        // Prepare target object store and base path
3198        let (target_store, target_base) = ObjectStore::from_uri_and_params(
3199            self.session.store_registry(),
3200            target_path,
3201            &store_params.clone().unwrap_or_default(),
3202        )
3203        .await?;
3204
3205        // Prevent cloning into an existing target dataset
3206        if self
3207            .commit_handler
3208            .resolve_latest_location(&target_base, &target_store)
3209            .await
3210            .is_ok()
3211        {
3212            return Err(Error::dataset_already_exists(target_path.to_string()));
3213        }
3214
3215        let build_absolute_path = |relative_path: &str, base: &Path| -> Path {
3216            let mut path = base.clone();
3217            for seg in relative_path.split('/') {
3218                if !seg.is_empty() {
3219                    path = path.clone().join(seg);
3220                }
3221            }
3222            path
3223        };
3224
3225        // When the source and target live in the same store we can keep the copy
3226        // server-side. Otherwise (e.g. cloning across accounts) we stream each file
3227        // from the source store to the target store.
3228        let same_store = src_ds.object_store.store_prefix == target_store.store_prefix;
3229
3230        // TODO: Leverage object store bulk copy for efficient same-store deep_clone.
3231        //
3232        // All cloud storage providers support batch copy APIs that would provide significant
3233        // performance improvements. We use single file copy before we have upstream support.
3234        //
3235        // Tracked by: https://github.com/lance-format/lance/issues/5435
3236        let io_parallelism = self.object_store.io_parallelism();
3237        let copy_futures = src_paths
3238            .iter()
3239            .map(|(relative_path, base)| {
3240                let source_store = Arc::clone(&src_ds.object_store);
3241                let target_store = Arc::clone(&target_store);
3242                let src_path = build_absolute_path(relative_path, base);
3243                let target_path = build_absolute_path(relative_path, &target_base);
3244                async move {
3245                    if same_store {
3246                        target_store.copy(&src_path, &target_path).await?;
3247                    } else {
3248                        let reader = source_store.open(&src_path).await?;
3249                        let mut writer = target_store.create(&target_path).await?;
3250                        writer.copy_from_reader(reader.as_ref()).await?;
3251                        writer.shutdown().await?;
3252                    }
3253                    Result::Ok(())
3254                }
3255            })
3256            .collect::<Vec<_>>();
3257
3258        futures::stream::iter(copy_futures)
3259            .buffer_unordered(io_parallelism)
3260            .collect::<Vec<_>>()
3261            .await
3262            .into_iter()
3263            .collect::<Result<Vec<_>>>()?;
3264
3265        // Record a Clone operation and commit via CommitBuilder
3266        let ref_name = src_ds.manifest.branch.clone();
3267        let ref_version = src_ds.manifest_location.version;
3268        let clone_op = Operation::Clone {
3269            is_shallow: false,
3270            ref_name,
3271            ref_version,
3272            ref_path: src_ds.uri().to_string(),
3273            branch_name: None,
3274        };
3275        let txn = Transaction::new(ref_version, clone_op, None);
3276        let builder = CommitBuilder::new(WriteDestination::Uri(target_path))
3277            .with_store_params(store_params.clone().unwrap_or_default())
3278            .with_object_store(target_store.clone())
3279            .with_source_store(src_ds.object_store.clone())
3280            .with_commit_handler(self.commit_handler.clone())
3281            .with_storage_format(self.manifest.data_storage_format.lance_file_version()?);
3282        let new_ds = builder.execute(txn).await?;
3283        Ok(new_ds)
3284    }
3285
3286    async fn resolve_reference(&self, reference: refs::Ref) -> Result<(Option<String>, u64)> {
3287        match reference {
3288            refs::Ref::Version(branch, version_number) => {
3289                if let Some(version_number) = version_number {
3290                    Ok((branch, version_number))
3291                } else {
3292                    let branch_location = self.branch_location().find_branch(branch.as_deref())?;
3293                    let version_number = self
3294                        .commit_handler
3295                        .resolve_latest_location(&branch_location.path, &self.object_store)
3296                        .await?
3297                        .version;
3298                    Ok((branch, version_number))
3299                }
3300            }
3301            refs::Ref::VersionNumber(version_number) => {
3302                Ok((self.manifest.branch.clone(), version_number))
3303            }
3304            refs::Ref::Tag(tag_name) => {
3305                let tag_contents = self.tags().get(tag_name.as_str()).await?;
3306                Ok((tag_contents.branch, tag_contents.version))
3307            }
3308        }
3309    }
3310
3311    /// Collect all (relative_path, path) of the dataset files.
3312    async fn collect_paths(&self) -> Result<Vec<(String, Path)>> {
3313        let mut file_paths: Vec<(String, Path)> = Vec::new();
3314        for fragment in self.manifest.fragments.iter() {
3315            if let Some(RowIdMeta::External(external_file)) = &fragment.row_id_meta {
3316                return Err(Error::internal(format!(
3317                    "External row_id_meta is not supported yet. external file path: {}",
3318                    external_file.path
3319                )));
3320            }
3321            for data_file in fragment.files.iter() {
3322                let base_root = if let Some(base_id) = data_file.base_id {
3323                    let base_path =
3324                        self.manifest.base_paths.get(&base_id).ok_or_else(|| {
3325                            Error::internal(format!("base_id {} not found", base_id))
3326                        })?;
3327                    Path::parse(base_path.path.as_str())?
3328                } else {
3329                    self.base.clone()
3330                };
3331                file_paths.push((
3332                    format!("{}/{}", DATA_DIR, data_file.path.clone()),
3333                    base_root,
3334                ));
3335            }
3336            if let Some(deletion_file) = &fragment.deletion_file {
3337                let base_root = if let Some(base_id) = deletion_file.base_id {
3338                    let base_path =
3339                        self.manifest.base_paths.get(&base_id).ok_or_else(|| {
3340                            Error::internal(format!("base_id {} not found", base_id))
3341                        })?;
3342                    Path::parse(base_path.path.as_str())?
3343                } else {
3344                    self.base.clone()
3345                };
3346                file_paths.push((
3347                    relative_deletion_file_path(fragment.id, deletion_file),
3348                    base_root,
3349                ));
3350            }
3351        }
3352
3353        let indices = read_manifest_indexes(
3354            self.object_store.as_ref(),
3355            &self.manifest_location,
3356            &self.manifest,
3357        )
3358        .await?;
3359
3360        for index in &indices {
3361            let base_root = if let Some(base_id) = index.base_id {
3362                let base_path = self
3363                    .manifest
3364                    .base_paths
3365                    .get(&base_id)
3366                    .ok_or_else(|| Error::internal(format!("base_id {} not found", base_id)))?;
3367                Path::parse(base_path.path.as_str())?
3368            } else {
3369                self.base.clone()
3370            };
3371            let index_root = base_root
3372                .clone()
3373                .join(INDICES_DIR)
3374                .join(index.uuid.to_string());
3375            let mut stream = self.object_store.read_dir_all(&index_root, None);
3376            while let Some(meta) = stream.next().await.transpose()? {
3377                if let Some(filename) = meta.location.filename() {
3378                    file_paths.push((
3379                        format!("{}/{}/{}", INDICES_DIR, index.uuid, filename),
3380                        base_root.clone(),
3381                    ));
3382                }
3383            }
3384        }
3385        Ok(file_paths)
3386    }
3387
3388    /// Run a SQL query against the dataset.
3389    /// The underlying SQL engine is DataFusion.
3390    /// Please refer to the DataFusion documentation for supported SQL syntax.
3391    pub fn sql(&self, sql: &str) -> SqlQueryBuilder {
3392        SqlQueryBuilder::new(self.clone(), sql)
3393    }
3394
3395    /// Returns true if Lance supports writing this datatype with nulls.
3396    pub(crate) fn lance_supports_nulls(&self, datatype: &DataType) -> bool {
3397        match self
3398            .manifest()
3399            .data_storage_format
3400            .lance_file_version()
3401            .unwrap_or(LanceFileVersion::Legacy)
3402            .resolve()
3403        {
3404            LanceFileVersion::Legacy => matches!(
3405                datatype,
3406                DataType::Utf8
3407                    | DataType::LargeUtf8
3408                    | DataType::Binary
3409                    | DataType::List(_)
3410                    | DataType::FixedSizeBinary(_)
3411                    | DataType::FixedSizeList(_, _)
3412            ),
3413            LanceFileVersion::V2_0 => !matches!(datatype, DataType::Struct(..)),
3414            _ => true,
3415        }
3416    }
3417}
3418
3419pub(crate) struct NewTransactionResult<'a> {
3420    pub dataset: BoxFuture<'a, Result<Dataset>>,
3421    pub new_transactions: BoxStream<'a, Result<(u64, Arc<Transaction>)>>,
3422}
3423
3424pub(crate) fn load_new_transactions(dataset: &Dataset) -> NewTransactionResult<'_> {
3425    // Resolve every manifest with version > our current version (the latest plus
3426    // the ones in between). On non-lexically-ordered stores this uses the version
3427    // hint to avoid an O(n) listing.
3428    let io_parallelism = dataset.object_store.as_ref().io_parallelism();
3429    let locations = dataset.commit_handler.list_manifest_locations_since(
3430        &dataset.base,
3431        dataset.object_store.as_ref(),
3432        dataset.manifest.version,
3433    );
3434
3435    // Will send the latest manifest via a channel.
3436    let (latest_tx, latest_rx) = tokio::sync::oneshot::channel();
3437    let mut latest_tx = Some(latest_tx);
3438
3439    let manifests = locations
3440        .map_ok(move |location| {
3441            let latest_tx = latest_tx.take();
3442            async move {
3443                let manifest = Dataset::get_manifest(
3444                    dataset.object_store.as_ref(),
3445                    &location,
3446                    &dataset.uri,
3447                    dataset.session.as_ref(),
3448                )
3449                .await?;
3450
3451                if let Some(latest_tx) = latest_tx {
3452                    // We ignore the error, since we don't care if the receiver is dropped.
3453                    let _ = latest_tx.send((manifest.clone(), location.clone()));
3454                }
3455
3456                Ok((manifest, location))
3457            }
3458        })
3459        .try_buffer_unordered(io_parallelism / 2);
3460    let transactions = manifests
3461        .map_ok(move |(manifest, location)| async move {
3462            let manifest_copy = manifest.clone();
3463            let tx_key = TransactionKey {
3464                version: manifest.version,
3465            };
3466            let transaction =
3467                if let Some(cached) = dataset.metadata_cache.get_with_key(&tx_key).await {
3468                    cached
3469                } else {
3470                    let dataset_version = Dataset::checkout_manifest(
3471                        dataset.object_store.clone(),
3472                        dataset.base.clone(),
3473                        dataset.uri.clone(),
3474                        manifest_copy.clone(),
3475                        location,
3476                        dataset.session(),
3477                        dataset.commit_handler.clone(),
3478                        dataset.file_reader_options.clone(),
3479                        dataset.store_params.as_deref().cloned(),
3480                        dataset.base_store_params.clone(),
3481                    )?;
3482                    let loaded =
3483                        Arc::new(dataset_version.read_transaction().await?.ok_or_else(|| {
3484                            Error::internal(format!(
3485                                "Dataset version {} does not have a transaction file",
3486                                manifest_copy.version
3487                            ))
3488                        })?);
3489                    dataset
3490                        .metadata_cache
3491                        .insert_with_key(&tx_key, loaded.clone())
3492                        .await;
3493                    loaded
3494                };
3495            Ok((manifest.version, transaction))
3496        })
3497        .try_buffer_unordered(io_parallelism / 2);
3498
3499    let dataset = async move {
3500        if let Ok((latest_manifest, location)) = latest_rx.await {
3501            // If we got the latest manifest, we can checkout the dataset.
3502            Dataset::checkout_manifest(
3503                dataset.object_store.clone(),
3504                dataset.base.clone(),
3505                dataset.uri.clone(),
3506                latest_manifest,
3507                location,
3508                dataset.session(),
3509                dataset.commit_handler.clone(),
3510                dataset.file_reader_options.clone(),
3511                dataset.store_params.as_deref().cloned(),
3512                dataset.base_store_params.clone(),
3513            )
3514        } else {
3515            // If we didn't get the latest manifest, we can still return the dataset
3516            // with the current manifest.
3517            Ok(dataset.clone())
3518        }
3519    }
3520    .boxed();
3521
3522    let new_transactions = transactions.boxed();
3523
3524    NewTransactionResult {
3525        dataset,
3526        new_transactions,
3527    }
3528}
3529
3530/// # Schema Evolution
3531///
3532/// Lance datasets support evolving the schema. Several operations are
3533/// supported that mirror common SQL operations:
3534///
3535/// - [Self::add_columns()]: Add new columns to the dataset, similar to `ALTER TABLE ADD COLUMN`.
3536/// - [Self::drop_columns()]: Drop columns from the dataset, similar to `ALTER TABLE DROP COLUMN`.
3537/// - [Self::alter_columns()]: Modify columns in the dataset, changing their name, type, or nullability.
3538///   Similar to `ALTER TABLE ALTER COLUMN`.
3539///
3540/// In addition, one operation is unique to Lance: [`merge`](Self::merge). This
3541/// operation allows inserting precomputed data into the dataset.
3542///
3543/// Because these operations change the schema of the dataset, they will conflict
3544/// with most other concurrent operations. Therefore, they should be performed
3545/// when no other write operations are being run.
3546impl Dataset {
3547    /// Append new columns to the dataset.
3548    pub async fn add_columns(
3549        &mut self,
3550        transforms: NewColumnTransform,
3551        read_columns: Option<Vec<String>>,
3552        batch_size: Option<u32>,
3553    ) -> Result<()> {
3554        schema_evolution::add_columns(self, transforms, read_columns, batch_size).await
3555    }
3556
3557    /// Modify columns in the dataset, changing their name, type, or nullability.
3558    ///
3559    /// If only changing the name or nullability of a column, this is a zero-copy
3560    /// operation and any indices will be preserved. If changing the type of a
3561    /// column, the data for that column will be rewritten and any indices will
3562    /// be dropped. The old column data will not be immediately deleted. To remove
3563    /// it, call [optimize::compact_files()] and then
3564    /// [cleanup::cleanup_old_versions()] on the dataset.
3565    pub async fn alter_columns(&mut self, alterations: &[ColumnAlteration]) -> Result<()> {
3566        schema_evolution::alter_columns(self, alterations).await
3567    }
3568
3569    /// Remove columns from the dataset.
3570    ///
3571    /// This is a metadata-only operation and does not remove the data from the
3572    /// underlying storage. In order to remove the data, you must subsequently
3573    /// call [optimize::compact_files()] to rewrite the data without the removed columns and
3574    /// then call [cleanup::cleanup_old_versions()] to remove the old files.
3575    pub async fn drop_columns(&mut self, columns: &[&str]) -> Result<()> {
3576        info!(target: TRACE_DATASET_EVENTS, event=DATASET_DROPPING_COLUMN_EVENT, uri = &self.uri, columns = columns.join(","));
3577        schema_evolution::drop_columns(self, columns).await
3578    }
3579
3580    /// Drop columns from the dataset and return updated dataset. Note that this
3581    /// is a zero-copy operation and column is not physically removed from the
3582    /// dataset.
3583    /// Parameters:
3584    /// - `columns`: the list of column names to drop.
3585    #[deprecated(since = "0.9.12", note = "Please use `drop_columns` instead.")]
3586    pub async fn drop(&mut self, columns: &[&str]) -> Result<()> {
3587        self.drop_columns(columns).await
3588    }
3589
3590    async fn merge_impl(
3591        &mut self,
3592        stream: Box<dyn RecordBatchReader + Send>,
3593        left_on: &str,
3594        right_on: &str,
3595    ) -> Result<()> {
3596        // Sanity check.
3597        if self.schema().field(left_on).is_none() && left_on != ROW_ID && left_on != ROW_ADDR {
3598            return Err(Error::invalid_input(format!(
3599                "Column {} does not exist in the left side dataset",
3600                left_on
3601            )));
3602        };
3603        let right_schema = stream.schema();
3604        if right_schema.field_with_name(right_on).is_err() {
3605            return Err(Error::invalid_input(format!(
3606                "Column {} does not exist in the right side dataset",
3607                right_on
3608            )));
3609        };
3610        for field in right_schema.fields() {
3611            if field.name() == right_on {
3612                // right_on is allowed to exist in the dataset, since it may be
3613                // the same as left_on.
3614                continue;
3615            }
3616            if self.schema().field(field.name()).is_some() {
3617                return Err(Error::invalid_input(format!(
3618                    "Column {} exists in both sides of the dataset",
3619                    field.name()
3620                )));
3621            }
3622        }
3623
3624        // Hash join
3625        let joiner = Arc::new(HashJoiner::try_new(stream, right_on).await?);
3626        // Final schema is union of current schema, plus the RHS schema without
3627        // the right_on key.
3628        let mut new_schema: Schema = self.schema().merge(joiner.out_schema().as_ref())?;
3629        new_schema.set_field_id(Some(self.manifest.max_field_id()));
3630
3631        // Write new data file to each fragment. Parallelism is done over columns,
3632        // so no parallelism done at this level.
3633        let updated_fragments: Vec<Fragment> = stream::iter(self.get_fragments())
3634            .then(|f| {
3635                let joiner = joiner.clone();
3636                async move { f.merge(left_on, &joiner).await.map(|f| f.metadata) }
3637            })
3638            .try_collect::<Vec<_>>()
3639            .await?;
3640
3641        let transaction = Transaction::new(
3642            self.manifest.version,
3643            Operation::Merge {
3644                fragments: updated_fragments,
3645                schema: new_schema,
3646            },
3647            None,
3648        );
3649
3650        self.apply_commit(transaction, &Default::default(), &Default::default())
3651            .await?;
3652
3653        Ok(())
3654    }
3655
3656    /// Merge this dataset with another arrow Table / Dataset, and returns a new version of dataset.
3657    ///
3658    /// Parameters:
3659    ///
3660    /// - `stream`: the stream of [`RecordBatch`] to merge.
3661    /// - `left_on`: the column name to join on the left side (self).
3662    /// - `right_on`: the column name to join on the right side (stream).
3663    ///
3664    /// Returns: a new version of dataset.
3665    ///
3666    /// It performs a left-join on the two datasets.
3667    pub async fn merge(
3668        &mut self,
3669        stream: impl RecordBatchReader + Send + 'static,
3670        left_on: &str,
3671        right_on: &str,
3672    ) -> Result<()> {
3673        let stream = Box::new(stream);
3674        self.merge_impl(stream, left_on, right_on).await
3675    }
3676
3677    /// Merge a distributed scalar index into a single root artifact and report
3678    /// progress via the supplied callback.
3679    pub async fn merge_index_metadata(
3680        &self,
3681        index_uuid: &Uuid,
3682        index_type: IndexType,
3683        _batch_readhead: Option<usize>,
3684        progress: Arc<dyn IndexBuildProgress>,
3685    ) -> Result<()> {
3686        let store = LanceIndexStore::from_dataset_for_new(self, index_uuid)?;
3687        let index_dir = self.indices_dir().join(index_uuid.to_string());
3688        match index_type {
3689            IndexType::Inverted => {
3690                // Call merge_index_files function for inverted index
3691                lance_index::scalar::inverted::builder::merge_index_files(
3692                    self.object_store.as_ref(),
3693                    &index_dir,
3694                    Arc::new(store),
3695                    progress,
3696                )
3697                .await
3698            }
3699            IndexType::BTree => {
3700                Err(Error::invalid_input(
3701                    "BTree distributed indexing no longer supports merge_index_metadata; \
3702                     build segments, optionally merge groups with merge_existing_index_segments(...), \
3703                     and commit with commit_existing_index_segments(...)"
3704                        .to_string(),
3705                ))
3706            }
3707            IndexType::Bitmap => {
3708                Err(Error::invalid_input(
3709                    "Bitmap distributed indexing no longer supports merge_index_metadata; \
3710                     build segments with create_index_uncommitted(...), merge them with \
3711                     merge_existing_index_segments(...), and commit with \
3712                     commit_existing_index_segments(...)"
3713                        .to_string(),
3714                ))
3715            }
3716            IndexType::IvfFlat | IndexType::IvfPq | IndexType::IvfSq | IndexType::Vector => {
3717                Err(Error::invalid_input(
3718                    "Vector distributed indexing no longer supports merge_index_metadata; \
3719                     build segments, optionally merge groups with merge_existing_index_segments(...), \
3720                     and commit with commit_existing_index_segments(...)"
3721                        .to_string(),
3722                ))
3723            }
3724            _ => Err(Error::invalid_input_source(Box::new(std::io::Error::new(
3725                std::io::ErrorKind::InvalidInput,
3726                format!("Unsupported index type (patched): {}", index_type),
3727            )))),
3728        }
3729    }
3730}
3731
3732/// # Dataset metadata APIs
3733///
3734/// There are four kinds of metadata on datasets:
3735///
3736///  - **Schema metadata**: metadata about the data itself.
3737///  - **Field metadata**: metadata about the dataset itself.
3738///  - **Dataset metadata**: metadata about the dataset. For example, this could
3739///    store a created_at date.
3740///  - **Dataset config**: configuration values controlling how engines should
3741///    manage the dataset. This configures things like auto-cleanup.
3742///
3743/// You can get
3744impl Dataset {
3745    /// Get dataset metadata.
3746    pub fn metadata(&self) -> &HashMap<String, String> {
3747        &self.manifest.table_metadata
3748    }
3749
3750    /// Get the dataset config from manifest
3751    pub fn config(&self) -> &HashMap<String, String> {
3752        &self.manifest.config
3753    }
3754
3755    /// Delete keys from the config.
3756    #[deprecated(
3757        note = "Use the new update_config(values, replace) method - pass None values to delete keys"
3758    )]
3759    pub async fn delete_config_keys(&mut self, delete_keys: &[&str]) -> Result<()> {
3760        let updates = delete_keys.iter().map(|key| (*key, None));
3761        self.update_config(updates).await?;
3762        Ok(())
3763    }
3764
3765    /// Update table metadata.
3766    ///
3767    /// Pass `None` for a value to remove that key.
3768    ///
3769    /// Use `.replace()` to replace the entire metadata map instead of merging.
3770    ///
3771    /// Returns the updated metadata map after the operation.
3772    ///
3773    /// ```
3774    /// # use lance::{Dataset, Result};
3775    /// # use lance::dataset::transaction::UpdateMapEntry;
3776    /// # async fn test_update_metadata(dataset: &mut Dataset) -> Result<()> {
3777    /// // Update single key
3778    /// dataset.update_metadata([("key", "value")]).await?;
3779    ///
3780    /// // Remove a key
3781    /// dataset.update_metadata([("to_delete", None)]).await?;
3782    ///
3783    /// // Clear all metadata
3784    /// dataset.update_metadata([] as [UpdateMapEntry; 0]).replace().await?;
3785    ///
3786    /// // Replace full metadata
3787    /// dataset.update_metadata([("k1", "v1"), ("k2", "v2")]).replace().await?;
3788    /// # Ok(())
3789    /// # }
3790    /// ```
3791    pub fn update_metadata(
3792        &mut self,
3793        values: impl IntoIterator<Item = impl Into<UpdateMapEntry>>,
3794    ) -> metadata::UpdateMetadataBuilder<'_> {
3795        metadata::UpdateMetadataBuilder::new(self, values, metadata::MetadataType::TableMetadata)
3796    }
3797
3798    /// Update config.
3799    ///
3800    /// Pass `None` for a value to remove that key.
3801    ///
3802    /// Use `.replace()` to replace the entire config map instead of merging.
3803    ///
3804    /// Returns the updated config map after the operation.
3805    ///
3806    /// ```
3807    /// # use lance::{Dataset, Result};
3808    /// # use lance::dataset::transaction::UpdateMapEntry;
3809    /// # async fn test_update_config(dataset: &mut Dataset) -> Result<()> {
3810    /// // Update single key
3811    /// dataset.update_config([("key", "value")]).await?;
3812    ///
3813    /// // Remove a key
3814    /// dataset.update_config([("to_delete", None)]).await?;
3815    ///
3816    /// // Clear all config
3817    /// dataset.update_config([] as [UpdateMapEntry; 0]).replace().await?;
3818    ///
3819    /// // Replace full config
3820    /// dataset.update_config([("k1", "v1"), ("k2", "v2")]).replace().await?;
3821    /// # Ok(())
3822    /// # }
3823    /// ```
3824    pub fn update_config(
3825        &mut self,
3826        values: impl IntoIterator<Item = impl Into<UpdateMapEntry>>,
3827    ) -> metadata::UpdateMetadataBuilder<'_> {
3828        metadata::UpdateMetadataBuilder::new(self, values, metadata::MetadataType::Config)
3829    }
3830
3831    /// Update schema metadata.
3832    ///
3833    /// Pass `None` for a value to remove that key.
3834    ///
3835    /// Use `.replace()` to replace the entire schema metadata map instead of merging.
3836    ///
3837    /// Returns the updated schema metadata map after the operation.
3838    ///
3839    /// ```
3840    /// # use lance::{Dataset, Result};
3841    /// # use lance::dataset::transaction::UpdateMapEntry;
3842    /// # async fn test_update_schema_metadata(dataset: &mut Dataset) -> Result<()> {
3843    /// // Update single key
3844    /// dataset.update_schema_metadata([("key", "value")]).await?;
3845    ///
3846    /// // Remove a key
3847    /// dataset.update_schema_metadata([("to_delete", None)]).await?;
3848    ///
3849    /// // Clear all schema metadata
3850    /// dataset.update_schema_metadata([] as [UpdateMapEntry; 0]).replace().await?;
3851    ///
3852    /// // Replace full schema metadata
3853    /// dataset.update_schema_metadata([("k1", "v1"), ("k2", "v2")]).replace().await?;
3854    /// # Ok(())
3855    /// # }
3856    /// ```
3857    pub fn update_schema_metadata(
3858        &mut self,
3859        values: impl IntoIterator<Item = impl Into<UpdateMapEntry>>,
3860    ) -> metadata::UpdateMetadataBuilder<'_> {
3861        metadata::UpdateMetadataBuilder::new(self, values, metadata::MetadataType::SchemaMetadata)
3862    }
3863
3864    /// Update schema metadata
3865    #[deprecated(note = "Use the new update_schema_metadata(values).replace() instead")]
3866    pub async fn replace_schema_metadata(
3867        &mut self,
3868        new_values: impl IntoIterator<Item = (String, String)>,
3869    ) -> Result<()> {
3870        let new_values = new_values
3871            .into_iter()
3872            .map(|(k, v)| (k, Some(v)))
3873            .collect::<HashMap<_, _>>();
3874        self.update_schema_metadata(new_values).replace().await?;
3875        Ok(())
3876    }
3877
3878    /// Update field metadata
3879    ///
3880    /// ```
3881    /// # use lance::{Dataset, Result};
3882    /// # use lance::dataset::transaction::UpdateMapEntry;
3883    /// # async fn test_update_field_metadata(dataset: &mut Dataset) -> Result<()> {
3884    /// // Update metadata by field path
3885    /// dataset.update_field_metadata()
3886    ///     .update("path.to_field", [("key", "value")])?
3887    ///     .await?;
3888    ///
3889    /// // Update metadata by field id
3890    /// dataset.update_field_metadata()
3891    ///     .update(12, [("key", "value")])?
3892    ///     .await?;
3893    ///
3894    /// // Clear field metadata
3895    /// dataset.update_field_metadata()
3896    ///     .replace("path.to_field", [] as [UpdateMapEntry; 0])?
3897    ///     .replace(12, [] as [UpdateMapEntry; 0])?
3898    ///     .await?;
3899    ///
3900    /// // Replace field metadata
3901    /// dataset.update_field_metadata()
3902    ///     .replace("field_name", [("k1", "v1"), ("k2", "v2")])?
3903    ///     .await?;
3904    /// # Ok(())
3905    /// # }
3906    /// ```
3907    pub fn update_field_metadata(&mut self) -> UpdateFieldMetadataBuilder<'_> {
3908        UpdateFieldMetadataBuilder::new(self)
3909    }
3910
3911    /// Update field metadata
3912    pub async fn replace_field_metadata(
3913        &mut self,
3914        new_values: impl IntoIterator<Item = (u32, HashMap<String, String>)>,
3915    ) -> Result<()> {
3916        let new_values = new_values.into_iter().collect::<HashMap<_, _>>();
3917        let field_metadata_updates = new_values
3918            .into_iter()
3919            .map(|(field_id, metadata)| {
3920                (
3921                    field_id as i32,
3922                    translate_schema_metadata_updates(&metadata),
3923                )
3924            })
3925            .collect();
3926        metadata::execute_metadata_update(
3927            self,
3928            Operation::UpdateConfig {
3929                config_updates: None,
3930                table_metadata_updates: None,
3931                schema_metadata_updates: None,
3932                field_metadata_updates,
3933            },
3934        )
3935        .await
3936    }
3937}
3938
3939#[async_trait::async_trait]
3940impl DatasetTakeRows for Dataset {
3941    fn schema(&self) -> &Schema {
3942        Self::schema(self)
3943    }
3944
3945    async fn take_rows(&self, row_ids: &[u64], projection: &Schema) -> Result<RecordBatch> {
3946        Self::take_rows(self, row_ids, projection.clone()).await
3947    }
3948}
3949
3950#[derive(Debug)]
3951pub(crate) struct ManifestWriteConfig {
3952    auto_set_feature_flags: bool,              // default true
3953    timestamp: Option<SystemTime>,             // default None
3954    use_stable_row_ids: bool,                  // default false
3955    use_legacy_format: Option<bool>,           // default None
3956    storage_format: Option<DataStorageFormat>, // default None
3957    disable_transaction_file: bool,            // default false
3958}
3959
3960impl Default for ManifestWriteConfig {
3961    fn default() -> Self {
3962        Self {
3963            auto_set_feature_flags: true,
3964            timestamp: None,
3965            use_stable_row_ids: false,
3966            disable_transaction_file: false,
3967            use_legacy_format: None,
3968            storage_format: None,
3969        }
3970    }
3971}
3972
3973impl ManifestWriteConfig {
3974    pub fn disable_transaction_file(&self) -> bool {
3975        self.disable_transaction_file
3976    }
3977}
3978
3979/// Commit a manifest file and create a copy at the latest manifest path.
3980#[allow(clippy::too_many_arguments)]
3981pub(crate) async fn write_manifest_file(
3982    object_store: &ObjectStore,
3983    commit_handler: &dyn CommitHandler,
3984    base_path: &Path,
3985    manifest: &mut Manifest,
3986    indices: Option<Vec<IndexMetadata>>,
3987    config: &ManifestWriteConfig,
3988    naming_scheme: ManifestNamingScheme,
3989    mut transaction: Option<&Transaction>,
3990) -> std::result::Result<ManifestLocation, CommitError> {
3991    if config.auto_set_feature_flags {
3992        // build_manifest may have already set FLAG_STABLE_ROW_IDS on the manifest.
3993        // Preserve it here so this second apply_feature_flags call does not clear it
3994        // when config.use_stable_row_ids is false (the ManifestWriteConfig default).
3995        let use_stable_row_ids = config.use_stable_row_ids || manifest.uses_stable_row_ids();
3996        apply_feature_flags(
3997            manifest,
3998            use_stable_row_ids,
3999            config.disable_transaction_file,
4000        )?;
4001    }
4002
4003    manifest.set_timestamp(timestamp_to_nanos(config.timestamp));
4004
4005    manifest.update_max_fragment_id();
4006
4007    commit_handler
4008        .commit(
4009            manifest,
4010            indices,
4011            base_path,
4012            object_store,
4013            write_manifest_file_to_path,
4014            naming_scheme,
4015            transaction.take().map(|tx| tx.into()),
4016        )
4017        .await
4018}
4019
4020impl Projectable for Dataset {
4021    fn schema(&self) -> &Schema {
4022        self.schema()
4023    }
4024}
4025
4026#[cfg(test)]
4027mod tests;