1use crate::blob::BlobFile;
7use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
8use arrow_schema::{Schema, SchemaRef};
9use async_trait::async_trait;
10use datafusion_execution::TaskContext;
11use datafusion_expr::Expr;
12use datafusion_physical_plan::ExecutionPlan;
13use datafusion_physical_plan::display::DisplayableExecutionPlan;
14use futures::StreamExt;
15use futures::stream::FuturesUnordered;
16pub use lance::dataset::ColumnAlteration;
17pub use lance::dataset::NewColumnTransform;
18pub use lance::dataset::ReadParams;
19pub use lance::dataset::Version;
20use lance::dataset::WriteMode;
21use lance::dataset::builder::DatasetBuilder;
22use lance::dataset::{InsertBuilder, WriteParams};
23use lance::index::DatasetIndexExt;
24use lance::index::scalar::load_segment_params;
25use lance::io::{ObjectStoreParams, WrappingObjectStore};
26use lance_datafusion::utils::StreamingWriteSource;
27use lance_index::IndexCriteria;
28use lance_io::object_store::{LanceNamespaceStorageOptionsProvider, StorageOptionsAccessor};
29pub use query::AnyQuery;
30
31use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
32use lance_index::scalar::InvertedIndexParams;
33use lance_index::scalar::inverted::query::collect_query_tokens;
34use lance_namespace::LanceNamespace;
35use lance_namespace::error::NamespaceError;
36use lance_namespace::models::DescribeTableRequest;
37use lance_table::format::Manifest;
38use lance_table::io::commit::CommitHandler;
39use lance_table::io::commit::ManifestNamingScheme;
40use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
41use serde::{Deserialize, Serialize};
42use std::collections::{HashMap, HashSet};
43use std::format;
44use std::path::Path;
45use std::sync::Arc;
46
47use crate::connection::NamespaceClientPushdownOperation;
48
49use crate::DistanceType;
50use crate::blob::BlobRangeRequest;
51use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions};
52use crate::database::Database;
53use crate::database::listing::LANCE_FILE_EXTENSION;
54use crate::database::read_freshness::TableFreshness;
55use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry};
56use crate::error::{Error, Result};
57use crate::index::IndexStatistics;
58use crate::index::{Index, IndexBuilder};
59use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType};
60use crate::job::Job;
61use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery};
62use crate::table::datafusion::insert::InsertExec;
63use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
64
65use self::dataset::DatasetConsistencyWrapper;
66use self::merge::MergeInsertBuilder;
67
68pub mod add_columns;
69mod add_data;
70pub mod branch_merge;
71pub mod checkpoint;
72mod create_index;
73pub mod datafusion;
74pub(crate) mod dataset;
75pub mod delete;
76pub mod lsm_stats;
77pub mod merge;
78pub mod optimize;
79mod primary_key;
80pub mod query;
81pub mod schema_evolution;
82pub mod update;
83pub mod write_progress;
84use crate::index::waiter::wait_for_index;
85pub use add_columns::AddColumnsBuilder;
86#[cfg(feature = "remote")]
87pub(crate) use add_data::PreprocessingOutput;
88pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior};
89pub use branch_merge::{
90 BranchDiff, ColumnChange, ColumnSummary, IndexSummary, MergeBlocker, MergeBlockerCode,
91 MergeBranchResult, MergeBranchStatus, MergePreview, RowCountSummary,
92};
93pub use chrono::Duration;
94pub use delete::DeleteResult;
95use futures::future::join_all;
96pub use lance::dataset::refs::{BranchContents, Ref, TagContents, Tags as LanceTags};
97pub use lance::dataset::scanner::DatasetRecordBatchStream;
98pub use lance_index::optimize::OptimizeOptions;
99pub use lsm_stats::{BucketStats, GenerationStats, LsmStats, MemtableStats};
100pub use optimize::{CompactionOptions, OptimizeAction, OptimizeStats};
101pub use schema_evolution::{
102 AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate,
103 UpdateFieldMetadataResult,
104};
105use serde_with::skip_serializing_none;
106pub use update::{UpdateBuilder, UpdateResult};
107
108fn find_namespace_error<'a>(
114 err: &'a (dyn std::error::Error + 'static),
115) -> Option<&'a NamespaceError> {
116 let mut current: Option<&(dyn std::error::Error + 'static)> = Some(err);
117 while let Some(e) = current {
118 if let Some(ns_err) = e.downcast_ref::<NamespaceError>() {
119 return Some(ns_err);
120 }
121 current = e.source();
122 }
123 None
124}
125
126pub(crate) fn map_namespace_lance_error(err: lance::Error, table_name: &str) -> Error {
131 if let Some(code) = find_namespace_error(&err).map(NamespaceError::code) {
132 match code {
133 lance_namespace::error::ErrorCode::TableNotFound => {
134 return Error::TableNotFound {
135 name: table_name.to_string(),
136 source: Box::new(err),
137 };
138 }
139 lance_namespace::error::ErrorCode::TableAlreadyExists => {
140 return Error::TableAlreadyExists {
141 name: table_name.to_string(),
142 };
143 }
144 _ => {}
145 }
146 }
147 match err {
148 lance::Error::Namespace { source, .. } => Error::Runtime {
149 message: format!("Namespace error: {}", source),
150 },
151 other => other.into(),
152 }
153}
154
155async fn map_dataset_not_found(
165 uri: &str,
166 name: &str,
167 params: ReadParams,
168 err: lance::Error,
169) -> Error {
170 let name = name.to_string();
171 let source = Box::new(err);
172 if table_dir_exists(uri, params).await.unwrap_or(false) {
173 Error::TableCorrupted { name, source }
174 } else {
175 Error::TableNotFound { name, source }
176 }
177}
178
179async fn table_dir_exists(uri: &str, params: ReadParams) -> Result<bool> {
188 let (object_store, path, _) = DatasetBuilder::from_uri(uri)
189 .with_read_params(params)
190 .build_object_store()
191 .await?;
192 if path.extension() != Some(LANCE_FILE_EXTENSION) {
195 return Ok(false);
196 }
197 let (Some(parent), Some(dir_name)) = (path.parent(), path.filename()) else {
198 return Ok(false);
199 };
200 let entries = object_store.read_dir(parent).await?;
201 Ok(entries.iter().any(|entry| entry.as_str() == dir_name))
202}
203
204#[derive(Debug, Clone, Serialize, Deserialize)]
206pub enum ColumnKind {
207 Physical,
209 Embedding(EmbeddingDefinition),
211}
212
213#[derive(Debug, Clone, Serialize, Deserialize)]
215pub struct ColumnDefinition {
216 pub kind: ColumnKind,
218}
219
220#[derive(Debug, Clone)]
221pub struct TableDefinition {
222 pub column_definitions: Vec<ColumnDefinition>,
223 pub schema: SchemaRef,
224}
225
226impl TableDefinition {
227 pub fn new(schema: SchemaRef, column_definitions: Vec<ColumnDefinition>) -> Self {
228 Self {
229 column_definitions,
230 schema,
231 }
232 }
233
234 pub fn new_from_schema(schema: SchemaRef) -> Self {
235 let column_definitions = schema
236 .fields()
237 .iter()
238 .map(|_| ColumnDefinition {
239 kind: ColumnKind::Physical,
240 })
241 .collect();
242 Self::new(schema, column_definitions)
243 }
244
245 pub fn try_from_rich_schema(schema: SchemaRef) -> Result<Self> {
246 let column_definitions = schema.metadata.get("lancedb::column_definitions");
247 if let Some(column_definitions) = column_definitions {
248 let column_definitions: Vec<ColumnDefinition> =
249 serde_json::from_str(column_definitions).map_err(|e| Error::Runtime {
250 message: format!("Failed to deserialize column definitions: {}", e),
251 })?;
252 Ok(Self::new(schema, column_definitions))
253 } else {
254 let column_definitions = schema
255 .fields()
256 .iter()
257 .map(|_| ColumnDefinition {
258 kind: ColumnKind::Physical,
259 })
260 .collect();
261 Ok(Self::new(schema, column_definitions))
262 }
263 }
264
265 pub fn into_rich_schema(self) -> SchemaRef {
266 let lancedb_metadata = serde_json::to_string(&self.column_definitions).unwrap();
269 let mut schema_with_metadata = (*self.schema).clone();
270 schema_with_metadata
271 .metadata
272 .insert("lancedb::column_definitions".to_string(), lancedb_metadata);
273 Arc::new(schema_with_metadata)
274 }
275}
276
277#[derive(Clone, Debug, Default)]
280#[allow(dead_code)] enum BadVectorHandling {
282 #[default]
284 Error,
285 Drop,
287 Fill(f32),
289 None,
291}
292
293#[derive(Clone, Debug, Default)]
295pub struct WriteOptions {
296 pub lance_write_params: Option<WriteParams>,
304}
305
306pub enum Filter {
308 Sql(String),
310 Datafusion(Expr),
312}
313
314pub enum Predicate<'a> {
320 String(&'a str),
322 Expr(&'a Expr),
324}
325
326impl<'a> From<&'a str> for Predicate<'a> {
327 fn from(s: &'a str) -> Self {
328 Predicate::String(s)
329 }
330}
331
332impl<'a> From<&'a String> for Predicate<'a> {
333 fn from(s: &'a String) -> Self {
334 Predicate::String(s.as_str())
335 }
336}
337
338impl<'a> From<&'a Expr> for Predicate<'a> {
339 fn from(e: &'a Expr) -> Self {
340 Predicate::Expr(e)
341 }
342}
343
344#[async_trait]
345pub trait Tags: Send + Sync {
346 async fn list(&self) -> Result<HashMap<String, TagContents>>;
348
349 async fn get_version(&self, tag: &str) -> Result<u64>;
351
352 async fn create(&mut self, tag: &str, version: u64) -> Result<()>;
354
355 async fn delete(&mut self, tag: &str) -> Result<()>;
357
358 async fn update(&mut self, tag: &str, version: u64) -> Result<()>;
360}
361
362pub use self::merge::MergeResult;
363
364#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
377pub enum LsmWriteSpec {
378 Bucket {
385 column: String,
386 num_buckets: u32,
387 maintained_indexes: Vec<String>,
390 writer_config_defaults: HashMap<String, String>,
392 },
393 Identity {
398 column: String,
399 maintained_indexes: Vec<String>,
402 writer_config_defaults: HashMap<String, String>,
404 },
405 Unsharded {
407 maintained_indexes: Vec<String>,
410 writer_config_defaults: HashMap<String, String>,
412 },
413}
414
415impl LsmWriteSpec {
416 pub fn bucket(column: impl Into<String>, num_buckets: u32) -> Self {
418 Self::Bucket {
419 column: column.into(),
420 num_buckets,
421 maintained_indexes: Vec::new(),
422 writer_config_defaults: HashMap::new(),
423 }
424 }
425
426 pub fn identity(column: impl Into<String>) -> Self {
437 Self::Identity {
438 column: column.into(),
439 maintained_indexes: Vec::new(),
440 writer_config_defaults: HashMap::new(),
441 }
442 }
443
444 pub fn unsharded() -> Self {
446 Self::Unsharded {
447 maintained_indexes: Vec::new(),
448 writer_config_defaults: HashMap::new(),
449 }
450 }
451
452 pub fn with_maintained_indexes<I, S>(mut self, indexes: I) -> Self
456 where
457 I: IntoIterator<Item = S>,
458 S: Into<String>,
459 {
460 let v: Vec<String> = indexes.into_iter().map(Into::into).collect();
461 match &mut self {
462 Self::Bucket {
463 maintained_indexes, ..
464 }
465 | Self::Identity {
466 maintained_indexes, ..
467 }
468 | Self::Unsharded {
469 maintained_indexes, ..
470 } => *maintained_indexes = v,
471 }
472 self
473 }
474
475 pub fn with_writer_config_defaults<I, K, V>(mut self, defaults: I) -> Self
480 where
481 I: IntoIterator<Item = (K, V)>,
482 K: Into<String>,
483 V: Into<String>,
484 {
485 let m: HashMap<String, String> = defaults
486 .into_iter()
487 .map(|(k, v)| (k.into(), v.into()))
488 .collect();
489 match &mut self {
490 Self::Bucket {
491 writer_config_defaults,
492 ..
493 }
494 | Self::Identity {
495 writer_config_defaults,
496 ..
497 }
498 | Self::Unsharded {
499 writer_config_defaults,
500 ..
501 } => *writer_config_defaults = m,
502 }
503 self
504 }
505
506 pub fn maintained_indexes(&self) -> &[String] {
508 match self {
509 Self::Bucket {
510 maintained_indexes, ..
511 }
512 | Self::Identity {
513 maintained_indexes, ..
514 }
515 | Self::Unsharded {
516 maintained_indexes, ..
517 } => maintained_indexes,
518 }
519 }
520
521 pub fn writer_config_defaults(&self) -> &HashMap<String, String> {
523 match self {
524 Self::Bucket {
525 writer_config_defaults,
526 ..
527 }
528 | Self::Identity {
529 writer_config_defaults,
530 ..
531 }
532 | Self::Unsharded {
533 writer_config_defaults,
534 ..
535 } => writer_config_defaults,
536 }
537 }
538}
539
540#[derive(Debug, Clone, PartialEq, Eq)]
542pub struct FtsToken {
543 pub text: String,
545 pub position: u32,
547}
548
549pub fn tokenize(query: &str, params: &InvertedIndexParams) -> Result<Vec<FtsToken>> {
555 let mut tokenizer = params.build().map_err(|err| Error::InvalidInput {
556 message: format!("Failed to build tokenizer: {}", err),
557 })?;
558 let tokens = collect_query_tokens(query, &mut tokenizer);
559 Ok((0..tokens.len())
560 .map(|idx| FtsToken {
561 text: tokens.get_token(idx).to_string(),
562 position: tokens.position(idx),
563 })
564 .collect())
565}
566
567#[async_trait]
572pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
573 fn as_any(&self) -> &dyn std::any::Any;
575 fn name(&self) -> &str;
577 fn namespace(&self) -> &[String];
579 fn id(&self) -> &str;
584 async fn schema(&self) -> Result<SchemaRef>;
586 async fn count_rows(&self, filter: Option<Filter>) -> Result<usize>;
588 async fn create_plan(
590 &self,
591 query: &AnyQuery,
592 options: QueryExecutionOptions,
593 ) -> Result<Arc<dyn ExecutionPlan>>;
594 async fn query(
596 &self,
597 query: &AnyQuery,
598 options: QueryExecutionOptions,
599 ) -> Result<DatasetRecordBatchStream>;
600 async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result<String> {
602 let plan = self.create_plan(query, Default::default()).await?;
603 let display = DisplayableExecutionPlan::new(plan.as_ref());
604
605 Ok(format!("{}", display.indent(verbose)))
606 }
607 async fn analyze_plan(
608 &self,
609 query: &AnyQuery,
610 options: QueryExecutionOptions,
611 ) -> Result<String>;
612
613 async fn add(&self, add: AddDataBuilder) -> Result<AddResult>;
615 async fn delete(&self, predicate: Predicate<'_>) -> Result<DeleteResult>;
617 async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult>;
619 async fn create_index(&self, index: IndexBuilder) -> Result<()>;
621
622 async fn create_index_async(&self, index: IndexBuilder) -> Result<Job>;
624 async fn list_indices(&self) -> Result<Vec<IndexConfig>>;
626 async fn drop_index(&self, name: &str) -> Result<()>;
628 async fn prewarm_index(&self, name: &str) -> Result<()>;
630 async fn prewarm_data(&self, columns: Option<Vec<String>>) -> Result<()>;
635 async fn index_stats(&self, index_name: &str) -> Result<Option<IndexStatistics>>;
637 async fn merge_insert(
639 &self,
640 params: MergeInsertBuilder,
641 new_data: Box<dyn RecordBatchReader + Send>,
642 ) -> Result<MergeResult>;
643 async fn set_unenforced_primary_key(&self, _columns: &[&str]) -> Result<()> {
653 Err(Error::NotSupported {
654 message: "set_unenforced_primary_key is not supported on this table type".into(),
655 })
656 }
657 async fn set_lsm_write_spec(&self, _spec: LsmWriteSpec) -> Result<()> {
665 Err(Error::NotSupported {
666 message: "set_lsm_write_spec is not supported on this table type".into(),
667 })
668 }
669 async fn unset_lsm_write_spec(&self) -> Result<()> {
676 Err(Error::NotSupported {
677 message: "unset_lsm_write_spec is not supported on this table type".into(),
678 })
679 }
680 async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
686 Err(Error::NotSupported {
687 message: "get_lsm_write_spec is not supported on this table type".into(),
688 })
689 }
690 async fn flush_lsm(&self) -> Result<()> {
694 Err(Error::NotSupported {
695 message: "flush_lsm is not supported on this table type".into(),
696 })
697 }
698 async fn compact_lsm(&self) -> Result<()> {
702 Err(Error::NotSupported {
703 message: "compact_lsm is not supported on this table type".into(),
704 })
705 }
706 async fn get_lsm_stats(&self, _include_generation_rows: bool) -> Result<Option<LsmStats>> {
711 Err(Error::NotSupported {
712 message: "get_lsm_stats is not supported on this table type".into(),
713 })
714 }
715 async fn close_lsm_writers(&self) -> Result<()> {
720 Ok(())
721 }
722 async fn blob_columns(&self) -> Result<Vec<String>> {
724 Err(Error::NotSupported {
725 message: "blob_columns is not supported on this table type".into(),
726 })
727 }
728 async fn fetch_blobs(&self, _column: &str, _row_ids: &[u64]) -> Result<LargeBinaryArray> {
730 Err(Error::NotSupported {
731 message: "fetch_blobs is not supported on this table type".into(),
732 })
733 }
734 async fn fetch_blob_ranges(
736 &self,
737 _column: &str,
738 _requests: &[BlobRangeRequest],
739 ) -> Result<LargeBinaryArray> {
740 Err(Error::NotSupported {
741 message: "fetch_blob_ranges is not supported on this table type".into(),
742 })
743 }
744 async fn fetch_blob_files(
746 &self,
747 _column: &str,
748 _row_ids: &[u64],
749 ) -> Result<Vec<Option<BlobFile>>> {
750 Err(Error::NotSupported {
751 message: "fetch_blob_files is not supported on this table type".into(),
752 })
753 }
754 async fn tags(&self) -> Result<Box<dyn Tags + '_>>;
756 async fn optimize(&self, action: OptimizeAction) -> Result<OptimizeStats>;
758 async fn add_columns(
760 &self,
761 transforms: NewColumnTransform,
762 read_columns: Option<Vec<String>>,
763 ) -> Result<AddColumnsResult>;
764 async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult>;
766 async fn drop_columns(&self, columns: &[&str]) -> Result<DropColumnsResult>;
768 async fn version(&self) -> Result<u64>;
770 async fn checkout(&self, version: u64) -> Result<()>;
772 async fn checkout_tag(&self, tag: &str) -> Result<()>;
775 async fn checkout_latest(&self) -> Result<()>;
777 async fn restore(&self) -> Result<()>;
779 async fn list_versions(&self) -> Result<Vec<Version>>;
781 async fn create_branch(
783 &self,
784 name: &str,
785 from: lance::dataset::refs::Ref,
786 ) -> Result<Arc<dyn BaseTable>>;
787 async fn checkout_branch(&self, name: &str) -> Result<Arc<dyn BaseTable>>;
789 async fn checkout_branch_version(
796 &self,
797 name: &str,
798 version: Option<u64>,
799 ) -> Result<Arc<dyn BaseTable>> {
800 let branch = self.checkout_branch(name).await?;
801 if let Some(version) = version {
802 branch.checkout(version).await?;
803 }
804 Ok(branch)
805 }
806 async fn list_branches(&self) -> Result<HashMap<String, BranchContents>>;
808 async fn delete_branch(&self, name: &str) -> Result<()>;
810 async fn diff_branch(&self, _from_branch: &str) -> Result<BranchDiff> {
812 Err(Error::NotSupported {
813 message: "diff_branch is only supported on remote tables".into(),
814 })
815 }
816 async fn merge_branch(&self, _from_branch: &str, _dry_run: bool) -> Result<MergeBranchResult> {
819 Err(Error::NotSupported {
820 message: "merge_branch is only supported on remote tables".into(),
821 })
822 }
823 fn current_branch(&self) -> Option<String>;
825 async fn table_definition(&self) -> Result<TableDefinition>;
827 async fn uri(&self) -> Result<String>;
829 #[deprecated(since = "0.25.0", note = "Use initial_storage_options() instead")]
831 async fn storage_options(&self) -> Option<HashMap<String, String>>;
832 async fn initial_storage_options(&self) -> Option<HashMap<String, String>>;
836 async fn latest_storage_options(&self) -> Result<Option<HashMap<String, String>>>;
841 async fn wait_for_index(
844 &self,
845 index_names: &[&str],
846 timeout: std::time::Duration,
847 ) -> Result<()>;
848 async fn stats(&self) -> Result<TableStatistics>;
850 async fn create_insert_exec(
855 &self,
856 _input: Arc<dyn datafusion_physical_plan::ExecutionPlan>,
857 _write_params: WriteParams,
858 ) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
859 Err(Error::NotSupported {
860 message: "create_insert_exec not implemented".to_string(),
861 })
862 }
863 async fn update_field_metadata(
869 &self,
870 _updates: &[FieldMetadataUpdate],
871 ) -> Result<UpdateFieldMetadataResult> {
872 Err(Error::NotSupported {
873 message: "update_field_metadata is not supported on this table type".into(),
874 })
875 }
876}
877
878#[derive(Clone, Debug)]
882pub struct Table {
883 inner: Arc<dyn BaseTable>,
884 database: Option<Arc<dyn Database>>,
885 embedding_registry: Arc<dyn EmbeddingRegistry>,
886}
887
888#[cfg(all(test, feature = "remote"))]
889mod test_utils {
890 use super::*;
891
892 impl Table {
893 pub fn new_with_handler<T>(
894 name: impl Into<String>,
895 handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
896 ) -> Self
897 where
898 T: Into<reqwest::Body>,
899 {
900 let inner = Arc::new(crate::remote::table::RemoteTable::new_mock(
901 name.into(),
902 handler.clone(),
903 None,
904 ));
905 let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock(handler));
906 Self {
907 inner,
908 database: Some(database),
909 embedding_registry: Arc::new(MemoryRegistry::new()),
911 }
912 }
913
914 pub fn new_with_handler_and_interval<T>(
915 name: impl Into<String>,
916 handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
917 read_consistency_interval: Option<std::time::Duration>,
918 ) -> Self
919 where
920 T: Into<reqwest::Body>,
921 {
922 let inner = Arc::new(
923 crate::remote::table::RemoteTable::new_mock_with_consistency_interval(
924 name.into(),
925 handler.clone(),
926 read_consistency_interval,
927 ),
928 );
929 let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock(handler));
930 Self {
931 inner,
932 database: Some(database),
933 embedding_registry: Arc::new(MemoryRegistry::new()),
935 }
936 }
937
938 pub fn new_with_handler_version<T>(
939 name: impl Into<String>,
940 version: semver::Version,
941 handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
942 ) -> Self
943 where
944 T: Into<reqwest::Body>,
945 {
946 let inner = Arc::new(crate::remote::table::RemoteTable::new_mock(
947 name.into(),
948 handler.clone(),
949 Some(version),
950 ));
951 let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock(handler));
952 Self {
953 inner,
954 database: Some(database),
955 embedding_registry: Arc::new(MemoryRegistry::new()),
957 }
958 }
959
960 pub fn new_with_handler_and_config<T>(
961 name: impl Into<String>,
962 handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
963 config: crate::remote::ClientConfig,
964 ) -> Self
965 where
966 T: Into<reqwest::Body>,
967 {
968 let inner = Arc::new(crate::remote::table::RemoteTable::new_mock_with_config(
969 name.into(),
970 handler.clone(),
971 config.clone(),
972 ));
973 let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock_with_config(
974 handler, config,
975 ));
976 Self {
977 inner,
978 database: Some(database),
979 embedding_registry: Arc::new(MemoryRegistry::new()),
981 }
982 }
983
984 pub fn new_with_handler_version_and_config<T>(
985 name: impl Into<String>,
986 version: semver::Version,
987 handler: impl Fn(reqwest::Request) -> http::Response<T> + Clone + Send + Sync + 'static,
988 config: crate::remote::ClientConfig,
989 ) -> Self
990 where
991 T: Into<reqwest::Body>,
992 {
993 let inner = Arc::new(
994 crate::remote::table::RemoteTable::new_mock_with_version_and_config(
995 name.into(),
996 handler.clone(),
997 Some(version),
998 config.clone(),
999 ),
1000 );
1001 let database = Arc::new(crate::remote::db::RemoteDatabase::new_mock_with_config(
1002 handler, config,
1003 ));
1004 Self {
1005 inner,
1006 database: Some(database),
1007 embedding_registry: Arc::new(MemoryRegistry::new()),
1009 }
1010 }
1011 }
1012}
1013
1014impl std::fmt::Display for Table {
1015 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1016 write!(f, "{}", self.inner)
1017 }
1018}
1019
1020impl From<Arc<dyn BaseTable>> for Table {
1021 fn from(inner: Arc<dyn BaseTable>) -> Self {
1022 Self {
1023 inner,
1024 database: None,
1025 embedding_registry: Arc::new(MemoryRegistry::new()),
1026 }
1027 }
1028}
1029
1030impl Table {
1031 pub fn new(inner: Arc<dyn BaseTable>, database: Arc<dyn Database>) -> Self {
1032 Self {
1033 inner,
1034 database: Some(database),
1035 embedding_registry: Arc::new(MemoryRegistry::new()),
1036 }
1037 }
1038
1039 pub fn base_table(&self) -> &Arc<dyn BaseTable> {
1040 &self.inner
1041 }
1042
1043 pub fn database(&self) -> &Arc<dyn Database> {
1044 self.database.as_ref().unwrap()
1045 }
1046
1047 pub fn embedding_registry(&self) -> &Arc<dyn EmbeddingRegistry> {
1048 &self.embedding_registry
1049 }
1050
1051 pub(crate) fn new_with_embedding_registry(
1052 inner: Arc<dyn BaseTable>,
1053 database: Arc<dyn Database>,
1054 embedding_registry: Arc<dyn EmbeddingRegistry>,
1055 ) -> Self {
1056 Self {
1057 inner,
1058 database: Some(database),
1059 embedding_registry,
1060 }
1061 }
1062
1063 pub fn as_native(&self) -> Option<&NativeTable> {
1068 self.inner.as_native()
1069 }
1070
1071 pub fn name(&self) -> &str {
1073 self.inner.name()
1074 }
1075
1076 pub fn namespace(&self) -> &[String] {
1078 self.inner.namespace()
1079 }
1080
1081 pub fn id(&self) -> &str {
1083 self.inner.id()
1084 }
1085
1086 pub fn dataset(&self) -> Option<&dataset::DatasetConsistencyWrapper> {
1090 self.inner.as_native().map(|t| &t.dataset)
1091 }
1092
1093 pub async fn schema(&self) -> Result<SchemaRef> {
1095 self.inner.schema().await
1096 }
1097
1098 pub async fn count_rows(&self, filter: Option<String>) -> Result<usize> {
1104 self.inner.count_rows(filter.map(Filter::Sql)).await
1105 }
1106
1107 pub async fn blob_columns(&self) -> Result<Vec<String>> {
1112 self.inner.blob_columns().await
1113 }
1114
1115 pub async fn fetch_blobs(
1147 &self,
1148 column: impl AsRef<str>,
1149 row_ids: &[u64],
1150 ) -> Result<LargeBinaryArray> {
1151 self.inner.fetch_blobs(column.as_ref(), row_ids).await
1152 }
1153
1154 pub async fn fetch_blob_ranges(
1185 &self,
1186 column: impl AsRef<str>,
1187 requests: impl IntoIterator<Item = BlobRangeRequest>,
1188 ) -> Result<LargeBinaryArray> {
1189 let requests = requests.into_iter().collect::<Vec<_>>();
1190 self.inner
1191 .fetch_blob_ranges(column.as_ref(), &requests)
1192 .await
1193 }
1194
1195 pub async fn fetch_blob_files(
1212 &self,
1213 column: impl AsRef<str>,
1214 row_ids: &[u64],
1215 ) -> Result<Vec<Option<BlobFile>>> {
1216 self.inner.fetch_blob_files(column.as_ref(), row_ids).await
1217 }
1218
1219 pub fn add<T: Scannable + 'static>(&self, data: T) -> AddDataBuilder {
1226 AddDataBuilder::new(
1227 self.inner.clone(),
1228 Box::new(data),
1229 Some(self.embedding_registry.clone()),
1230 )
1231 }
1232
1233 pub fn update(&self) -> UpdateBuilder {
1248 UpdateBuilder::new(self.inner.clone())
1249 }
1250
1251 pub async fn delete(&self, predicate: impl Into<Predicate<'_>>) -> Result<DeleteResult> {
1304 self.inner.delete(predicate.into()).await
1305 }
1306
1307 pub fn create_index(&self, columns: &[impl AsRef<str>], index: Index) -> IndexBuilder {
1369 IndexBuilder::new(
1370 self.inner.clone(),
1371 columns
1372 .iter()
1373 .map(|val| val.as_ref().to_string())
1374 .collect::<Vec<_>>(),
1375 index,
1376 )
1377 }
1378
1379 pub fn create_index_with_timeout(
1382 &self,
1383 columns: &[impl AsRef<str>],
1384 index: Index,
1385 wait_timeout: Option<std::time::Duration>,
1386 ) -> IndexBuilder {
1387 let mut builder = IndexBuilder::new(
1388 self.inner.clone(),
1389 columns
1390 .iter()
1391 .map(|val| val.as_ref().to_string())
1392 .collect::<Vec<_>>(),
1393 index,
1394 );
1395 if let Some(timeout) = wait_timeout {
1396 builder = builder.wait_timeout(timeout);
1397 }
1398 builder
1399 }
1400
1401 pub fn merge_insert(&self, on: &[&str]) -> MergeInsertBuilder {
1477 MergeInsertBuilder::new(
1478 self.inner.clone(),
1479 on.iter().map(|s| s.to_string()).collect(),
1480 )
1481 }
1482
1483 pub fn query(&self) -> Query {
1571 Query::new(self.inner.clone())
1572 }
1573
1574 pub fn take_offsets(&self, offsets: Vec<u64>) -> TakeQuery {
1597 TakeQuery::from_offsets(self.inner.clone(), offsets)
1598 }
1599
1600 pub fn take_row_ids(&self, row_ids: Vec<u64>) -> TakeQuery {
1619 TakeQuery::from_row_ids(self.inner.clone(), row_ids)
1620 }
1621
1622 pub fn vector_search(&self, query: impl IntoQueryVector) -> Result<VectorQuery> {
1628 self.query().nearest_to(query)
1629 }
1630
1631 pub async fn optimize(&self, action: OptimizeAction) -> Result<OptimizeStats> {
1648 self.inner.optimize(action).await
1649 }
1650
1651 pub fn add_columns(&self) -> AddColumnsBuilder {
1653 AddColumnsBuilder::new(self.inner.clone())
1654 }
1655
1656 pub async fn alter_columns(
1658 &self,
1659 alterations: &[ColumnAlteration],
1660 ) -> Result<AlterColumnsResult> {
1661 self.inner.alter_columns(alterations).await
1662 }
1663
1664 pub async fn update_field_metadata(
1666 &self,
1667 updates: &[FieldMetadataUpdate],
1668 ) -> Result<UpdateFieldMetadataResult> {
1669 self.inner.update_field_metadata(updates).await
1670 }
1671
1672 pub async fn drop_columns(&self, columns: &[&str]) -> Result<DropColumnsResult> {
1674 self.inner.drop_columns(columns).await
1675 }
1676
1677 pub async fn set_unenforced_primary_key<I, S>(&self, columns: I) -> Result<()>
1690 where
1691 I: IntoIterator<Item = S>,
1692 S: Into<String>,
1693 {
1694 let owned: Vec<String> = columns.into_iter().map(Into::into).collect();
1695 let borrowed: Vec<&str> = owned.iter().map(String::as_str).collect();
1696 self.inner.set_unenforced_primary_key(&borrowed).await
1697 }
1698
1699 pub async fn set_lsm_write_spec(&self, spec: LsmWriteSpec) -> Result<()> {
1722 self.inner.set_lsm_write_spec(spec).await
1723 }
1724
1725 pub async fn unset_lsm_write_spec(&self) -> Result<()> {
1730 self.inner.unset_lsm_write_spec().await
1731 }
1732
1733 pub async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
1753 self.inner.get_lsm_write_spec().await
1754 }
1755
1756 pub async fn checkpoint_lsm(&self) -> Result<()> {
1790 checkpoint::checkpoint_lsm(self).await
1791 }
1792
1793 pub async fn flush_lsm(&self) -> Result<()> {
1801 self.inner.flush_lsm().await
1802 }
1803
1804 pub async fn compact_lsm(&self) -> Result<()> {
1810 self.inner.compact_lsm().await
1811 }
1812
1813 pub async fn get_lsm_stats(&self, include_generation_rows: bool) -> Result<Option<LsmStats>> {
1832 self.inner.get_lsm_stats(include_generation_rows).await
1833 }
1834
1835 pub async fn close_lsm_writers(&self) -> Result<()> {
1842 self.inner.close_lsm_writers().await
1843 }
1844
1845 pub async fn version(&self) -> Result<u64> {
1852 self.inner.version().await
1853 }
1854
1855 pub async fn checkout(&self, version: u64) -> Result<()> {
1870 self.inner.checkout(version).await
1871 }
1872
1873 pub async fn checkout_tag(&self, tag: &str) -> Result<()> {
1888 self.inner.checkout_tag(tag).await
1889 }
1890
1891 pub async fn checkout_latest(&self) -> Result<()> {
1896 self.inner.checkout_latest().await
1897 }
1898
1899 pub async fn restore(&self) -> Result<()> {
1910 self.inner.restore().await
1911 }
1912
1913 pub async fn list_versions(&self) -> Result<Vec<Version>> {
1915 self.inner.list_versions().await
1916 }
1917
1918 pub async fn list_indices(&self) -> Result<Vec<IndexConfig>> {
1920 self.inner.list_indices().await
1921 }
1922
1923 pub async fn tokenize(&self, query: &str, index_name: &str) -> Result<Vec<FtsToken>> {
1929 let indices = self.inner.list_indices().await?;
1930 let matches = indices
1931 .iter()
1932 .filter(|idx| idx.name == index_name)
1933 .collect::<Vec<_>>();
1934 let index = match matches.as_slice() {
1935 [index] => *index,
1936 [] => {
1937 return Err(Error::InvalidInput {
1938 message: format!("No index named '{}'", index_name),
1939 });
1940 }
1941 _ => {
1942 return Err(Error::InvalidInput {
1943 message: format!("Index name '{}' is ambiguous", index_name),
1944 });
1945 }
1946 };
1947 if index.index_type != IndexType::FTS {
1948 return Err(Error::InvalidInput {
1949 message: format!("Index '{}' is not a full text search index", index_name),
1950 });
1951 }
1952 self.tokenize_with_index(query, index, index_name)
1953 }
1954
1955 pub async fn tokenize_with_column(&self, query: &str, column: &str) -> Result<Vec<FtsToken>> {
1963 let schema = self.inner.schema().await?;
1964 let (column, _) = resolve_arrow_field_path(schema.as_ref(), column)?;
1965 let indices = self.inner.list_indices().await?;
1966 let matches = indices
1967 .iter()
1968 .filter(|idx| {
1969 idx.index_type == IndexType::FTS
1970 && idx.columns.len() == 1
1971 && idx.columns[0] == column
1972 })
1973 .collect::<Vec<_>>();
1974 let index = match matches.as_slice() {
1975 [index] => *index,
1976 [] => {
1977 return Err(Error::InvalidInput {
1978 message: format!("Column '{}' does not have a full text search index", column),
1979 });
1980 }
1981 _ => {
1982 return Err(Error::InvalidInput {
1983 message: format!(
1984 "Column '{}' has multiple full text search indexes; tokenization by column is ambiguous",
1985 column
1986 ),
1987 });
1988 }
1989 };
1990 self.tokenize(query, &index.name).await
1991 }
1992
1993 fn tokenize_with_index(
1994 &self,
1995 query: &str,
1996 index: &IndexConfig,
1997 index_name: &str,
1998 ) -> Result<Vec<FtsToken>> {
1999 let selector_description = format!("index name '{}'", index_name);
2000 let details = index
2001 .index_details
2002 .as_deref()
2003 .ok_or_else(|| Error::InvalidInput {
2004 message: format!(
2005 "Full text search index '{}' for {} does not include tokenizer details",
2006 index.name, selector_description
2007 ),
2008 })?;
2009 let params = serde_json::from_str::<InvertedIndexParams>(details).map_err(|err| {
2010 Error::InvalidInput {
2011 message: format!(
2012 "Failed to parse tokenizer details for full text search index '{}' for {}: {}",
2013 index.name, selector_description, err
2014 ),
2015 }
2016 })?;
2017 tokenize(query, ¶ms).map_err(|err| match err {
2018 Error::InvalidInput { message } => Error::InvalidInput {
2019 message: format!(
2020 "{} for full text search index '{}' for {}",
2021 message, index.name, selector_description
2022 ),
2023 },
2024 err => err,
2025 })
2026 }
2027
2028 pub async fn uri(&self) -> Result<String> {
2033 self.inner.uri().await
2034 }
2035
2036 #[deprecated(since = "0.25.0", note = "Use initial_storage_options() instead")]
2040 pub async fn storage_options(&self) -> Option<HashMap<String, String>> {
2041 #[allow(deprecated)]
2042 self.inner.storage_options().await
2043 }
2044
2045 pub async fn initial_storage_options(&self) -> Option<HashMap<String, String>> {
2051 self.inner.initial_storage_options().await
2052 }
2053
2054 pub async fn latest_storage_options(&self) -> Result<Option<HashMap<String, String>>> {
2062 self.inner.latest_storage_options().await
2063 }
2064
2065 pub async fn index_stats(
2068 &self,
2069 index_name: impl AsRef<str>,
2070 ) -> Result<Option<IndexStatistics>> {
2071 self.inner.index_stats(index_name.as_ref()).await
2072 }
2073
2074 pub async fn drop_index(&self, name: &str) -> Result<()> {
2083 self.inner.drop_index(name).await
2084 }
2085
2086 pub async fn prewarm_index(&self, name: &str) -> Result<()> {
2101 self.inner.prewarm_index(name).await
2102 }
2103
2104 pub async fn prewarm_data(&self, columns: Option<Vec<String>>) -> Result<()> {
2122 self.inner.prewarm_data(columns).await
2123 }
2124
2125 pub async fn wait_for_index(
2128 &self,
2129 index_names: &[&str],
2130 timeout: std::time::Duration,
2131 ) -> Result<()> {
2132 self.inner.wait_for_index(index_names, timeout).await
2133 }
2134
2135 pub async fn tags(&self) -> Result<Box<dyn Tags + '_>> {
2137 self.inner.tags().await
2138 }
2139
2140 pub async fn create_branch(
2142 &self,
2143 name: &str,
2144 from: impl Into<lance::dataset::refs::Ref>,
2145 ) -> Result<Self> {
2146 let inner = self.inner.create_branch(name, from.into()).await?;
2147 Ok(Self {
2148 inner,
2149 database: self.database.clone(),
2150 embedding_registry: self.embedding_registry.clone(),
2151 })
2152 }
2153
2154 pub async fn checkout_branch(&self, name: &str, version: Option<u64>) -> Result<Self> {
2168 let inner = self.inner.checkout_branch_version(name, version).await?;
2169 Ok(Self {
2170 inner,
2171 database: self.database.clone(),
2172 embedding_registry: self.embedding_registry.clone(),
2173 })
2174 }
2175
2176 pub async fn list_branches(&self) -> Result<HashMap<String, BranchContents>> {
2178 self.inner.list_branches().await
2179 }
2180
2181 pub async fn delete_branch(&self, name: &str) -> Result<()> {
2183 self.inner.delete_branch(name).await
2184 }
2185
2186 pub async fn diff_branch(&self, from_branch: &str) -> Result<BranchDiff> {
2188 self.inner.diff_branch(from_branch).await
2189 }
2190
2191 pub async fn merge_branch(
2194 &self,
2195 from_branch: &str,
2196 dry_run: bool,
2197 ) -> Result<MergeBranchResult> {
2198 self.inner.merge_branch(from_branch, dry_run).await
2199 }
2200
2201 pub fn current_branch(&self) -> Option<String> {
2203 self.inner.current_branch()
2204 }
2205
2206 pub async fn stats(&self) -> Result<TableStatistics> {
2208 self.inner.stats().await
2209 }
2210}
2211
2212pub struct NativeTags {
2213 dataset: dataset::DatasetConsistencyWrapper,
2214}
2215#[async_trait]
2216impl Tags for NativeTags {
2217 async fn list(&self) -> Result<HashMap<String, TagContents>> {
2218 let dataset = self.dataset.get().await?;
2219 Ok(dataset.tags().list().await?)
2220 }
2221
2222 async fn get_version(&self, tag: &str) -> Result<u64> {
2223 let dataset = self.dataset.get().await?;
2224 Ok(dataset.tags().get_version(tag).await?)
2225 }
2226
2227 async fn create(&mut self, tag: &str, version: u64) -> Result<()> {
2228 let dataset = self.dataset.get().await?;
2229 dataset.tags().create(tag, version).await?;
2230 Ok(())
2231 }
2232
2233 async fn delete(&mut self, tag: &str) -> Result<()> {
2234 let dataset = self.dataset.get().await?;
2235 dataset.tags().delete(tag).await?;
2236 Ok(())
2237 }
2238
2239 async fn update(&mut self, tag: &str, version: u64) -> Result<()> {
2240 let dataset = self.dataset.get().await?;
2241 dataset.tags().update(tag, version).await?;
2242 Ok(())
2243 }
2244}
2245
2246pub trait NativeTableExt {
2247 fn as_native(&self) -> Option<&NativeTable>;
2249}
2250
2251impl NativeTableExt for Arc<dyn BaseTable> {
2252 fn as_native(&self) -> Option<&NativeTable> {
2253 self.as_any().downcast_ref::<NativeTable>()
2254 }
2255}
2256
2257#[derive(Clone)]
2259pub struct NativeTable {
2260 name: String,
2261 namespace: Vec<String>,
2262 id: String,
2263 uri: String,
2264 pub(crate) dataset: dataset::DatasetConsistencyWrapper,
2265 read_consistency_interval: Option<std::time::Duration>,
2268 pub(crate) namespace_client: Option<Arc<dyn LanceNamespace>>,
2271 pub(crate) pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2274 freshness: Option<TableFreshness>,
2276}
2277
2278impl std::fmt::Debug for NativeTable {
2279 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2280 f.debug_struct("NativeTable")
2281 .field("name", &self.name)
2282 .field("namespace", &self.namespace)
2283 .field("id", &self.id)
2284 .field("uri", &self.uri)
2285 .field("read_consistency_interval", &self.read_consistency_interval)
2286 .field("namespace_client", &self.namespace_client)
2287 .field("pushdown_operations", &self.pushdown_operations)
2288 .finish()
2289 }
2290}
2291
2292impl std::fmt::Display for NativeTable {
2293 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2294 write!(
2295 f,
2296 "NativeTable({}, uri={}, read_consistency_interval={})",
2297 self.name,
2298 self.uri,
2299 match self.read_consistency_interval {
2300 None => {
2301 "None".to_string()
2302 }
2303 Some(duration) => {
2304 format!("{}s", duration.as_secs_f64())
2305 }
2306 }
2307 )
2308 }
2309}
2310
2311impl NativeTable {
2312 pub async fn open(uri: &str) -> Result<Self> {
2323 let name = Self::get_table_name(uri)?;
2324 Self::open_with_params(
2325 uri,
2326 &name,
2327 vec![],
2328 None,
2329 None,
2330 None,
2331 None,
2332 HashSet::new(),
2333 None,
2334 )
2335 .await
2336 }
2337
2338 #[allow(clippy::too_many_arguments)]
2354 pub async fn open_with_params(
2355 uri: &str,
2356 name: &str,
2357 namespace: Vec<String>,
2358 write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2359 params: Option<ReadParams>,
2360 read_consistency_interval: Option<std::time::Duration>,
2361 namespace_client: Option<Arc<dyn LanceNamespace>>,
2362 pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2363 managed_versioning: Option<bool>,
2364 ) -> Result<Self> {
2365 let params = params.unwrap_or_default();
2366 let params = match write_store_wrapper.clone() {
2368 Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
2369 None => params,
2370 };
2371
2372 let mut table_id = namespace.clone();
2374 table_id.push(name.to_string());
2375
2376 let managed_versioning = match managed_versioning {
2379 Some(value) => value,
2380 None if namespace_client.is_some() => {
2381 let ns_client = namespace_client.as_ref().unwrap();
2382 let describe_request = DescribeTableRequest {
2383 id: Some(table_id.clone()),
2384 ..Default::default()
2385 };
2386 let response = ns_client
2387 .describe_table(describe_request)
2388 .await
2389 .map_err(|e| Error::Runtime {
2390 message: format!(
2391 "Failed to describe table via namespace client: {}. \
2392 If you don't need managed versioning, don't pass namespace_client.",
2393 e
2394 ),
2395 })?;
2396 response.managed_versioning == Some(true)
2397 }
2398 None => false,
2399 };
2400
2401 let recovery_params = params.clone();
2403 let mut builder = DatasetBuilder::from_uri(uri).with_read_params(params);
2404
2405 if managed_versioning && let Some(ref ns_client) = namespace_client {
2407 let external_store = LanceNamespaceExternalManifestStore::for_table_uri(
2408 ns_client.clone(),
2409 table_id.clone(),
2410 uri,
2411 )?;
2412 let commit_handler: Arc<dyn CommitHandler> = Arc::new(ExternalManifestCommitHandler {
2413 external_manifest_store: Arc::new(external_store),
2414 });
2415 builder = builder.with_commit_handler(commit_handler);
2416 }
2417
2418 let dataset = match builder.load().await {
2419 Ok(dataset) => dataset,
2420 Err(e @ lance::Error::DatasetNotFound { .. }) => {
2421 return Err(map_dataset_not_found(uri, name, recovery_params, e).await);
2422 }
2423 Err(e) => return Err(e.into()),
2424 };
2425
2426 let dataset = DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval);
2427 let id = Self::build_id(&namespace, name);
2428
2429 Ok(Self {
2430 name: name.to_string(),
2431 namespace,
2432 id,
2433 uri: uri.to_string(),
2434 dataset,
2435 read_consistency_interval,
2436 namespace_client,
2437 pushdown_operations,
2438 freshness: None,
2439 })
2440 }
2441
2442 pub fn with_namespace_client(mut self, namespace_client: Arc<dyn LanceNamespace>) -> Self {
2446 self.namespace_client = Some(namespace_client);
2447 self
2448 }
2449
2450 pub(crate) fn with_freshness(mut self, freshness: TableFreshness) -> Self {
2452 self.freshness = Some(freshness);
2453 self
2454 }
2455
2456 fn with_dataset(&self, dataset: dataset::DatasetConsistencyWrapper) -> Self {
2459 Self {
2460 name: self.name.clone(),
2461 namespace: self.namespace.clone(),
2462 id: self.id.clone(),
2463 uri: self.uri.clone(),
2464 dataset,
2465 read_consistency_interval: self.read_consistency_interval,
2466 namespace_client: self.namespace_client.clone(),
2467 pushdown_operations: self.pushdown_operations.clone(),
2468 freshness: self.freshness.clone(),
2469 }
2470 }
2471
2472 fn bump_freshness(&self) {
2474 if let Some(freshness) = &self.freshness {
2475 freshness.bump();
2476 }
2477 }
2478
2479 fn validate_branch_name(name: &str, field: &str) -> Result<()> {
2480 if name.is_empty() {
2481 return Err(Error::InvalidInput {
2482 message: format!("{field} must be a non-empty string"),
2483 });
2484 }
2485 Ok(())
2486 }
2487
2488 #[allow(clippy::too_many_arguments)]
2510 pub async fn open_from_namespace(
2511 namespace_client: Arc<dyn LanceNamespace>,
2512 name: &str,
2513 namespace: Vec<String>,
2514 write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2515 params: Option<ReadParams>,
2516 read_consistency_interval: Option<std::time::Duration>,
2517 pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2518 session: Option<Arc<lance::session::Session>>,
2519 ) -> Result<Self> {
2520 let mut params = params.unwrap_or_default();
2521
2522 if let Some(sess) = session {
2524 params.session(sess);
2525 }
2526
2527 let params = match write_store_wrapper.clone() {
2529 Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
2530 None => params,
2531 };
2532
2533 let mut table_id = namespace.clone();
2535 table_id.push(name.to_string());
2536
2537 let builder = DatasetBuilder::from_namespace(namespace_client.clone(), table_id)
2540 .await
2541 .map_err(|e| map_namespace_lance_error(e, name))?;
2542
2543 let dataset = builder
2544 .with_read_params(params)
2545 .load()
2546 .await
2547 .map_err(|e| match e {
2548 lance::Error::DatasetNotFound { .. } => Error::TableNotFound {
2549 name: name.to_string(),
2550 source: Box::new(e),
2551 },
2552 e => e.into(),
2553 })?;
2554
2555 let uri = dataset.uri().to_string();
2556 let dataset = DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval);
2557 let id = Self::build_id(&namespace, name);
2558
2559 let stored_namespace_client =
2560 if pushdown_operations.contains(&NamespaceClientPushdownOperation::QueryTable) {
2561 Some(namespace_client)
2562 } else {
2563 None
2564 };
2565
2566 Ok(Self {
2567 name: name.to_string(),
2568 namespace,
2569 id,
2570 uri,
2571 dataset,
2572 read_consistency_interval,
2573 namespace_client: stored_namespace_client,
2574 pushdown_operations,
2575 freshness: None,
2576 })
2577 }
2578
2579 fn get_table_name(uri: &str) -> Result<String> {
2580 let path = Path::new(uri);
2581 let name = path
2582 .file_stem()
2583 .ok_or(Error::TableNotFound {
2584 name: uri.to_string(),
2585 source: format!("Could not extract table name from URI: '{}'", uri).into(),
2586 })?
2587 .to_str()
2588 .ok_or(Error::InvalidTableName {
2589 name: uri.to_string(),
2590 reason: "Table name is not valid URL".to_string(),
2591 })?;
2592 Ok(name.to_string())
2593 }
2594
2595 fn build_id(namespace: &[String], name: &str) -> String {
2596 if namespace.is_empty() {
2597 name.to_string()
2598 } else {
2599 let mut parts = namespace.to_vec();
2600 parts.push(name.to_string());
2601 parts.join("$")
2602 }
2603 }
2604
2605 #[allow(clippy::too_many_arguments)]
2622 pub async fn create(
2623 uri: &str,
2624 name: &str,
2625 namespace: Vec<String>,
2626 batches: impl StreamingWriteSource,
2627 write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2628 params: Option<WriteParams>,
2629 read_consistency_interval: Option<std::time::Duration>,
2630 namespace_client: Option<Arc<dyn LanceNamespace>>,
2631 pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2632 ) -> Result<Self> {
2633 let params = params.unwrap_or(WriteParams {
2635 ..Default::default()
2636 });
2637 let params = match write_store_wrapper.clone() {
2639 Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
2640 None => params,
2641 };
2642
2643 let insert_builder = InsertBuilder::new(uri).with_params(¶ms);
2644 let dataset = insert_builder
2645 .execute_stream(batches)
2646 .await
2647 .map_err(|e| match e {
2648 lance::Error::DatasetAlreadyExists { .. } => Error::TableAlreadyExists {
2649 name: name.to_string(),
2650 },
2651 e => e.into(),
2652 })?;
2653
2654 let id = Self::build_id(&namespace, name);
2655
2656 Ok(Self {
2657 name: name.to_string(),
2658 namespace,
2659 id,
2660 uri: uri.to_string(),
2661 dataset: DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval),
2662 read_consistency_interval,
2663 namespace_client,
2664 pushdown_operations,
2665 freshness: None,
2666 })
2667 }
2668
2669 #[allow(clippy::too_many_arguments)]
2670 pub async fn create_empty(
2671 uri: &str,
2672 name: &str,
2673 namespace: Vec<String>,
2674 schema: SchemaRef,
2675 write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2676 params: Option<WriteParams>,
2677 read_consistency_interval: Option<std::time::Duration>,
2678 namespace_client: Option<Arc<dyn LanceNamespace>>,
2679 pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2680 ) -> Result<Self> {
2681 let data: Box<dyn Scannable> = Box::new(RecordBatch::new_empty(schema));
2682 Self::create(
2683 uri,
2684 name,
2685 namespace,
2686 data,
2687 write_store_wrapper,
2688 params,
2689 read_consistency_interval,
2690 namespace_client,
2691 pushdown_operations,
2692 )
2693 .await
2694 }
2695
2696 #[allow(clippy::too_many_arguments)]
2720 pub async fn create_from_namespace(
2721 namespace_client: Arc<dyn LanceNamespace>,
2722 uri: &str,
2723 name: &str,
2724 namespace: Vec<String>,
2725 batches: impl StreamingWriteSource,
2726 write_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
2727 params: Option<WriteParams>,
2728 read_consistency_interval: Option<std::time::Duration>,
2729 pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
2730 session: Option<Arc<lance::session::Session>>,
2731 ) -> Result<Self> {
2732 let mut table_id = namespace.clone();
2734 table_id.push(name.to_string());
2735
2736 let storage_options_provider = Arc::new(LanceNamespaceStorageOptionsProvider::new(
2738 namespace_client.clone(),
2739 table_id,
2740 ));
2741
2742 let mut params = params.unwrap_or_default();
2744
2745 if let Some(sess) = session {
2747 params.session = Some(sess);
2748 }
2749
2750 let store_params = params
2752 .store_params
2753 .get_or_insert_with(ObjectStoreParams::default);
2754 let accessor = match store_params.storage_options().cloned() {
2755 Some(options) => {
2756 StorageOptionsAccessor::with_initial_and_provider(options, storage_options_provider)
2757 }
2758 None => StorageOptionsAccessor::with_provider(storage_options_provider),
2759 };
2760 store_params.storage_options_accessor = Some(Arc::new(accessor));
2761
2762 let params = match write_store_wrapper.clone() {
2764 Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
2765 None => params,
2766 };
2767
2768 let insert_builder = InsertBuilder::new(uri).with_params(¶ms);
2769 let dataset = insert_builder
2770 .execute_stream(batches)
2771 .await
2772 .map_err(|e| match e {
2773 lance::Error::DatasetAlreadyExists { .. } => Error::TableAlreadyExists {
2774 name: name.to_string(),
2775 },
2776 e => e.into(),
2777 })?;
2778
2779 let id = Self::build_id(&namespace, name);
2780
2781 let stored_namespace_client =
2782 if pushdown_operations.contains(&NamespaceClientPushdownOperation::QueryTable) {
2783 Some(namespace_client)
2784 } else {
2785 None
2786 };
2787
2788 Ok(Self {
2789 name: name.to_string(),
2790 namespace,
2791 id,
2792 uri: uri.to_string(),
2793 dataset: DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval),
2794 read_consistency_interval,
2795 namespace_client: stored_namespace_client,
2796 pushdown_operations,
2797 freshness: None,
2798 })
2799 }
2800
2801 pub async fn merge(
2803 &mut self,
2804 batches: impl RecordBatchReader + Send + 'static,
2805 left_on: &str,
2806 right_on: &str,
2807 ) -> Result<()> {
2808 self.dataset.ensure_mutable()?;
2809 let mut dataset = (*self.dataset.get().await?).clone();
2810 dataset.merge(batches, left_on, right_on).await?;
2811 self.dataset.update(dataset);
2812 Ok(())
2813 }
2814
2815 pub async fn count_fragments(&self) -> Result<usize> {
2817 Ok(self.dataset.get().await?.count_fragments())
2818 }
2819
2820 pub async fn count_deleted_rows(&self) -> Result<usize> {
2821 Ok(self.dataset.get().await?.count_deleted_rows().await?)
2822 }
2823
2824 pub async fn num_small_files(&self, max_rows_per_group: usize) -> Result<usize> {
2825 Ok(self
2826 .dataset
2827 .get()
2828 .await?
2829 .num_small_files(max_rows_per_group)
2830 .await)
2831 }
2832 pub async fn uses_v2_manifest_paths(&self) -> Result<bool> {
2837 let dataset = self.dataset.get().await?;
2838 Ok(dataset.manifest_location().naming_scheme == ManifestNamingScheme::V2)
2839 }
2840
2841 pub async fn migrate_manifest_paths_v2(&self) -> Result<()> {
2856 self.dataset.ensure_mutable()?;
2857 let mut dataset = (*self.dataset.get().await?).clone();
2858 dataset.migrate_manifest_paths_v2().await?;
2859 self.dataset.update(dataset);
2860 Ok(())
2861 }
2862
2863 pub async fn manifest(&self) -> Result<Manifest> {
2865 let dataset = self.dataset.get().await?;
2866 Ok(dataset.manifest().clone())
2867 }
2868
2869 pub async fn update_config(
2871 &self,
2872 upsert_values: impl IntoIterator<Item = (String, String)>,
2873 ) -> Result<()> {
2874 self.dataset.ensure_mutable()?;
2875 let mut dataset = (*self.dataset.get().await?).clone();
2876 dataset.update_config(upsert_values).await?;
2877 self.dataset.update(dataset);
2878 Ok(())
2879 }
2880
2881 pub async fn delete_config_keys(&self, delete_keys: &[&str]) -> Result<()> {
2883 self.dataset.ensure_mutable()?;
2884 let mut dataset = (*self.dataset.get().await?).clone();
2885 #[allow(deprecated)]
2887 dataset.delete_config_keys(delete_keys).await?;
2888 self.dataset.update(dataset);
2889 Ok(())
2890 }
2891
2892 pub async fn replace_schema_metadata(
2894 &self,
2895 upsert_values: impl IntoIterator<Item = (String, String)>,
2896 ) -> Result<()> {
2897 self.dataset.ensure_mutable()?;
2898 let mut dataset = (*self.dataset.get().await?).clone();
2899 #[allow(deprecated)]
2901 dataset.replace_schema_metadata(upsert_values).await?;
2902 self.dataset.update(dataset);
2903 Ok(())
2904 }
2905
2906 #[deprecated(since = "0.33.1", note = "Use `update_field_metadata` instead")]
2914 pub async fn replace_field_metadata(
2915 &self,
2916 new_values: impl IntoIterator<Item = (u32, HashMap<String, String>)>,
2917 ) -> Result<()> {
2918 self.dataset.ensure_mutable()?;
2919 let mut dataset = (*self.dataset.get().await?).clone();
2920 dataset.replace_field_metadata(new_values).await?;
2921 self.dataset.update(dataset);
2922 Ok(())
2923 }
2924}
2925
2926#[async_trait::async_trait]
2927impl BaseTable for NativeTable {
2928 fn as_any(&self) -> &dyn std::any::Any {
2929 self
2930 }
2931
2932 fn name(&self) -> &str {
2933 self.name.as_str()
2934 }
2935
2936 fn namespace(&self) -> &[String] {
2937 &self.namespace
2938 }
2939
2940 fn id(&self) -> &str {
2941 &self.id
2942 }
2943
2944 async fn version(&self) -> Result<u64> {
2945 Ok(self.dataset.get().await?.version().version)
2946 }
2947
2948 async fn checkout(&self, version: u64) -> Result<()> {
2949 self.dataset.as_time_travel(version).await
2950 }
2951
2952 async fn checkout_tag(&self, tag: &str) -> Result<()> {
2953 self.dataset.as_time_travel(tag).await
2954 }
2955
2956 async fn checkout_latest(&self) -> Result<()> {
2957 self.bump_freshness();
2959 self.dataset.as_latest().await?;
2960 self.dataset.reload().await
2961 }
2962
2963 async fn create_branch(
2964 &self,
2965 name: &str,
2966 from: lance::dataset::refs::Ref,
2967 ) -> Result<Arc<dyn BaseTable>> {
2968 Self::validate_branch_name(name, "branch name")?;
2969 if let lance::dataset::refs::Ref::Version(Some(from_branch), _) = &from {
2970 Self::validate_branch_name(from_branch, "from_ref")?;
2971 }
2972 let mut ds = (*self.dataset.get().await?).clone();
2973 let branch_ds = ds.create_branch(name, from, None).await?;
2974 let dataset = dataset::DatasetConsistencyWrapper::new_latest(
2975 branch_ds,
2976 self.read_consistency_interval,
2977 );
2978 Ok(Arc::new(self.with_dataset(dataset)))
2979 }
2980
2981 async fn checkout_branch(&self, name: &str) -> Result<Arc<dyn BaseTable>> {
2982 Self::validate_branch_name(name, "branch name")?;
2983 let branch_ds = self.dataset.get().await?.checkout_branch(name).await?;
2984 let dataset = dataset::DatasetConsistencyWrapper::new_latest(
2985 branch_ds,
2986 self.read_consistency_interval,
2987 );
2988 Ok(Arc::new(self.with_dataset(dataset)))
2989 }
2990
2991 async fn checkout_branch_version(
2992 &self,
2993 name: &str,
2994 version: Option<u64>,
2995 ) -> Result<Arc<dyn BaseTable>> {
2996 let Some(version) = version else {
2997 return self.checkout_branch(name).await;
2998 };
2999 Self::validate_branch_name(name, "branch name")?;
3000 let branch_ds = self
3002 .dataset
3003 .get()
3004 .await?
3005 .checkout_version((name, version))
3006 .await?;
3007 let dataset = dataset::DatasetConsistencyWrapper::new_time_travel(
3008 branch_ds,
3009 self.read_consistency_interval,
3010 );
3011 Ok(Arc::new(self.with_dataset(dataset)))
3012 }
3013
3014 async fn list_branches(&self) -> Result<HashMap<String, BranchContents>> {
3015 Ok(self.dataset.get().await?.list_branches().await?)
3016 }
3017
3018 async fn delete_branch(&self, name: &str) -> Result<()> {
3019 Self::validate_branch_name(name, "branch name")?;
3020 let mut ds = (*self.dataset.get().await?).clone();
3021 ds.delete_branch(name).await?;
3022 Ok(())
3023 }
3024
3025 fn current_branch(&self) -> Option<String> {
3026 self.dataset.current_branch()
3027 }
3028
3029 async fn list_versions(&self) -> Result<Vec<Version>> {
3030 Ok(self.dataset.get().await?.versions().await?)
3031 }
3032
3033 async fn restore(&self) -> Result<()> {
3034 let version = self
3035 .dataset
3036 .time_travel_version()
3037 .ok_or_else(|| Error::InvalidInput {
3038 message: "you must run checkout before running restore".to_string(),
3039 })?;
3040 {
3041 let mut dataset = (*self.dataset.get().await?).clone();
3043 debug_assert_eq!(dataset.version().version, version);
3044 dataset.restore().await?;
3045 }
3046 self.bump_freshness();
3048 self.dataset.as_latest().await?;
3049 Ok(())
3050 }
3051
3052 async fn schema(&self) -> Result<SchemaRef> {
3053 let lance_schema = self.dataset.get().await?.schema().clone();
3054 Ok(Arc::new(Schema::from(&lance_schema)))
3055 }
3056
3057 async fn table_definition(&self) -> Result<TableDefinition> {
3058 let schema = self.schema().await?;
3059 TableDefinition::try_from_rich_schema(schema)
3060 }
3061
3062 async fn count_rows(&self, filter: Option<Filter>) -> Result<usize> {
3063 let dataset = self.dataset.get().await?;
3064 match filter {
3065 None => Ok(dataset.count_rows(None).await?),
3066 Some(Filter::Sql(sql)) => Ok(dataset.count_rows(Some(sql)).await?),
3067 Some(Filter::Datafusion(_)) => Err(Error::NotSupported {
3068 message: "Datafusion filters are not yet supported".to_string(),
3069 }),
3070 }
3071 }
3072
3073 async fn add(&self, mut add: AddDataBuilder) -> Result<AddResult> {
3074 let table_def = self.table_definition().await?;
3075
3076 self.dataset.ensure_mutable()?;
3077 let ds_wrapper = self.dataset.clone();
3078 let ds = self.dataset.get().await?;
3079
3080 let table_schema = Schema::from(&ds.schema().clone());
3081
3082 let num_partitions = if let Some(parallelism) = add.write_parallelism {
3083 parallelism
3084 } else {
3085 let mut peeked = PeekedScannable::new(add.data);
3088 let n = if let Some(first_batch) = peeked.peek().await {
3089 let max_partitions = lance_core::utils::tokio::get_num_compute_intensive_cpus();
3090 estimate_write_partitions(
3091 first_batch.get_array_memory_size(),
3092 first_batch.num_rows(),
3093 peeked.num_rows(),
3094 max_partitions,
3095 )
3096 } else {
3097 1
3098 };
3099 add.data = Box::new(peeked);
3100 n
3101 };
3102
3103 let output = add.into_plan(&table_schema, &table_def)?;
3104
3105 let lance_params = output
3106 .write_options
3107 .lance_write_params
3108 .unwrap_or(WriteParams {
3109 mode: match output.mode {
3110 AddDataMode::Append => WriteMode::Append,
3111 AddDataMode::Overwrite => WriteMode::Overwrite,
3112 },
3113 ..Default::default()
3114 });
3115
3116 let plan = if num_partitions > 1 {
3118 Arc::new(
3119 datafusion_physical_plan::repartition::RepartitionExec::try_new(
3120 output.plan,
3121 datafusion_physical_plan::Partitioning::RoundRobinBatch(num_partitions),
3122 )?,
3123 ) as Arc<dyn ExecutionPlan>
3124 } else {
3125 output.plan
3126 };
3127
3128 let insert_exec = Arc::new(InsertExec::new_with_tracker(
3129 ds_wrapper.clone(),
3130 ds,
3131 plan,
3132 lance_params,
3133 output.tracker.clone(),
3134 ));
3135
3136 let tracker_for_tasks = output.tracker.clone();
3137 if let Some(ref t) = tracker_for_tasks {
3138 t.set_total_tasks(num_partitions);
3139 }
3140 let _finish = write_progress::FinishOnDrop(output.tracker);
3141
3142 let task_ctx = Arc::new(TaskContext::default());
3144 let handles = FuturesUnordered::new();
3145 for partition in 0..num_partitions {
3146 let exec = insert_exec.clone();
3147 let ctx = task_ctx.clone();
3148 let tracker = tracker_for_tasks.clone();
3149 handles.push(tokio::spawn(async move {
3150 let _guard = tracker.as_ref().map(|t| t.track_task());
3151 let mut stream = exec
3152 .execute(partition, ctx)
3153 .map_err(|e| -> Error { e.into() })?;
3154 while let Some(batch) = stream.next().await {
3155 batch.map_err(|e| -> Error { e.into() })?;
3156 }
3157 Ok::<_, Error>(())
3158 }));
3159 }
3160 for handle in handles {
3161 handle.await.map_err(|e| Error::Runtime {
3162 message: format!("Insert task panicked: {}", e),
3163 })??;
3164 }
3165
3166 let version = ds_wrapper.get().await?.manifest().version;
3167 self.bump_freshness();
3168 Ok(AddResult { version })
3169 }
3170
3171 async fn create_index(&self, opts: IndexBuilder) -> Result<()> {
3172 let prepared = self.prepare_index(&opts).await?;
3173 self.build_index(opts, prepared).await
3174 }
3175
3176 async fn create_index_async(&self, opts: IndexBuilder) -> Result<Job> {
3177 let prepared = self.prepare_index(&opts).await?;
3180 let table = self.clone();
3181 Ok(Job::spawned(tokio::spawn(async move {
3182 table.build_index(opts, prepared).await
3183 })))
3184 }
3185
3186 async fn drop_index(&self, index_name: &str) -> Result<()> {
3187 self.dataset.ensure_mutable()?;
3188 let mut dataset = (*self.dataset.get().await?).clone();
3189 dataset.drop_index(index_name).await?;
3190 self.dataset.update(dataset);
3191 Ok(())
3192 }
3193
3194 async fn prewarm_index(&self, index_name: &str) -> Result<()> {
3195 let dataset = self.dataset.get().await?;
3196 Ok(dataset.prewarm_index(index_name).await?)
3197 }
3198
3199 async fn prewarm_data(&self, _columns: Option<Vec<String>>) -> Result<()> {
3200 Err(Error::NotSupported {
3201 message: "prewarm_data is currently only supported on remote tables.".into(),
3202 })
3203 }
3204
3205 async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult> {
3206 let result = update::execute_update(self, update).await?;
3208 self.bump_freshness();
3209 Ok(result)
3210 }
3211
3212 async fn create_plan(
3213 &self,
3214 query: &AnyQuery,
3215 options: QueryExecutionOptions,
3216 ) -> Result<Arc<dyn ExecutionPlan>> {
3217 query::create_plan(self, query, options).await
3218 }
3219
3220 async fn query(
3221 &self,
3222 query: &AnyQuery,
3223 options: QueryExecutionOptions,
3224 ) -> Result<DatasetRecordBatchStream> {
3225 query::execute_query(self, query, options).await
3226 }
3227
3228 async fn analyze_plan(
3229 &self,
3230 query: &AnyQuery,
3231 options: QueryExecutionOptions,
3232 ) -> Result<String> {
3233 query::analyze_query_plan(self, query, options).await
3234 }
3235
3236 async fn merge_insert(
3237 &self,
3238 params: MergeInsertBuilder,
3239 new_data: Box<dyn RecordBatchReader + Send>,
3240 ) -> Result<MergeResult> {
3241 let result = merge::execute_merge_insert(self, params, new_data).await?;
3242 self.bump_freshness();
3243 Ok(result)
3244 }
3245
3246 async fn set_unenforced_primary_key(&self, columns: &[&str]) -> Result<()> {
3247 primary_key::set_unenforced_primary_key(self, columns).await
3248 }
3249
3250 async fn set_lsm_write_spec(&self, spec: LsmWriteSpec) -> Result<()> {
3251 merge::lsm::set_lsm_write_spec(self, spec).await
3252 }
3253
3254 async fn unset_lsm_write_spec(&self) -> Result<()> {
3255 merge::lsm::unset_lsm_write_spec(self).await
3256 }
3257
3258 async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
3259 merge::lsm::get_lsm_write_spec(self).await
3260 }
3261
3262 async fn close_lsm_writers(&self) -> Result<()> {
3263 merge::lsm::close_lsm_writers(self).await
3264 }
3265
3266 async fn blob_columns(&self) -> Result<Vec<String>> {
3267 let schema = self.schema().await?;
3268 Ok(crate::blob::blob_column_names(schema.as_ref()))
3269 }
3270
3271 async fn fetch_blobs(&self, column: &str, row_ids: &[u64]) -> Result<LargeBinaryArray> {
3272 let dataset = self.dataset.get().await?;
3273 crate::blob::take_blobs_aligned(&dataset, column, row_ids).await
3274 }
3275
3276 async fn fetch_blob_ranges(
3277 &self,
3278 column: &str,
3279 requests: &[BlobRangeRequest],
3280 ) -> Result<LargeBinaryArray> {
3281 let dataset = self.dataset.get().await?;
3282 crate::blob::take_blob_ranges_aligned(&dataset, column, requests).await
3283 }
3284
3285 async fn fetch_blob_files(
3286 &self,
3287 column: &str,
3288 row_ids: &[u64],
3289 ) -> Result<Vec<Option<BlobFile>>> {
3290 let dataset = self.dataset.get().await?;
3291 crate::blob::take_blob_files_aligned(&dataset, column, row_ids).await
3292 }
3293
3294 async fn delete(&self, predicate: Predicate<'_>) -> Result<DeleteResult> {
3296 let result = delete::execute_delete(self, predicate).await?;
3297 self.bump_freshness();
3298 Ok(result)
3299 }
3300
3301 async fn tags(&self) -> Result<Box<dyn Tags + '_>> {
3302 Ok(Box::new(NativeTags {
3303 dataset: self.dataset.clone(),
3304 }))
3305 }
3306
3307 async fn optimize(&self, action: OptimizeAction) -> Result<OptimizeStats> {
3308 optimize::execute_optimize(self, action).await
3310 }
3311
3312 async fn add_columns(
3313 &self,
3314 transforms: NewColumnTransform,
3315 read_columns: Option<Vec<String>>,
3316 ) -> Result<AddColumnsResult> {
3317 let result = schema_evolution::execute_add_columns(self, transforms, read_columns).await?;
3318 self.bump_freshness();
3319 Ok(result)
3320 }
3321
3322 async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult> {
3323 let result = schema_evolution::execute_alter_columns(self, alterations).await?;
3324 self.bump_freshness();
3325 Ok(result)
3326 }
3327
3328 async fn update_field_metadata(
3329 &self,
3330 updates: &[FieldMetadataUpdate],
3331 ) -> Result<UpdateFieldMetadataResult> {
3332 let result = schema_evolution::execute_update_field_metadata(self, updates).await?;
3333 self.bump_freshness();
3334 Ok(result)
3335 }
3336
3337 async fn drop_columns(&self, columns: &[&str]) -> Result<DropColumnsResult> {
3338 let result = schema_evolution::execute_drop_columns(self, columns).await?;
3339 self.bump_freshness();
3340 Ok(result)
3341 }
3342
3343 async fn list_indices(&self) -> Result<Vec<IndexConfig>> {
3344 let dataset = self.dataset.get().await?;
3345 let total_rows = dataset.count_rows(None).await? as u64;
3346 let descriptions = dataset.describe_indices(None).await?;
3347 let mut indices: Vec<IndexConfig> = descriptions
3348 .iter()
3349 .filter_map(|idx_desc| {
3350 let index_type: crate::index::IndexType = idx_desc
3351 .index_type()
3352 .parse()
3353 .unwrap_or(crate::index::IndexType::Unknown);
3354 if index_type == crate::index::IndexType::Unknown {
3355 log::debug!(
3359 "Skipping unrecognized index '{}' (type '{}') in list_indices",
3360 idx_desc.name(),
3361 idx_desc.index_type(),
3362 );
3363 return None;
3364 }
3365
3366 let field_ids = idx_desc.field_ids();
3367 let mut columns = Vec::with_capacity(field_ids.len());
3368 for field_id in field_ids {
3369 let field_path = match dataset.schema().field_path(*field_id as i32) {
3370 Ok(field_path) => field_path,
3371 Err(e) => {
3372 log::warn!(
3373 "Failed to resolve field path for index {} field id {}: {}",
3374 idx_desc.name(),
3375 field_id,
3376 e
3377 );
3378 return None;
3379 }
3380 };
3381 columns.push(field_path);
3382 }
3383
3384 let segments = idx_desc.segments();
3385 let index_uuid = segments.first().map(|seg| seg.uuid.to_string());
3386 let created_at = segments.iter().filter_map(|seg| seg.created_at).min();
3387 let index_version = segments.first().map(|seg| seg.index_version);
3388 let num_indexed_rows = idx_desc.rows_indexed();
3389
3390 Some(IndexConfig {
3391 name: idx_desc.name().to_string(),
3392 index_type,
3393 columns,
3394 index_uuid,
3395 type_url: Some(idx_desc.type_url().to_string()),
3396 created_at,
3397 num_indexed_rows: Some(num_indexed_rows),
3398 num_unindexed_rows: Some(total_rows.saturating_sub(num_indexed_rows)),
3399 size_bytes: idx_desc.total_size_bytes(),
3400 num_segments: Some(segments.len() as u32),
3401 index_version,
3402 index_details: idx_desc.details().ok(),
3403 })
3404 })
3405 .collect();
3406
3407 for index in indices
3408 .iter_mut()
3409 .filter(|index| index.index_type == crate::index::IndexType::FTS)
3410 {
3411 let Some(description) = descriptions
3412 .iter()
3413 .find(|description| description.name() == index.name)
3414 else {
3415 continue;
3416 };
3417 let segments = description.segments();
3418 let Some(segment) = segments.first() else {
3419 continue;
3420 };
3421 let params = load_segment_params(&dataset, segment).await?;
3422 let details = serde_json::to_string(¶ms).map_err(|source| Error::Other {
3423 message: format!(
3424 "Failed to serialize full text search configuration for index '{}'",
3425 index.name
3426 ),
3427 source: Some(Box::new(source)),
3428 })?;
3429 index.index_details = Some(details);
3430 }
3431 Ok(indices)
3432 }
3433
3434 async fn uri(&self) -> Result<String> {
3435 Ok(self.uri.clone())
3436 }
3437
3438 async fn storage_options(&self) -> Option<HashMap<String, String>> {
3439 self.initial_storage_options().await
3440 }
3441
3442 async fn initial_storage_options(&self) -> Option<HashMap<String, String>> {
3443 self.dataset
3444 .get()
3445 .await
3446 .ok()
3447 .and_then(|dataset| dataset.initial_storage_options().cloned())
3448 }
3449
3450 async fn latest_storage_options(&self) -> Result<Option<HashMap<String, String>>> {
3451 let dataset = self.dataset.get().await?;
3452 Ok(dataset.latest_storage_options().await?.map(|o| o.0))
3453 }
3454
3455 async fn index_stats(&self, index_name: &str) -> Result<Option<IndexStatistics>> {
3456 let dataset = self.dataset.get().await?;
3461
3462 let mut descriptions = dataset
3463 .describe_indices(Some(IndexCriteria::default().with_name(index_name)))
3464 .await?;
3465 let Some(description) = descriptions.pop() else {
3466 return Ok(None);
3467 };
3468
3469 let index_type: crate::index::IndexType = description
3470 .index_type()
3471 .parse()
3472 .unwrap_or(crate::index::IndexType::Unknown);
3473
3474 let is_vector = matches!(
3475 index_type,
3476 crate::index::IndexType::IvfFlat
3477 | crate::index::IndexType::IvfSq
3478 | crate::index::IndexType::IvfPq
3479 | crate::index::IndexType::IvfRq
3480 | crate::index::IndexType::IvfHnswPq
3481 | crate::index::IndexType::IvfHnswSq
3482 | crate::index::IndexType::IvfHnswFlat
3483 );
3484
3485 let distance_type = description.details().ok().and_then(|json| {
3488 #[derive(serde::Deserialize)]
3489 struct Details {
3490 metric_type: Option<String>,
3491 }
3492 serde_json::from_str::<Details>(&json)
3493 .ok()
3494 .and_then(|d| d.metric_type)
3495 .and_then(|m| match m.to_uppercase().as_str() {
3496 "L2" => Some(DistanceType::L2),
3497 "COSINE" => Some(DistanceType::Cosine),
3498 "DOT" => Some(DistanceType::Dot),
3499 "HAMMING" => Some(DistanceType::Hamming),
3500 _ => None,
3501 })
3502 });
3503
3504 if is_vector && distance_type.is_none() {
3507 let stats = dataset.index_statistics(index_name).await?;
3508 let mut stats: IndexStatisticsImpl =
3509 serde_json::from_str(&stats).map_err(|e| Error::InvalidInput {
3510 message: format!("error deserializing index statistics: {}", e),
3511 })?;
3512 let first_index = stats.indices.pop().ok_or_else(|| Error::InvalidInput {
3513 message: "index statistics is empty".to_string(),
3514 })?;
3515 return Ok(Some(IndexStatistics {
3516 num_indexed_rows: stats.num_indexed_rows,
3517 num_unindexed_rows: stats.num_unindexed_rows,
3518 index_type,
3519 distance_type: first_index.metric_type,
3520 num_indices: stats.num_indices,
3521 }));
3522 }
3523
3524 let num_indexed_rows = description.rows_indexed() as usize;
3525 let total_rows = dataset.count_rows(None).await?;
3526 let num_unindexed_rows = total_rows.saturating_sub(num_indexed_rows);
3527 Ok(Some(IndexStatistics {
3528 num_indexed_rows,
3529 num_unindexed_rows,
3530 index_type,
3531 distance_type,
3532 num_indices: Some(description.metadata().len() as u32),
3533 }))
3534 }
3535
3536 async fn wait_for_index(
3539 &self,
3540 index_names: &[&str],
3541 timeout: std::time::Duration,
3542 ) -> Result<()> {
3543 wait_for_index(self, index_names, timeout).await
3544 }
3545
3546 async fn stats(&self) -> Result<TableStatistics> {
3547 let num_rows = self.count_rows(None).await?;
3548 let num_indices = self.list_indices().await?.len();
3549 let ds = self.dataset.get().await?;
3550 let mut total_bytes = ds.manifest().summary().total_files_size as usize;
3558 for frag in ds.manifest().fragments.iter() {
3559 for overlay in &frag.overlays {
3560 if let Some(size) = overlay.data_file.file_size_bytes.get() {
3561 total_bytes += size.get() as usize;
3562 }
3563 }
3564 }
3565 for index in ds.load_indices().await?.iter() {
3566 total_bytes += index.total_size_bytes().unwrap_or(0) as usize;
3567 }
3568
3569 let frags = ds.get_fragments();
3570 let mut sorted_sizes = join_all(
3571 frags
3572 .iter()
3573 .map(|frag| async move { frag.physical_rows().await.unwrap_or(0) }),
3574 )
3575 .await;
3576 sorted_sizes.sort();
3577
3578 let small_frag_threshold = 100000;
3579 let num_fragments = sorted_sizes.len();
3580 let num_small_fragments = sorted_sizes
3581 .iter()
3582 .filter(|&&size| size < small_frag_threshold)
3583 .count();
3584
3585 let p25 = *sorted_sizes.get(num_fragments / 4).unwrap_or(&0);
3586 let p50 = *sorted_sizes.get(num_fragments / 2).unwrap_or(&0);
3587 let p75 = *sorted_sizes.get(num_fragments * 3 / 4).unwrap_or(&0);
3588 let p99 = *sorted_sizes.get(num_fragments * 99 / 100).unwrap_or(&0);
3589 let min = sorted_sizes.first().copied().unwrap_or(0);
3590 let max = sorted_sizes.last().copied().unwrap_or(0);
3591 let mean = sorted_sizes
3592 .iter()
3593 .copied()
3594 .sum::<usize>()
3595 .checked_div(num_fragments)
3596 .unwrap_or(0);
3597
3598 let frag_stats = FragmentStatistics {
3599 num_fragments,
3600 num_small_fragments,
3601 lengths: FragmentSummaryStats {
3602 min,
3603 max,
3604 mean,
3605 p25,
3606 p50,
3607 p75,
3608 p99,
3609 },
3610 };
3611 let stats = TableStatistics {
3612 total_bytes,
3613 num_rows,
3614 num_indices,
3615 fragment_stats: frag_stats,
3616 };
3617 Ok(stats)
3618 }
3619
3620 async fn create_insert_exec(
3621 &self,
3622 input: Arc<dyn datafusion_physical_plan::ExecutionPlan>,
3623 write_params: WriteParams,
3624 ) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
3625 let ds = self.dataset.get().await?;
3626 let dataset = Arc::new((*ds).clone());
3627 Ok(Arc::new(datafusion::insert::InsertExec::new(
3628 self.dataset.clone(),
3629 dataset,
3630 input,
3631 write_params,
3632 )))
3633 }
3634}
3635
3636#[skip_serializing_none]
3637#[derive(Debug, Deserialize, PartialEq)]
3638pub struct TableStatistics {
3639 pub total_bytes: usize,
3646
3647 pub num_rows: usize,
3649
3650 pub num_indices: usize,
3652
3653 pub fragment_stats: FragmentStatistics,
3655}
3656
3657#[skip_serializing_none]
3658#[derive(Debug, Deserialize, PartialEq)]
3659pub struct FragmentStatistics {
3660 pub num_fragments: usize,
3662
3663 pub num_small_fragments: usize,
3665
3666 pub lengths: FragmentSummaryStats,
3668 }
3672
3673#[skip_serializing_none]
3674#[derive(Debug, Deserialize, PartialEq)]
3675pub struct FragmentSummaryStats {
3676 pub min: usize,
3677 pub max: usize,
3678 pub mean: usize,
3679 pub p25: usize,
3680 pub p50: usize,
3681 pub p75: usize,
3682 pub p99: usize,
3683}
3684
3685#[cfg(test)]
3686#[allow(deprecated)]
3687mod tests {
3688 use std::sync::Arc;
3689 use std::sync::atomic::{AtomicBool, Ordering};
3690 use std::time::Duration;
3691
3692 use arrow_array::{
3693 Int32Array, RecordBatch, RecordBatchIterator, RecordBatchReader, StringArray,
3694 };
3695 use arrow_schema::{DataType, Field, Schema};
3696 use futures::TryStreamExt;
3697 use lance::Dataset;
3698 use lance::io::{ObjectStoreParams, WrappingObjectStore};
3699 use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION;
3700 use tempfile::tempdir;
3701
3702 use super::*;
3703 use crate::connect;
3704 use crate::connection::ConnectBuilder;
3705 use crate::io::object_store::io_tracking::IoTrackingStore;
3706 use crate::query::Select;
3707 use crate::query::{ExecutableQuery, QueryBase};
3708 use crate::test_utils::connection::new_test_connection;
3709
3710 #[test]
3711 fn test_tokenize_uses_explicit_simple_tokenizer() {
3712 let params =
3713 crate::index::scalar::FtsIndexBuilder::default().base_tokenizer("simple".to_string());
3714 let tokens = crate::tokenize("Running in cafés", ¶ms).unwrap();
3715
3716 assert_eq!(
3717 tokens,
3718 vec![
3719 FtsToken {
3720 text: "run".to_string(),
3721 position: 0,
3722 },
3723 FtsToken {
3724 text: "cafe".to_string(),
3725 position: 2,
3726 },
3727 ]
3728 );
3729 }
3730
3731 #[tokio::test]
3732 async fn test_open() {
3733 let tmp_dir = tempdir().unwrap();
3734 let dataset_path = tmp_dir.path().join("test.lance");
3735
3736 let batch = make_test_batches();
3737 let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
3738 Dataset::write(reader, dataset_path.to_str().unwrap(), None)
3739 .await
3740 .unwrap();
3741
3742 let table = NativeTable::open(dataset_path.to_str().unwrap())
3743 .await
3744 .unwrap();
3745
3746 assert_eq!(table.name, "test")
3747 }
3748
3749 #[tokio::test]
3750 async fn test_open_not_found() {
3751 let tmp_dir = tempdir().unwrap();
3752 let uri = tmp_dir.path().to_str().unwrap();
3753 let table = NativeTable::open(uri).await;
3754 assert!(matches!(table.unwrap_err(), Error::TableNotFound { .. }));
3755 }
3756
3757 #[tokio::test]
3758 async fn test_open_not_found_missing_lance_dir() {
3759 let tmp_dir = tempdir().unwrap();
3760 let dataset_path = tmp_dir.path().join("test.lance");
3761
3762 let err = NativeTable::open(dataset_path.to_str().unwrap())
3763 .await
3764 .unwrap_err();
3765 assert!(
3766 matches!(&err, Error::TableNotFound { name, .. } if name == "test"),
3767 "got {err:?}"
3768 );
3769 }
3770
3771 async fn write_then_corrupt_table(dir: &std::path::Path, remove_all: bool) -> String {
3776 let dataset_path = dir.join("test.lance");
3777 let uri = dataset_path.to_str().unwrap().to_string();
3778
3779 let batch = make_test_batches();
3780 let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
3781 Dataset::write(reader, &uri, None).await.unwrap();
3782
3783 if remove_all {
3784 for entry in std::fs::read_dir(&dataset_path).unwrap() {
3785 let entry = entry.unwrap();
3786 if entry.file_type().unwrap().is_dir() {
3787 std::fs::remove_dir_all(entry.path()).unwrap();
3788 } else {
3789 std::fs::remove_file(entry.path()).unwrap();
3790 }
3791 }
3792 assert_eq!(std::fs::read_dir(&dataset_path).unwrap().count(), 0);
3793 } else {
3794 let versions = dataset_path.join("_versions");
3795 assert!(versions.is_dir(), "expected manifests under {versions:?}");
3796 std::fs::remove_dir_all(&versions).unwrap();
3797 assert!(std::fs::read_dir(&dataset_path).unwrap().count() > 0);
3798 }
3799
3800 uri
3801 }
3802
3803 #[tokio::test]
3804 async fn test_open_corrupt_empty_dir() {
3805 let tmp_dir = tempdir().unwrap();
3806 let uri = write_then_corrupt_table(tmp_dir.path(), true).await;
3807
3808 let err = NativeTable::open(&uri).await.unwrap_err();
3809 assert!(
3810 matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
3811 "got {err:?}"
3812 );
3813 }
3814
3815 #[tokio::test]
3816 async fn test_open_corrupt_missing_manifest() {
3817 let tmp_dir = tempdir().unwrap();
3818 let uri = write_then_corrupt_table(tmp_dir.path(), false).await;
3819
3820 let err = NativeTable::open(&uri).await.unwrap_err();
3821 assert!(
3822 matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
3823 "got {err:?}"
3824 );
3825 }
3826
3827 #[tokio::test]
3830 async fn test_open_table_corrupt_is_still_listed() {
3831 let tmp_dir = tempdir().unwrap();
3832 let db = connect(tmp_dir.path().to_str().unwrap())
3833 .execute()
3834 .await
3835 .unwrap();
3836
3837 write_then_corrupt_table(tmp_dir.path(), true).await;
3838
3839 assert_eq!(
3840 db.table_names().execute().await.unwrap(),
3841 vec!["test".to_string()]
3842 );
3843 let err = db.open_table("test").execute().await.unwrap_err();
3844 assert!(
3845 matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
3846 "got {err:?}"
3847 );
3848 assert!(
3849 err.to_string().contains("exists but could not be loaded"),
3850 "got {err}"
3851 );
3852 }
3853
3854 #[test]
3855 #[cfg(not(windows))]
3856 fn test_object_store_path() {
3857 use std::path::Path as StdPath;
3858 let p = StdPath::new("s3://bucket/path/to/file");
3859 let c = p.join("subfile");
3860 assert_eq!(c.to_str().unwrap(), "s3://bucket/path/to/file/subfile");
3861 }
3862
3863 #[tokio::test]
3864 async fn test_count_rows() {
3865 let tmp_dir = tempdir().unwrap();
3866 let uri = tmp_dir.path().to_str().unwrap();
3867
3868 let batch = make_test_batches();
3869 let reader: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
3870 vec![Ok(batch.clone())],
3871 batch.schema(),
3872 ));
3873 let table = NativeTable::create(
3874 uri,
3875 "test",
3876 vec![],
3877 reader,
3878 None,
3879 None,
3880 None,
3881 None,
3882 HashSet::new(),
3883 )
3884 .await
3885 .unwrap();
3886
3887 assert_eq!(table.count_rows(None).await.unwrap(), 10);
3888 assert_eq!(
3889 table
3890 .count_rows(Some(Filter::Sql("i >= 5".to_string())))
3891 .await
3892 .unwrap(),
3893 5
3894 );
3895 }
3896
3897 #[derive(Default, Debug)]
3898 struct NoOpCacheWrapper {
3899 called: AtomicBool,
3900 }
3901
3902 impl NoOpCacheWrapper {
3903 fn called(&self) -> bool {
3904 self.called.load(Ordering::Relaxed)
3905 }
3906 }
3907
3908 impl WrappingObjectStore for NoOpCacheWrapper {
3909 fn wrap(
3910 &self,
3911 _store_prefix: &str,
3912 original: Arc<dyn object_store::ObjectStore>,
3913 ) -> Arc<dyn object_store::ObjectStore> {
3914 self.called.store(true, Ordering::Relaxed);
3915 original
3916 }
3917 }
3918
3919 #[tokio::test]
3920 async fn test_open_table_options() {
3921 let tmp_dir = tempdir().unwrap();
3922 let dataset_path = tmp_dir.path().join("test.lance");
3923 let uri = dataset_path.to_str().unwrap();
3924 let conn = connect(uri).execute().await.unwrap();
3925
3926 let batches = make_test_batches();
3927
3928 conn.create_table("my_table", batches)
3929 .execute()
3930 .await
3931 .unwrap();
3932
3933 let wrapper = Arc::new(NoOpCacheWrapper::default());
3934
3935 let object_store_params = ObjectStoreParams {
3936 object_store_wrapper: Some(wrapper.clone()),
3937 ..Default::default()
3938 };
3939 let param = ReadParams {
3940 store_options: Some(object_store_params),
3941 ..Default::default()
3942 };
3943 assert!(!wrapper.called());
3944 conn.open_table("my_table")
3945 .lance_read_params(param)
3946 .execute()
3947 .await
3948 .unwrap();
3949 assert!(wrapper.called());
3950 }
3951
3952 fn make_test_batches() -> RecordBatch {
3953 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
3954 RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from_iter_values(0..10))]).unwrap()
3955 }
3956
3957 #[tokio::test]
3958 async fn test_tags() {
3959 let tmp_dir = tempdir().unwrap();
3960 let uri = tmp_dir.path().to_str().unwrap();
3961
3962 let conn = ConnectBuilder::new(uri)
3963 .read_consistency_interval(Duration::from_secs(0))
3964 .execute()
3965 .await
3966 .unwrap();
3967 let table = conn
3968 .create_table("my_table", some_sample_data())
3969 .execute()
3970 .await
3971 .unwrap();
3972 assert_eq!(table.version().await.unwrap(), 1);
3973 table.add(some_sample_data()).execute().await.unwrap();
3974 assert_eq!(table.version().await.unwrap(), 2);
3975 let mut tags_manager = table.tags().await.unwrap();
3976 let tags = tags_manager.list().await.unwrap();
3977 assert!(tags.is_empty(), "Tags should be empty initially");
3978 let tag1 = "tag1";
3979 tags_manager.create(tag1, 1).await.unwrap();
3980 assert_eq!(tags_manager.get_version(tag1).await.unwrap(), 1);
3981 let tags = tags_manager.list().await.unwrap();
3982 assert_eq!(tags.len(), 1);
3983 assert!(tags.contains_key(tag1));
3984 assert_eq!(tags.get(tag1).unwrap().version, 1);
3985 tags_manager.create("tag2", 2).await.unwrap();
3986 assert_eq!(tags_manager.get_version("tag2").await.unwrap(), 2);
3987 let tags = tags_manager.list().await.unwrap();
3988 assert_eq!(tags.len(), 2);
3989 assert!(tags.contains_key(tag1));
3990 assert_eq!(tags.get(tag1).unwrap().version, 1);
3991 assert!(tags.contains_key("tag2"));
3992 assert_eq!(tags.get("tag2").unwrap().version, 2);
3993 table.add(some_sample_data()).execute().await.unwrap();
3995 tags_manager.update(tag1, 3).await.unwrap();
3996 assert_eq!(tags_manager.get_version(tag1).await.unwrap(), 3);
3997 tags_manager.delete("tag2").await.unwrap();
3998 let tags = tags_manager.list().await.unwrap();
3999 assert_eq!(tags.len(), 1);
4000 assert!(tags.contains_key(tag1));
4001 assert_eq!(tags.get(tag1).unwrap().version, 3);
4002 table.add(some_sample_data()).execute().await.unwrap();
4004 assert_eq!(table.version().await.unwrap(), 4);
4005 table.checkout_tag(tag1).await.unwrap();
4006 assert_eq!(table.version().await.unwrap(), 3);
4007 table.checkout_latest().await.unwrap();
4008 assert_eq!(table.version().await.unwrap(), 4);
4009 }
4010
4011 #[tokio::test]
4012 async fn test_branches() {
4013 let tmp_dir = tempdir().unwrap();
4014 let uri = tmp_dir.path().to_str().unwrap();
4015
4016 let conn = ConnectBuilder::new(uri)
4017 .read_consistency_interval(Duration::from_secs(0))
4018 .execute()
4019 .await
4020 .unwrap();
4021
4022 let table = conn
4024 .create_table("my_table", some_sample_data())
4025 .execute()
4026 .await
4027 .unwrap();
4028 assert_eq!(table.count_rows(None).await.unwrap(), 1);
4029 assert_eq!(table.current_branch(), None);
4030 let main_version = table.version().await.unwrap();
4031
4032 let branch = table.create_branch("exp", main_version).await.unwrap();
4034 assert_eq!(branch.current_branch().as_deref(), Some("exp"));
4035 assert_eq!(branch.count_rows(None).await.unwrap(), 1);
4036
4037 branch.add(some_sample_data()).execute().await.unwrap();
4039 assert_eq!(branch.count_rows(None).await.unwrap(), 2);
4040 assert_eq!(
4041 table.count_rows(None).await.unwrap(),
4042 1,
4043 "main must be untouched by branch writes"
4044 );
4045
4046 let branches = table.list_branches().await.unwrap();
4048 assert!(branches.contains_key("exp"));
4049
4050 let checked_out = table.checkout_branch("exp", None).await.unwrap();
4052 assert_eq!(checked_out.current_branch().as_deref(), Some("exp"));
4053 assert_eq!(checked_out.count_rows(None).await.unwrap(), 2);
4054
4055 let opened = conn
4057 .open_table("my_table")
4058 .branch("exp")
4059 .execute()
4060 .await
4061 .unwrap();
4062 assert_eq!(opened.current_branch().as_deref(), Some("exp"));
4063 assert_eq!(opened.count_rows(None).await.unwrap(), 2);
4064
4065 table.delete_branch("exp").await.unwrap();
4067 let branches = table.list_branches().await.unwrap();
4068 assert!(!branches.contains_key("exp"));
4069 }
4070
4071 #[tokio::test]
4072 async fn test_branch_version_checkout() {
4073 let tmp_dir = tempdir().unwrap();
4074 let uri = tmp_dir.path().to_str().unwrap();
4075
4076 let conn = ConnectBuilder::new(uri)
4077 .read_consistency_interval(Duration::from_secs(0))
4078 .execute()
4079 .await
4080 .unwrap();
4081
4082 let table = conn
4084 .create_table("my_table", sample_rows(vec![0]))
4085 .execute()
4086 .await
4087 .unwrap();
4088 let fork_point = table.version().await.unwrap();
4089
4090 let branch = table.create_branch("exp", fork_point).await.unwrap();
4093 let exp_fork = branch.version().await.unwrap(); branch.add(sample_rows(vec![1])).execute().await.unwrap(); let exp_v2 = branch.version().await.unwrap();
4096 branch.add(sample_rows(vec![2])).execute().await.unwrap(); table
4100 .add(sample_rows(vec![100, 101, 102]))
4101 .execute()
4102 .await
4103 .unwrap(); let main_v2 = table.version().await.unwrap();
4105 assert_eq!(
4106 exp_v2, main_v2,
4107 "branch and main must share the version number for this test to mean anything"
4108 );
4109
4110 let pinned = conn
4114 .open_table("my_table")
4115 .branch("exp")
4116 .version(exp_v2)
4117 .execute()
4118 .await
4119 .unwrap();
4120 assert_eq!(pinned.current_branch().as_deref(), Some("exp"));
4121 assert_eq!(pinned.count_rows(None).await.unwrap(), 2);
4123 assert_eq!(
4125 pinned.count_rows(Some("i = 1".to_string())).await.unwrap(),
4126 1
4127 );
4128 assert_eq!(
4129 pinned
4130 .count_rows(Some("i = 100".to_string()))
4131 .await
4132 .unwrap(),
4133 0
4134 );
4135
4136 let pinned_direct = table.checkout_branch("exp", Some(exp_v2)).await.unwrap();
4138 assert_eq!(pinned_direct.current_branch().as_deref(), Some("exp"));
4139 assert_eq!(pinned_direct.count_rows(None).await.unwrap(), 2);
4140
4141 let head = conn
4143 .open_table("my_table")
4144 .branch("exp")
4145 .execute()
4146 .await
4147 .unwrap();
4148 assert_eq!(head.count_rows(None).await.unwrap(), 3);
4149 assert_eq!(table.count_rows(None).await.unwrap(), 4);
4150
4151 assert!(pinned.add(sample_rows(vec![9])).execute().await.is_err());
4153
4154 let old_main = conn
4158 .open_table("my_table")
4159 .version(fork_point)
4160 .execute()
4161 .await
4162 .unwrap();
4163 assert_eq!(old_main.current_branch(), None);
4164 assert_eq!(old_main.count_rows(None).await.unwrap(), 1);
4165 let shared_on_main = conn
4166 .open_table("my_table")
4167 .version(exp_v2)
4168 .execute()
4169 .await
4170 .unwrap();
4171 assert_eq!(shared_on_main.current_branch(), None);
4172 assert_eq!(shared_on_main.count_rows(None).await.unwrap(), 4);
4173
4174 assert!(
4176 conn.open_table("my_table")
4177 .version(9999)
4178 .execute()
4179 .await
4180 .is_err()
4181 );
4182
4183 assert!(
4186 conn.open_table("my_table")
4187 .branch("exp")
4188 .version(9999)
4189 .execute()
4190 .await
4191 .is_err()
4192 );
4193
4194 let exp_at_fork = conn
4197 .open_table("my_table")
4198 .branch("exp")
4199 .version(exp_fork)
4200 .execute()
4201 .await
4202 .unwrap();
4203 assert_eq!(exp_at_fork.current_branch().as_deref(), Some("exp"));
4204 assert_eq!(exp_at_fork.count_rows(None).await.unwrap(), 1);
4205
4206 pinned.checkout_latest().await.unwrap();
4209 assert_eq!(pinned.current_branch().as_deref(), Some("exp"));
4210 assert_eq!(pinned.count_rows(None).await.unwrap(), 3); pinned.add(sample_rows(vec![3])).execute().await.unwrap();
4212 assert_eq!(pinned.count_rows(None).await.unwrap(), 4); }
4214
4215 #[tokio::test]
4216 async fn test_branch_version_two_branches() {
4217 let tmp_dir = tempdir().unwrap();
4218 let uri = tmp_dir.path().to_str().unwrap();
4219 let conn = ConnectBuilder::new(uri)
4220 .read_consistency_interval(Duration::from_secs(0))
4221 .execute()
4222 .await
4223 .unwrap();
4224
4225 let table = conn
4226 .create_table("my_table", sample_rows(vec![0]))
4227 .execute()
4228 .await
4229 .unwrap();
4230 let fork_point = table.version().await.unwrap();
4231
4232 let exp1 = table.create_branch("exp1", fork_point).await.unwrap();
4235 let exp2 = table.create_branch("exp2", fork_point).await.unwrap();
4236 exp1.add(sample_rows(vec![10])).execute().await.unwrap();
4237 exp2.add(sample_rows(vec![20])).execute().await.unwrap();
4238 let v1 = exp1.version().await.unwrap();
4239 let v2 = exp2.version().await.unwrap();
4240 assert_eq!(v1, v2, "both branches must reach the same version number");
4241
4242 let at1 = table.checkout_branch("exp1", Some(v1)).await.unwrap();
4244 assert_eq!(at1.count_rows(Some("i = 10".to_string())).await.unwrap(), 1);
4245 assert_eq!(at1.count_rows(Some("i = 20".to_string())).await.unwrap(), 0);
4246 let at2 = table.checkout_branch("exp2", Some(v2)).await.unwrap();
4247 assert_eq!(at2.count_rows(Some("i = 20".to_string())).await.unwrap(), 1);
4248 assert_eq!(at2.count_rows(Some("i = 10".to_string())).await.unwrap(), 0);
4249 }
4250
4251 #[tokio::test]
4252 async fn test_branch_name_validation() {
4253 let tmp_dir = tempdir().unwrap();
4254 let uri = tmp_dir.path().to_str().unwrap();
4255 let conn = ConnectBuilder::new(uri).execute().await.unwrap();
4256 let table = conn
4257 .create_table("my_table", some_sample_data())
4258 .execute()
4259 .await
4260 .unwrap();
4261
4262 assert!(matches!(
4264 table.create_branch("", 1u64).await,
4265 Err(Error::InvalidInput { .. })
4266 ));
4267 assert!(matches!(
4268 table.checkout_branch("", None).await,
4269 Err(Error::InvalidInput { .. })
4270 ));
4271 assert!(matches!(
4272 table.delete_branch("").await,
4273 Err(Error::InvalidInput { .. })
4274 ));
4275 assert!(matches!(
4277 table
4278 .create_branch(
4279 "ok",
4280 lance::dataset::refs::Ref::Version(Some(String::new()), None)
4281 )
4282 .await,
4283 Err(Error::InvalidInput { .. })
4284 ));
4285 }
4286
4287 #[tokio::test]
4288 async fn test_branch_handle_tracks_concurrent_writes() {
4289 let tmp_dir = tempdir().unwrap();
4290 let uri = tmp_dir.path().to_str().unwrap();
4291
4292 let conn = ConnectBuilder::new(uri)
4294 .read_consistency_interval(Duration::from_secs(0))
4295 .execute()
4296 .await
4297 .unwrap();
4298 let table = conn
4299 .create_table("my_table", some_sample_data())
4300 .execute()
4301 .await
4302 .unwrap();
4303 let v1 = table.version().await.unwrap();
4304
4305 let writer = table.create_branch("exp", v1).await.unwrap();
4307 let reader = conn
4308 .open_table("my_table")
4309 .branch("exp")
4310 .execute()
4311 .await
4312 .unwrap();
4313 assert_eq!(reader.count_rows(None).await.unwrap(), 1);
4314
4315 writer.add(some_sample_data()).execute().await.unwrap();
4318 assert_eq!(reader.count_rows(None).await.unwrap(), 2);
4319 assert_eq!(table.count_rows(None).await.unwrap(), 1);
4321 }
4322
4323 #[tokio::test]
4324 async fn test_branch_handle_without_consistency_interval_is_pinned() {
4325 let tmp_dir = tempdir().unwrap();
4326 let uri = tmp_dir.path().to_str().unwrap();
4327
4328 let conn = ConnectBuilder::new(uri).execute().await.unwrap();
4330 let table = conn
4331 .create_table("my_table", some_sample_data())
4332 .execute()
4333 .await
4334 .unwrap();
4335 let v1 = table.version().await.unwrap();
4336
4337 let writer = table.create_branch("exp", v1).await.unwrap();
4338 let reader = conn
4339 .open_table("my_table")
4340 .branch("exp")
4341 .execute()
4342 .await
4343 .unwrap();
4344 assert_eq!(reader.count_rows(None).await.unwrap(), 1);
4345
4346 writer.add(some_sample_data()).execute().await.unwrap();
4349 assert_eq!(reader.count_rows(None).await.unwrap(), 1);
4350
4351 reader.checkout_latest().await.unwrap();
4353 assert_eq!(reader.count_rows(None).await.unwrap(), 2);
4354 }
4355
4356 #[tokio::test]
4357 async fn test_dynamic_select() {
4358 let tc = new_test_connection().await.unwrap();
4359 let db = tc.connection;
4360
4361 let table = db
4362 .create_table("test", some_sample_data())
4363 .execute()
4364 .await
4365 .unwrap();
4366
4367 let query = table.query().select(Select::dynamic(&[("i_alias", "i")]));
4368
4369 let result = query.execute().await;
4370 let batches = result
4371 .expect("should have result")
4372 .try_collect::<Vec<_>>()
4373 .await
4374 .unwrap();
4375
4376 for batch in batches {
4377 assert!(batch.column_by_name("i_alias").is_some());
4378 }
4379 }
4380
4381 fn some_sample_data() -> Box<dyn arrow_array::RecordBatchReader + Send> {
4382 let batch = RecordBatch::try_new(
4383 Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])),
4384 vec![Arc::new(Int32Array::from(vec![1]))],
4385 )
4386 .unwrap();
4387 let schema = batch.schema().clone();
4388 let batch = Ok(batch);
4389
4390 Box::new(RecordBatchIterator::new(vec![batch], schema))
4391 }
4392
4393 fn sample_rows(values: Vec<i32>) -> Box<dyn arrow_array::RecordBatchReader + Send> {
4396 let batch = RecordBatch::try_new(
4397 Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])),
4398 vec![Arc::new(Int32Array::from(values))],
4399 )
4400 .unwrap();
4401 let schema = batch.schema().clone();
4402
4403 Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema))
4404 }
4405
4406 #[tokio::test]
4407 async fn test_read_consistency_interval() {
4408 use crate::utils::background_cache::clock;
4409
4410 let intervals = vec![
4411 None,
4412 Some(0),
4413 Some(100), ];
4415
4416 for interval in intervals {
4417 let data = some_sample_data();
4418
4419 let tmp_dir = tempdir().unwrap();
4420 let uri = tmp_dir.path().to_str().unwrap();
4421
4422 let conn1 = ConnectBuilder::new(uri).execute().await.unwrap();
4423 let table1 = conn1
4424 .create_empty_table("my_table", RecordBatchReader::schema(&data))
4425 .execute()
4426 .await
4427 .unwrap();
4428
4429 let mut conn2 = ConnectBuilder::new(uri);
4430 if let Some(interval) = interval {
4431 conn2 = conn2.read_consistency_interval(std::time::Duration::from_millis(interval));
4432 }
4433 let conn2 = conn2.execute().await.unwrap();
4434 let table2 = conn2.open_table("my_table").execute().await.unwrap();
4435
4436 clock::pin();
4441
4442 assert_eq!(table1.count_rows(None).await.unwrap(), 0);
4443 assert_eq!(table2.count_rows(None).await.unwrap(), 0);
4444
4445 table1.add(data).execute().await.unwrap();
4446 assert_eq!(table1.count_rows(None).await.unwrap(), 1);
4447
4448 match interval {
4449 None => {
4450 assert_eq!(table2.count_rows(None).await.unwrap(), 0);
4451 table2.checkout_latest().await.unwrap();
4452 assert_eq!(table2.count_rows(None).await.unwrap(), 1);
4453 }
4454 Some(0) => {
4455 assert_eq!(table2.count_rows(None).await.unwrap(), 1);
4456 }
4457 Some(100) => {
4458 assert_eq!(table2.count_rows(None).await.unwrap(), 0);
4459 clock::advance_by(Duration::from_millis(100));
4460 assert_eq!(table2.count_rows(None).await.unwrap(), 1);
4461 }
4462 _ => unreachable!(),
4463 }
4464 }
4465 }
4466
4467 #[tokio::test]
4468 async fn test_time_travel_write() {
4469 let tmp_dir = tempdir().unwrap();
4470 let uri = tmp_dir.path().to_str().unwrap();
4471
4472 let conn = ConnectBuilder::new(uri)
4473 .read_consistency_interval(Duration::from_secs(0))
4474 .execute()
4475 .await
4476 .unwrap();
4477 let table = conn
4478 .create_table("my_table", some_sample_data())
4479 .execute()
4480 .await
4481 .unwrap();
4482 let version = table.version().await.unwrap();
4483 table.add(some_sample_data()).execute().await.unwrap();
4484 table.checkout(version).await.unwrap();
4485 assert!(table.add(some_sample_data()).execute().await.is_err())
4486 }
4487
4488 #[tokio::test]
4489 async fn test_update_dataset_config() {
4490 let tmp_dir = tempdir().unwrap();
4491 let uri = tmp_dir.path().to_str().unwrap();
4492
4493 let conn = ConnectBuilder::new(uri)
4494 .read_consistency_interval(Duration::from_secs(0))
4495 .execute()
4496 .await
4497 .unwrap();
4498
4499 let table = conn
4500 .create_table("my_table", some_sample_data())
4501 .execute()
4502 .await
4503 .unwrap();
4504 let native_tbl = table.as_native().unwrap();
4505
4506 let manifest = native_tbl.manifest().await.unwrap();
4507 let base_config_len = manifest.config.len();
4508
4509 native_tbl
4510 .update_config(vec![("test_key1".to_string(), "test_val1".to_string())])
4511 .await
4512 .unwrap();
4513
4514 let manifest = native_tbl.manifest().await.unwrap();
4515 assert_eq!(manifest.config.len(), 1 + base_config_len);
4516 assert_eq!(
4517 manifest.config.get("test_key1"),
4518 Some(&"test_val1".to_string())
4519 );
4520
4521 native_tbl
4522 .update_config(vec![("test_key2".to_string(), "test_val2".to_string())])
4523 .await
4524 .unwrap();
4525 let manifest = native_tbl.manifest().await.unwrap();
4526 assert_eq!(manifest.config.len(), 2 + base_config_len);
4527 assert_eq!(
4528 manifest.config.get("test_key1"),
4529 Some(&"test_val1".to_string())
4530 );
4531 assert_eq!(
4532 manifest.config.get("test_key2"),
4533 Some(&"test_val2".to_string())
4534 );
4535
4536 native_tbl
4537 .update_config(vec![(
4538 "test_key2".to_string(),
4539 "test_val2_update".to_string(),
4540 )])
4541 .await
4542 .unwrap();
4543 let manifest = native_tbl.manifest().await.unwrap();
4544 assert_eq!(manifest.config.len(), 2 + base_config_len);
4545 assert_eq!(
4546 manifest.config.get("test_key1"),
4547 Some(&"test_val1".to_string())
4548 );
4549 assert_eq!(
4550 manifest.config.get("test_key2"),
4551 Some(&"test_val2_update".to_string())
4552 );
4553
4554 native_tbl.delete_config_keys(&["test_key1"]).await.unwrap();
4555 let manifest = native_tbl.manifest().await.unwrap();
4556 assert_eq!(manifest.config.len(), 1 + base_config_len);
4557 assert_eq!(
4558 manifest.config.get("test_key2"),
4559 Some(&"test_val2_update".to_string())
4560 );
4561 }
4562
4563 #[tokio::test]
4564 async fn test_schema_metadata_config() {
4565 let tmp_dir = tempdir().unwrap();
4566 let uri = tmp_dir.path().to_str().unwrap();
4567
4568 let conn = ConnectBuilder::new(uri)
4569 .read_consistency_interval(Duration::from_secs(0))
4570 .execute()
4571 .await
4572 .unwrap();
4573 let table = conn
4574 .create_table("my_table", some_sample_data())
4575 .execute()
4576 .await
4577 .unwrap();
4578
4579 let native_tbl = table.as_native().unwrap();
4580 let schema = native_tbl.schema().await.unwrap();
4581 let metadata = schema.metadata();
4582 assert_eq!(metadata.len(), 0);
4583
4584 native_tbl
4585 .replace_schema_metadata(vec![("test_key1".to_string(), "test_val1".to_string())])
4586 .await
4587 .unwrap();
4588
4589 let schema = native_tbl.schema().await.unwrap();
4590 let metadata = schema.metadata();
4591 assert_eq!(metadata.len(), 1);
4592 assert_eq!(metadata.get("test_key1"), Some(&"test_val1".to_string()));
4593
4594 native_tbl
4595 .replace_schema_metadata(vec![
4596 ("test_key1".to_string(), "test_val1_update".to_string()),
4597 ("test_key2".to_string(), "test_val2".to_string()),
4598 ])
4599 .await
4600 .unwrap();
4601 let schema = native_tbl.schema().await.unwrap();
4602 let metadata = schema.metadata();
4603 assert_eq!(metadata.len(), 2);
4604 assert_eq!(
4605 metadata.get("test_key1"),
4606 Some(&"test_val1_update".to_string())
4607 );
4608 assert_eq!(metadata.get("test_key2"), Some(&"test_val2".to_string()));
4609
4610 native_tbl
4611 .replace_schema_metadata(vec![(
4612 "test_key2".to_string(),
4613 "test_val2_update".to_string(),
4614 )])
4615 .await
4616 .unwrap();
4617 let schema = native_tbl.schema().await.unwrap();
4618 let metadata = schema.metadata();
4619 assert_eq!(
4620 metadata.get("test_key2"),
4621 Some(&"test_val2_update".to_string())
4622 );
4623 }
4624
4625 #[tokio::test]
4626 pub async fn test_field_metadata_update() {
4627 let tmp_dir = tempdir().unwrap();
4628 let uri = tmp_dir.path().to_str().unwrap();
4629
4630 let conn = ConnectBuilder::new(uri)
4631 .read_consistency_interval(Duration::from_secs(0))
4632 .execute()
4633 .await
4634 .unwrap();
4635 let table = conn
4636 .create_table("my_table", some_sample_data())
4637 .execute()
4638 .await
4639 .unwrap();
4640
4641 let native_tbl = table.as_native().unwrap();
4642 let schema = native_tbl.manifest().await.unwrap().schema;
4643
4644 let field = schema.field("i").unwrap();
4645 assert_eq!(field.metadata.len(), 0);
4646
4647 native_tbl
4648 .replace_schema_metadata(vec![(
4649 "test_key2".to_string(),
4650 "test_val2_update".to_string(),
4651 )])
4652 .await
4653 .unwrap();
4654
4655 let schema = native_tbl.schema().await.unwrap();
4656 let metadata = schema.metadata();
4657 assert_eq!(metadata.len(), 1);
4658 assert_eq!(
4659 metadata.get("test_key2"),
4660 Some(&"test_val2_update".to_string())
4661 );
4662
4663 native_tbl
4664 .update_field_metadata(&[
4665 FieldMetadataUpdate::new("i").set("test_field_key1", "test_field_val1")
4666 ])
4667 .await
4668 .unwrap();
4669
4670 let schema = native_tbl.manifest().await.unwrap().schema;
4671 let field = schema.field("i").unwrap();
4672 assert_eq!(field.metadata.len(), 1);
4673 assert_eq!(
4674 field.metadata.get("test_field_key1"),
4675 Some(&"test_field_val1".to_string())
4676 );
4677 }
4678
4679 #[tokio::test]
4680 async fn test_set_unenforced_primary_key() {
4681 let tmp_dir = tempdir().unwrap();
4682 let uri = tmp_dir.path().to_str().unwrap();
4683
4684 let schema = Arc::new(Schema::new(vec![
4685 Field::new("id", DataType::Int64, false),
4686 Field::new("name", DataType::Utf8, true),
4687 Field::new("score", DataType::Float64, true),
4688 ]));
4689 let batch = RecordBatch::try_new(
4690 schema.clone(),
4691 vec![
4692 Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
4693 Arc::new(StringArray::from(vec!["a", "b", "c"])),
4694 Arc::new(arrow_array::Float64Array::from(vec![1.0, 2.0, 3.0])),
4695 ],
4696 )
4697 .unwrap();
4698 let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4699 Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4700
4701 let conn = ConnectBuilder::new(uri)
4702 .read_consistency_interval(Duration::from_secs(0))
4703 .execute()
4704 .await
4705 .unwrap();
4706 let table = conn.create_table("t", reader).execute().await.unwrap();
4707
4708 let err = table
4710 .set_unenforced_primary_key(Vec::<&str>::new())
4711 .await
4712 .expect_err("empty input should be rejected");
4713 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4714
4715 let err = table
4717 .set_unenforced_primary_key(["id", "name"])
4718 .await
4719 .expect_err("compound primary key should be rejected");
4720 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4721
4722 let err = table
4724 .set_unenforced_primary_key(["nonexistent"])
4725 .await
4726 .expect_err("nonexistent column should be rejected");
4727 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4728
4729 let err = table
4731 .set_unenforced_primary_key(["score"])
4732 .await
4733 .expect_err("Float64 should be rejected");
4734 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4735
4736 let lance_schema = table.as_native().unwrap().manifest().await.unwrap().schema;
4738 assert!(lance_schema.unenforced_primary_key().is_empty());
4739
4740 table.set_unenforced_primary_key(["id"]).await.unwrap();
4742 let lance_schema = table.as_native().unwrap().manifest().await.unwrap().schema;
4743 let pk = lance_schema.unenforced_primary_key();
4744 assert_eq!(pk.len(), 1);
4745 assert_eq!(pk[0].name, "id");
4746 assert_eq!(
4748 pk[0].metadata.get(LANCE_UNENFORCED_PRIMARY_KEY_POSITION),
4749 Some(&"1".to_string())
4750 );
4751
4752 let err = table
4755 .set_unenforced_primary_key(["id"])
4756 .await
4757 .expect_err("re-setting the same primary key should be rejected");
4758 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4759 let err = table
4760 .set_unenforced_primary_key(["name"])
4761 .await
4762 .expect_err("changing the primary key should be rejected");
4763 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4764
4765 let lance_schema = table.as_native().unwrap().manifest().await.unwrap().schema;
4767 let pk = lance_schema.unenforced_primary_key();
4768 assert_eq!(pk.len(), 1);
4769 assert_eq!(pk[0].name, "id");
4770 }
4771
4772 #[tokio::test]
4773 async fn test_set_unenforced_primary_key_concurrent() {
4774 let tmp_dir = tempdir().unwrap();
4775 let uri = tmp_dir.path().to_str().unwrap();
4776
4777 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
4778 let batch = RecordBatch::try_new(
4779 schema.clone(),
4780 vec![Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3]))],
4781 )
4782 .unwrap();
4783 let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4784 Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4785
4786 let conn = ConnectBuilder::new(uri)
4790 .read_consistency_interval(Duration::from_secs(3600))
4791 .execute()
4792 .await
4793 .unwrap();
4794 conn.create_table("t", reader).execute().await.unwrap();
4795
4796 let table_a = conn.open_table("t").execute().await.unwrap();
4797 let table_b = conn.open_table("t").execute().await.unwrap();
4798
4799 table_a.set_unenforced_primary_key(["id"]).await.unwrap();
4801
4802 let err = table_b
4808 .set_unenforced_primary_key(["id"])
4809 .await
4810 .expect_err("concurrent primary key commit on a stale base should fail");
4811 assert!(
4812 !matches!(err, Error::InvalidInput { .. }),
4813 "expected a commit-time conflict, not an up-front input error: {:?}",
4814 err
4815 );
4816
4817 let fresh = conn.open_table("t").execute().await.unwrap();
4819 let lance_schema = fresh.as_native().unwrap().manifest().await.unwrap().schema;
4820 let pk = lance_schema.unenforced_primary_key();
4821 assert_eq!(pk.len(), 1);
4822 assert_eq!(pk[0].name, "id");
4823 }
4824
4825 #[tokio::test]
4826 async fn test_set_lsm_write_spec() {
4827 use arrow_array::StringArray;
4828 use lance::dataset::mem_wal::DatasetMemWalExt;
4829
4830 let tmp_dir = tempdir().unwrap();
4831 let uri = tmp_dir.path().to_str().unwrap();
4832
4833 let schema = Arc::new(Schema::new(vec![
4834 Field::new("id", DataType::Int64, false),
4835 Field::new("name", DataType::Utf8, true),
4836 ]));
4837 let batch = RecordBatch::try_new(
4838 schema.clone(),
4839 vec![
4840 Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
4841 Arc::new(StringArray::from(vec!["a", "b", "c"])),
4842 ],
4843 )
4844 .unwrap();
4845 let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4846 Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4847
4848 let conn = ConnectBuilder::new(uri)
4849 .read_consistency_interval(Duration::from_secs(0))
4850 .execute()
4851 .await
4852 .unwrap();
4853 let table = conn.create_table("t", reader).execute().await.unwrap();
4854
4855 for bad in [0u32, 1025] {
4857 let err = table
4858 .set_lsm_write_spec(LsmWriteSpec::bucket("id", bad))
4859 .await
4860 .expect_err("should reject");
4861 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4862 }
4863
4864 table
4866 .set_lsm_write_spec(LsmWriteSpec::bucket("id", 4))
4867 .await
4868 .unwrap();
4869
4870 let native_tbl = table.as_native().unwrap();
4871 let dataset = native_tbl.dataset.get().await.unwrap();
4872 let details = dataset
4873 .mem_wal_index_details()
4874 .await
4875 .unwrap()
4876 .expect("MemWAL index should be initialized");
4877 assert_eq!(details.num_shards, 4);
4878 assert_eq!(details.sharding_specs.len(), 1);
4879 let installed = &details.sharding_specs[0];
4880 assert_eq!(installed.fields.len(), 1);
4881 let f = &installed.fields[0];
4882 assert_eq!(f.transform.as_deref(), Some("bucket"));
4883 assert_eq!(
4884 f.parameters.get("num_buckets").map(String::as_str),
4885 Some("4")
4886 );
4887 assert_eq!(f.parameters.len(), 1);
4889
4890 let err = table
4892 .set_lsm_write_spec(LsmWriteSpec::bucket("id", 8))
4893 .await
4894 .expect_err("mutation should be rejected");
4895 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
4896 }
4897
4898 #[tokio::test]
4899 async fn test_set_lsm_write_spec_unsharded() {
4900 use lance::dataset::mem_wal::DatasetMemWalExt;
4901
4902 let tmp_dir = tempdir().unwrap();
4903 let uri = tmp_dir.path().to_str().unwrap();
4904
4905 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
4906 let batch = RecordBatch::try_new(
4907 schema.clone(),
4908 vec![Arc::new(arrow_array::Int64Array::from(vec![1]))],
4909 )
4910 .unwrap();
4911 let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4912 Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4913 let conn = ConnectBuilder::new(uri)
4914 .read_consistency_interval(Duration::from_secs(0))
4915 .execute()
4916 .await
4917 .unwrap();
4918 let table = conn.create_table("t", reader).execute().await.unwrap();
4919
4920 table
4921 .set_lsm_write_spec(LsmWriteSpec::unsharded())
4922 .await
4923 .unwrap();
4924
4925 let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
4926 let details = dataset
4927 .mem_wal_index_details()
4928 .await
4929 .unwrap()
4930 .expect("MemWAL index should be initialized");
4931 assert_eq!(details.num_shards, 1);
4932 assert_eq!(details.sharding_specs.len(), 1);
4933 let f = &details.sharding_specs[0].fields[0];
4934 assert_eq!(f.transform.as_deref(), Some("unsharded"));
4935 assert!(f.source_ids.is_empty());
4936 }
4937
4938 #[tokio::test]
4939 async fn test_set_lsm_write_spec_identity() {
4940 use lance::dataset::mem_wal::DatasetMemWalExt;
4941
4942 let tmp_dir = tempdir().unwrap();
4943 let uri = tmp_dir.path().to_str().unwrap();
4944
4945 let schema = Arc::new(Schema::new(vec![
4946 Field::new("id", DataType::Int64, false),
4947 Field::new("region", DataType::Utf8, true),
4948 ]));
4949 let batch = RecordBatch::try_new(
4950 schema.clone(),
4951 vec![
4952 Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
4953 Arc::new(StringArray::from(vec!["a", "b", "c"])),
4954 ],
4955 )
4956 .unwrap();
4957 let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
4958 Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
4959 let conn = ConnectBuilder::new(uri)
4960 .read_consistency_interval(Duration::from_secs(0))
4961 .execute()
4962 .await
4963 .unwrap();
4964 let table = conn.create_table("t", reader).execute().await.unwrap();
4965
4966 table
4967 .set_lsm_write_spec(
4968 LsmWriteSpec::identity("region")
4969 .with_writer_config_defaults([("durable_write", "false")]),
4970 )
4971 .await
4972 .unwrap();
4973
4974 let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
4975 let details = dataset
4976 .mem_wal_index_details()
4977 .await
4978 .unwrap()
4979 .expect("MemWAL index should be initialized");
4980 assert_eq!(details.num_shards, 0);
4982 assert_eq!(details.sharding_specs.len(), 1);
4983 let f = &details.sharding_specs[0].fields[0];
4984 assert_eq!(f.transform.as_deref(), Some("identity"));
4985 assert_eq!(
4987 details
4988 .writer_config_defaults
4989 .get("durable_write")
4990 .map(String::as_str),
4991 Some("false")
4992 );
4993 }
4994
4995 #[tokio::test]
4996 async fn test_unset_lsm_write_spec() {
4997 use lance::dataset::mem_wal::DatasetMemWalExt;
4998
4999 let tmp_dir = tempdir().unwrap();
5000 let uri = tmp_dir.path().to_str().unwrap();
5001
5002 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
5003 let batch = RecordBatch::try_new(
5004 schema.clone(),
5005 vec![Arc::new(arrow_array::Int64Array::from(vec![1]))],
5006 )
5007 .unwrap();
5008 let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
5009 Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
5010 let conn = ConnectBuilder::new(uri)
5011 .read_consistency_interval(Duration::from_secs(0))
5012 .execute()
5013 .await
5014 .unwrap();
5015 let table = conn.create_table("t", reader).execute().await.unwrap();
5016
5017 table.unset_lsm_write_spec().await.unwrap_err();
5019
5020 table
5022 .set_lsm_write_spec(LsmWriteSpec::bucket("id", 4))
5023 .await
5024 .unwrap();
5025 {
5026 let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
5027 assert!(dataset.mem_wal_index_details().await.unwrap().is_some());
5028 }
5029
5030 table.unset_lsm_write_spec().await.unwrap();
5031 {
5032 let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
5033 assert!(dataset.mem_wal_index_details().await.unwrap().is_none());
5034 }
5035
5036 table.unset_lsm_write_spec().await.unwrap_err();
5038 table
5039 .set_lsm_write_spec(LsmWriteSpec::bucket("id", 8))
5040 .await
5041 .unwrap();
5042 {
5043 let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
5044 assert!(dataset.mem_wal_index_details().await.unwrap().is_some());
5045 }
5046 }
5047
5048 #[tokio::test]
5049 async fn test_get_lsm_write_spec() {
5050 let tmp_dir = tempdir().unwrap();
5051 let uri = tmp_dir.path().to_str().unwrap();
5052
5053 let schema = Arc::new(Schema::new(vec![
5054 Field::new("id", DataType::Int64, false),
5055 Field::new("region", DataType::Utf8, true),
5056 ]));
5057 let batch = RecordBatch::try_new(
5058 schema.clone(),
5059 vec![
5060 Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
5061 Arc::new(StringArray::from(vec!["a", "b", "c"])),
5062 ],
5063 )
5064 .unwrap();
5065 let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
5066 Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
5067 let conn = ConnectBuilder::new(uri)
5068 .read_consistency_interval(Duration::from_secs(0))
5069 .execute()
5070 .await
5071 .unwrap();
5072 let table = conn.create_table("t", reader).execute().await.unwrap();
5073
5074 assert_eq!(table.get_lsm_write_spec().await.unwrap(), None);
5076
5077 table
5079 .create_index(&["id"], Index::Auto)
5080 .execute()
5081 .await
5082 .unwrap();
5083 let idx_name = table.list_indices().await.unwrap()[0].name.clone();
5084
5085 let spec = LsmWriteSpec::bucket("id", 4)
5088 .with_maintained_indexes([idx_name])
5089 .with_writer_config_defaults([("durable_write", "false")]);
5090 table.set_lsm_write_spec(spec.clone()).await.unwrap();
5091 assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
5092
5093 table.unset_lsm_write_spec().await.unwrap();
5095 assert_eq!(table.get_lsm_write_spec().await.unwrap(), None);
5096
5097 let spec = LsmWriteSpec::identity("region");
5099 table.set_lsm_write_spec(spec.clone()).await.unwrap();
5100 assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
5101 table.unset_lsm_write_spec().await.unwrap();
5102
5103 let spec = LsmWriteSpec::unsharded();
5105 table.set_lsm_write_spec(spec.clone()).await.unwrap();
5106 assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
5107 }
5108
5109 #[tokio::test]
5110 pub async fn test_stats() {
5111 let tmp_dir = tempdir().unwrap();
5112 let uri = tmp_dir.path().to_str().unwrap();
5113
5114 let conn = ConnectBuilder::new(uri).execute().await.unwrap();
5115
5116 let schema = Arc::new(Schema::new(vec![
5117 Field::new("id", DataType::Int32, false),
5118 Field::new("foo", DataType::Int32, true),
5119 ]));
5120 let batch = RecordBatch::try_new(
5121 schema.clone(),
5122 vec![
5123 Arc::new(Int32Array::from_iter_values(0..100)),
5124 Arc::new(Int32Array::from_iter_values(0..100)),
5125 ],
5126 )
5127 .unwrap();
5128
5129 let table = conn
5130 .create_table("test_stats", batch.clone())
5131 .execute()
5132 .await
5133 .unwrap();
5134 for _ in 0..10 {
5135 let batch = RecordBatch::try_new(
5136 schema.clone(),
5137 vec![
5138 Arc::new(Int32Array::from_iter_values(0..15)),
5139 Arc::new(Int32Array::from_iter_values(0..15)),
5140 ],
5141 )
5142 .unwrap();
5143 table.add(batch.clone()).execute().await.unwrap();
5144 }
5145
5146 let empty_table = conn
5147 .create_table("test_stats_empty", RecordBatch::new_empty(batch.schema()))
5148 .execute()
5149 .await
5150 .unwrap();
5151
5152 let res = table.stats().await.unwrap();
5153 println!("{:#?}", res);
5154 assert_eq!(
5159 res,
5160 TableStatistics {
5161 num_rows: 250,
5162 num_indices: 0,
5163 total_bytes: 8925,
5164 fragment_stats: FragmentStatistics {
5165 num_fragments: 11,
5166 num_small_fragments: 11,
5167 lengths: FragmentSummaryStats {
5168 min: 15,
5169 max: 100,
5170 mean: 22,
5171 p25: 15,
5172 p50: 15,
5173 p75: 15,
5174 p99: 100,
5175 },
5176 },
5177 }
5178 );
5179 let res = empty_table.stats().await.unwrap();
5180 println!("{:#?}", res);
5181 assert_eq!(
5182 res,
5183 TableStatistics {
5184 num_rows: 0,
5185 num_indices: 0,
5186 total_bytes: 0,
5187 fragment_stats: FragmentStatistics {
5188 num_fragments: 0,
5189 num_small_fragments: 0,
5190 lengths: FragmentSummaryStats {
5191 min: 0,
5192 max: 0,
5193 mean: 0,
5194 p25: 0,
5195 p50: 0,
5196 p75: 0,
5197 p99: 0,
5198 },
5199 },
5200 }
5201 )
5202 }
5203
5204 #[tokio::test]
5207 pub async fn test_stats_includes_index_and_overlay_files() {
5208 use lance::dataset::WriteDestination;
5209 use lance::dataset::transaction::{DataOverlayGroup, Operation};
5210 use lance_file::version::{ConcreteFileVersion, LanceFileVersion};
5211 use lance_file::writer::{FileWriter, FileWriterOptions};
5212 use lance_io::utils::CachedFileSize;
5213 use lance_table::format::DataFile;
5214 use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage};
5215 use roaring::RoaringBitmap;
5216
5217 let tmp_dir = tempdir().unwrap();
5218 let uri = tmp_dir.path().to_str().unwrap();
5219 let conn = ConnectBuilder::new(uri)
5220 .read_consistency_interval(Duration::from_secs(0))
5221 .execute()
5222 .await
5223 .unwrap();
5224
5225 let schema = Arc::new(Schema::new(vec![
5226 Field::new("id", DataType::Int32, false),
5227 Field::new("foo", DataType::Int32, true),
5228 ]));
5229 let batch = RecordBatch::try_new(
5230 schema.clone(),
5231 vec![
5232 Arc::new(Int32Array::from_iter_values(0..100)),
5233 Arc::new(Int32Array::from_iter_values(0..100)),
5234 ],
5235 )
5236 .unwrap();
5237 let table = conn
5238 .create_table("test_stats_extra_files", batch)
5239 .execute()
5240 .await
5241 .unwrap();
5242
5243 let data_only = table.stats().await.unwrap().total_bytes;
5244 assert!(data_only > 0);
5245
5246 table
5249 .create_index(&["id"], Index::Auto)
5250 .execute()
5251 .await
5252 .unwrap();
5253 let with_index = table.stats().await.unwrap().total_bytes;
5254 let dataset = {
5255 let native = table.as_native().unwrap();
5256 (*native.dataset.get().await.unwrap()).clone()
5257 };
5258 let index_bytes: usize = dataset
5259 .load_indices()
5260 .await
5261 .unwrap()
5262 .iter()
5263 .map(|idx| idx.total_size_bytes().unwrap_or(0) as usize)
5264 .sum();
5265 assert!(index_bytes > 0);
5266 assert_eq!(with_index, data_only + index_bytes);
5267
5268 let read_version = dataset.version().version;
5273 let fragment_id = dataset.get_fragments()[0].id() as u64;
5274 let foo_field_id = dataset.schema().field("foo").unwrap().id;
5275 let overlay_schema = dataset.schema().project_by_ids(&[foo_field_id], true);
5276 let file_version = ConcreteFileVersion::from(LanceFileVersion::Stable);
5277
5278 let filename = "overlay.lance".to_string();
5279 let store = dataset.object_store(None).await.unwrap();
5280 let path = dataset.data_dir().child(filename.clone());
5281 let obj_writer = store.create(&path).await.unwrap();
5282 let mut writer = FileWriter::try_new(
5283 obj_writer,
5284 overlay_schema,
5285 FileWriterOptions {
5286 format_version: Some(file_version.into()),
5287 ..Default::default()
5288 },
5289 )
5290 .unwrap();
5291 writer
5292 .write_column(0, Arc::new(Int32Array::from(vec![1000, 1001, 1002])) as _)
5293 .await
5294 .unwrap();
5295 let summary = writer.finish().await.unwrap();
5296 let overlay_bytes = summary.size_bytes as usize;
5297 assert!(overlay_bytes > 0);
5298
5299 let mut data_file = DataFile::new_unstarted(filename, file_version);
5300 data_file.fields = writer
5301 .field_id_to_column_indices()
5302 .iter()
5303 .map(|(field_id, _)| *field_id as i32)
5304 .collect::<Vec<_>>()
5305 .into();
5306 data_file.column_indices = writer
5307 .field_id_to_column_indices()
5308 .iter()
5309 .map(|(_, column_index)| *column_index as i32)
5310 .collect::<Vec<_>>()
5311 .into();
5312 data_file.file_size_bytes = CachedFileSize::new(summary.size_bytes);
5313
5314 let overlay = DataOverlayFile {
5315 data_file,
5316 coverage: OverlayCoverage::dense(RoaringBitmap::from_iter(0..3)),
5317 committed_version: 0,
5318 };
5319 Dataset::commit(
5320 WriteDestination::Dataset(Arc::new(dataset)),
5321 Operation::DataOverlay {
5322 groups: vec![DataOverlayGroup {
5323 fragment_id,
5324 overlays: vec![overlay],
5325 }],
5326 },
5327 Some(read_version),
5328 None,
5329 None,
5330 Arc::new(Default::default()),
5331 false,
5332 )
5333 .await
5334 .unwrap();
5335
5336 table.checkout_latest().await.unwrap();
5337 let with_overlay = table.stats().await.unwrap().total_bytes;
5338 assert_eq!(with_overlay, with_index + overlay_bytes);
5339 }
5340
5341 #[tokio::test]
5344 pub async fn test_stats_does_not_read_data_files() {
5345 let tmp_dir = tempdir().unwrap();
5346 let uri = tmp_dir.path().to_str().unwrap();
5347
5348 let conn = ConnectBuilder::new(uri).execute().await.unwrap();
5349
5350 let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
5351 let batch = RecordBatch::try_new(
5352 schema.clone(),
5353 vec![Arc::new(Int32Array::from_iter_values(0..10))],
5354 )
5355 .unwrap();
5356
5357 conn.create_table("test_stats_io", batch.clone())
5358 .execute()
5359 .await
5360 .unwrap();
5361 let table = conn.open_table("test_stats_io").execute().await.unwrap();
5362 const NUM_APPENDS: usize = 20;
5363 for _ in 0..NUM_APPENDS {
5364 table.add(batch.clone()).execute().await.unwrap();
5365 }
5366
5367 let (wrapper, io_stats) = IoTrackingStore::new_wrapper();
5370 let table = conn
5371 .open_table("test_stats_io")
5372 .lance_read_params(ReadParams {
5373 store_options: Some(ObjectStoreParams {
5374 object_store_wrapper: Some(wrapper),
5375 ..Default::default()
5376 }),
5377 ..Default::default()
5378 })
5379 .execute()
5380 .await
5381 .unwrap();
5382 io_stats.lock().unwrap().read_iops = 0;
5383
5384 let stats = table.stats().await.unwrap();
5385 let read_iops = io_stats.lock().unwrap().read_iops;
5386
5387 assert_eq!(stats.fragment_stats.num_fragments, NUM_APPENDS + 1);
5388 assert!(stats.total_bytes > 0);
5389 assert!(
5391 read_iops < stats.fragment_stats.num_fragments as u64,
5392 "stats() issued {} read IOPs across {} fragments",
5393 read_iops,
5394 stats.fragment_stats.num_fragments
5395 );
5396 }
5397}