1use 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
155pub const DEFAULT_INDEX_CACHE_SIZE: usize = 6 * 1024 * 1024 * 1024;
158pub const DEFAULT_METADATA_CACHE_SIZE: usize = 1024 * 1024 * 1024;
162
163#[derive(Clone)]
165pub struct Dataset {
166 pub(crate) object_store: Arc<ObjectStore>,
169 pub(crate) commit_handler: Arc<dyn CommitHandler>,
170 uri: String,
175 pub(crate) base: Path,
176 pub manifest: Arc<Manifest>,
177 pub(crate) manifest_location: ManifestLocation,
180 pub(crate) session: Arc<Session>,
181 pub refs: Refs,
182
183 pub(crate) fragment_bitmap: Arc<RoaringBitmap>,
185
186 pub(crate) index_cache: Arc<DSIndexCache>,
188 pub(crate) metadata_cache: Arc<DSMetadataCache>,
189
190 pub(crate) file_reader_options: Option<FileReaderOptions>,
192
193 pub(crate) store_params: Option<Box<ObjectStoreParams>>,
196 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#[derive(Deserialize, Serialize, Debug)]
214pub struct Version {
215 pub version: u64,
217
218 pub timestamp: DateTime<Utc>,
220
221 pub metadata: BTreeMap<String, String>,
223}
224
225impl 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#[derive(Debug, Clone)]
242pub struct VersionTransaction {
243 pub version: u64,
245
246 pub timestamp: DateTime<Utc>,
248
249 pub transaction: Option<Transaction>,
251}
252
253#[derive(Clone, Debug)]
255pub struct ReadParams {
256 pub index_cache_size_bytes: usize,
259
260 pub metadata_cache_size_bytes: usize,
263
264 pub session: Option<Arc<Session>>,
268
269 pub store_options: Option<ObjectStoreParams>,
270
271 pub commit_handler: Option<Arc<dyn CommitHandler>>,
288
289 pub file_reader_options: Option<FileReaderOptions>,
296}
297
298impl ReadParams {
299 #[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; 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 #[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; self.metadata_cache_size_bytes = cache_size * assumed_entry_size;
323 self
324 }
325
326 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 pub fn session(&mut self, session: Arc<Session>) -> &mut Self {
334 self.session = Some(session);
335 self
336 }
337
338 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 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 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 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 ProjectionPlan::from_schema(dataset, schema.as_ref())
419 } else {
420 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 #[instrument]
451 pub async fn open(uri: &str) -> Result<Self> {
452 DatasetBuilder::from_uri(uri).load().await
453 }
454
455 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 pub fn statistics(&self) -> DatasetStatistics<'_> {
482 DatasetStatistics::new(self)
483 }
484
485 pub fn branches(&self) -> Branches<'_> {
486 self.refs.branches()
487 }
488
489 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 pub async fn checkout_branch(&self, branch: &str) -> Result<Self> {
506 self.checkout_by_ref(None, Some(branch)).await
507 }
508
509 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 .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 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 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 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 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 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 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 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 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 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 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 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) = ¶ms {
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 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 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 if response.managed_versioning == Some(true) {
953 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 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 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 let commit_handler: Option<Arc<dyn CommitHandler>> =
1016 if response.managed_versioning == Some(true) {
1017 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 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 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 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 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 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 pub fn manifest(&self) -> &Manifest {
1126 &self.manifest
1127 }
1128
1129 pub fn manifest_location(&self) -> &ManifestLocation {
1130 &self.manifest_location
1131 }
1132
1133 pub fn delta(&self) -> delta::DatasetDeltaBuilder {
1148 delta::DatasetDeltaBuilder::new(self.clone())
1149 }
1150
1151 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 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 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 async fn read_transaction_from_storage(
1224 &self,
1225 manifest: &Manifest,
1226 manifest_location: &ManifestLocation,
1227 ) -> Result<Option<Transaction>> {
1228 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 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 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 pub async fn read_version_transaction(&self, version: u64) -> Result<VersionTransaction> {
1286 let manifest_location = self
1288 .commit_handler
1289 .resolve_version_location(&self.base, version, &self.object_store.inner)
1290 .await?;
1291
1292 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 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 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 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 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 #[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 #[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 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 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 false,
1573 )
1574 .await
1575 }
1576
1577 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 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 pub fn scan(&self) -> Scanner {
1640 Scanner::new(Arc::new(self.clone()))
1641 }
1642
1643 #[instrument(skip_all)]
1647 pub async fn count_rows(&self, filter: Option<String>) -> Result<usize> {
1648 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() .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 #[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 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 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 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 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 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 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 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 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 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 pub async fn truncate_table(&mut self) -> Result<()> {
1963 self.delete("true").await.map(|_| ())
1964 }
1965
1966 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 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 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 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 #[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 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 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 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 pub async fn latest_storage_options(&self) -> Result<Option<StorageOptions>> {
2166 if let Some(accessor) = self.storage_options_accessor() {
2168 let options = accessor.get_storage_options().await?;
2169 return Ok(Some(options));
2170 }
2171
2172 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 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 let file_size = object_store.size(&filepath).await?;
2219
2220 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 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 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 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 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 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 pub fn version_id(&self) -> u64 {
2588 self.manifest.version
2589 }
2590
2591 pub fn version(&self) -> Version {
2596 Version::from(self.manifest.as_ref())
2597 }
2598
2599 pub async fn index_cache_entry_count(&self) -> usize {
2601 self.session.index_cache.size().await
2602 }
2603
2604 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 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 versions.sort_by_key(|v| v.version);
2630
2631 Ok(versions)
2632 }
2633
2634 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 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 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 #[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 pub fn schema(&self) -> &Schema {
2705 &self.manifest.schema
2706 }
2707
2708 pub fn empty_projection(self: &Arc<Self>) -> Projection {
2711 Projection::empty(self.clone())
2712 }
2713
2714 pub fn full_projection(self: &Arc<Self>) -> Projection {
2716 Projection::full(self.clone())
2717 }
2718
2719 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 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 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 async fn filter_addr_or_ids(&self, addr_or_ids: &[u64], addrs: &[u64]) -> Result<Vec<u64>> {
2854 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 let sorted_addrs = perm.apply_slice(addrs);
2871
2872 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 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 if let Some(dv) = dv.as_ref() {
2906 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 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 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 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 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 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 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 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 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 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 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 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 pub async fn migrate_manifest_paths_v2(&mut self) -> Result<()> {
3128 migrate_scheme_to_v2(self.object_store.as_ref(), &self.base).await?;
3129 let latest_version = self.latest_version_id().await?;
3131 *self = self.checkout_version(latest_version).await?;
3132 Ok(())
3133 }
3134
3135 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 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 let src_ds = self.checkout_version(version).await?;
3195 let src_paths = src_ds.collect_paths().await?;
3196
3197 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 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 let same_store = src_ds.object_store.store_prefix == target_store.store_prefix;
3229
3230 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 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 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 pub fn sql(&self, sql: &str) -> SqlQueryBuilder {
3392 SqlQueryBuilder::new(self.clone(), sql)
3393 }
3394
3395 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 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 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 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 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 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
3530impl Dataset {
3547 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 pub async fn alter_columns(&mut self, alterations: &[ColumnAlteration]) -> Result<()> {
3566 schema_evolution::alter_columns(self, alterations).await
3567 }
3568
3569 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 #[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 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 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 let joiner = Arc::new(HashJoiner::try_new(stream, right_on).await?);
3626 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 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 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 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 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
3732impl Dataset {
3745 pub fn metadata(&self) -> &HashMap<String, String> {
3747 &self.manifest.table_metadata
3748 }
3749
3750 pub fn config(&self) -> &HashMap<String, String> {
3752 &self.manifest.config
3753 }
3754
3755 #[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 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 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 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 #[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 pub fn update_field_metadata(&mut self) -> UpdateFieldMetadataBuilder<'_> {
3908 UpdateFieldMetadataBuilder::new(self)
3909 }
3910
3911 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, timestamp: Option<SystemTime>, use_stable_row_ids: bool, use_legacy_format: Option<bool>, storage_format: Option<DataStorageFormat>, disable_transaction_file: bool, }
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#[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 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;